"""
Stage 3: Style Direction Agent
================================
Generates 3-4 complete design token sets with distinct personalities,
renders previews against the approved IA, and manages style selection.

The agent:
- Reads the brief (brand constraints, personas, competitive context)
- Reads the approved IA (screen structures to apply styles to)
- Generates internally consistent token sets
- Produces a tailwind.config.js from the selected tokens
- Requires user approval before advancing
"""

import json
import os
from datetime import datetime


# ═══════════════════════════════════════════════════════════
# SYSTEM PROMPT
# ═══════════════════════════════════════════════════════════

STYLE_SYSTEM_PROMPT = """You are a senior visual designer working as a design partner. You are in Stage 3: Style Direction — your job is to define the visual language for the product through design tokens.

You have the approved brief and information architecture injected below. Use them as context for your design decisions.

<role>
You are not generating generic themes. You are crafting distinct visual personalities that each solve the design problem differently. Every token choice should have a rationale tied to the persona, the brand, or the product's competitive positioning.

You are opinionated about:
- Typography that serves readability AND personality
- Color palettes that create emotional resonance, not just aesthetic pleasantness  
- Spacing systems that match the content density needs
- The relationship between visual style and perceived product quality
</role>

<process>
1. ANALYZE the brief for style-relevant signals:
   - Brand keywords and personality (e.g., "fresh, approachable, not clinical")
   - Persona context (when/where/how they use the product)
   - Competitive differentiation (how should this FEEL different from competitors?)
   - Platform constraints (mobile-first has different typography needs than desktop-first)

2. GENERATE 3-4 STYLE DIRECTIONS, each as a complete token set.
   Each direction must:
   - Have a clear name and personality statement (e.g., "Warm & Grounded" or "Sharp & Efficient")
   - Solve the design problem from a different angle
   - Be internally consistent (colors, type, spacing, and radii should feel like one system)
   - Include responsive overrides for mobile
   - Pass WCAG AA contrast requirements (verify your foreground/background pairings)

   CRITICAL: Directions must be MEANINGFULLY different, not just recolored versions of each other.
   Differentiate across MULTIPLE dimensions:
   - One might be spacious with a serif heading font and rounded shapes
   - Another might be dense with a geometric sans and sharp corners
   - Another might be bold with high-contrast colors and heavy type weights
   
   Each direction should make a user say "those feel like different products" not "those are the same layout in different colors."

3. For each direction, GENERATE a React + Tailwind preview component that demonstrates:
   - A card component with image, heading, body text, and button
   - Typography hierarchy (heading, subheading, body, caption)
   - The color palette applied (primary, secondary, background, surface, text levels)
   - Spacing rhythm
   - Border and shadow treatment
   
   This preview should be a single self-contained React component that can be rendered directly.

4. PRESENT all directions with:
   - The name and personality statement
   - The complete token set (JSON)
   - The preview component (React/Tailwind)
   - Key rationale: why this direction fits the product
   - Tradeoffs: what this direction optimizes for vs. what it sacrifices
   - Your recommendation and why

5. After the user selects (or mixes), generate the final tailwind.config.js.
</process>

<token_set_format>
Generate token sets in this exact JSON format:

```json
{
  "meta": {
    "name": "Direction name",
    "personality": "2-3 sentence personality statement",
    "optimizes_for": "What this direction prioritizes",
    "tradeoff": "What it sacrifices"
  },
  "color": {
    "primitive": {
      "gray-50": "#fafafa",
      "gray-100": "#f5f5f5",
      "gray-200": "#e5e5e5",
      "gray-300": "#d4d4d4",
      "gray-400": "#a3a3a3",
      "gray-500": "#737373",
      "gray-600": "#525252",
      "gray-700": "#404040",
      "gray-800": "#262626",
      "gray-900": "#171717",
      "primary-50": "#...",
      "primary-100": "#...",
      "primary-200": "#...",
      "primary-300": "#...",
      "primary-400": "#...",
      "primary-500": "#...",
      "primary-600": "#...",
      "primary-700": "#...",
      "primary-800": "#...",
      "primary-900": "#...",
      "accent-500": "#...",
      "accent-600": "#...",
      "error-500": "#...",
      "success-500": "#...",
      "warning-500": "#..."
    },
    "semantic": {
      "background": "{gray-50 or custom}",
      "surface": "{white or custom}",
      "surface-elevated": "{white with shadow}",
      "primary": "{primary-600}",
      "primary-hover": "{primary-700}",
      "secondary": "{gray-600}",
      "accent": "{accent-500}",
      "text-primary": "{gray-900}",
      "text-secondary": "{gray-500}",
      "text-tertiary": "{gray-400}",
      "text-inverse": "{white}",
      "border": "{gray-200}",
      "border-strong": "{gray-300}",
      "error": "{error-500}",
      "success": "{success-500}",
      "warning": "{warning-500}"
    }
  },
  "typography": {
    "font_family": {
      "heading": "Font Name",
      "body": "Font Name",
      "mono": "Font Name"
    },
    "scale_ratio": 1.25,
    "base_size_px": 16,
    "sizes": {
      "xs": 12, "sm": 14, "base": 16, "lg": 20,
      "xl": 25, "2xl": 31, "3xl": 39, "4xl": 49
    },
    "line_heights": {
      "tight": 1.2,
      "snug": 1.35,
      "normal": 1.5,
      "relaxed": 1.65
    },
    "weights": {
      "normal": 400,
      "medium": 500,
      "semibold": 600,
      "bold": 700
    },
    "letter_spacing": {
      "tight": "-0.02em",
      "normal": "0",
      "wide": "0.02em",
      "wider": "0.05em"
    }
  },
  "spacing": {
    "base_unit": 8,
    "scale": [0, 2, 4, 8, 12, 16, 20, 24, 32, 40, 48, 64, 80, 96, 128]
  },
  "borders": {
    "radius": {
      "none": "0px",
      "sm": "4px",
      "md": "8px",
      "lg": "12px",
      "xl": "16px",
      "2xl": "24px",
      "full": "9999px"
    },
    "width": {
      "default": "1px",
      "thick": "2px"
    }
  },
  "shadows": {
    "none": "none",
    "sm": "...",
    "md": "...",
    "lg": "...",
    "xl": "..."
  },
  "motion": {
    "duration": {
      "fast": "100ms",
      "normal": "200ms",
      "slow": "300ms",
      "slower": "500ms"
    },
    "easing": {
      "default": "cubic-bezier(0.4, 0, 0.2, 1)",
      "in": "cubic-bezier(0.4, 0, 1, 1)",
      "out": "cubic-bezier(0, 0, 0.2, 1)",
      "bounce": "cubic-bezier(0.34, 1.56, 0.64, 1)"
    }
  },
  "responsive_overrides": {
    "mobile": {
      "typography": {
        "base_size_px": 15,
        "scale_ratio": 1.2
      },
      "spacing": {
        "scale": [0, 2, 4, 8, 12, 16, 20, 24, 32, 40, 48, 64]
      }
    }
  }
}
```
</token_set_format>

<preview_component_format>
For each style direction, generate a self-contained React component that demonstrates the visual system.

The component should:
- Import no external dependencies beyond React and Tailwind
- Use inline style overrides for the custom fonts (since we can't modify tailwind config in preview)
- Show a realistic slice of the product UI (e.g., a recipe card for a food app, a product card for e-commerce)
- Demonstrate the full type scale (heading, subhead, body, caption)
- Show primary button, secondary button, and input field styles
- Use the token colors explicitly
- Be renderable as-is in a React preview environment

Use CSS custom properties for the token values so the component is self-documenting:

```jsx
const tokens = {
  colors: { primary: '#...', background: '#...', ... },
  fonts: { heading: '...', body: '...' },
  spacing: { sm: 8, md: 16, lg: 24, ... },
  radius: { md: '8px', lg: '12px', ... },
};
```
</preview_component_format>

<tailwind_config_format>
After the user selects a direction, generate a complete tailwind.config.js:

```javascript
/** @type {import('tailwindcss').Config} */
module.exports = {
  theme: {
    extend: {
      colors: {
        primary: { 50: '#...', 100: '#...', /* ... */ 900: '#...' },
        accent: { 500: '#...', 600: '#...' },
        surface: '#...',
        background: '#...',
      },
      fontFamily: {
        heading: ['Font Name', 'sans-serif'],
        body: ['Font Name', 'sans-serif'],
      },
      fontSize: { /* from token scale */ },
      spacing: { /* from token scale */ },
      borderRadius: { /* from token radii */ },
      boxShadow: { /* from token shadows */ },
    },
  },
};
```
</tailwind_config_format>

<opinion_guidelines>
Push back when:
- The user wants a style that contradicts their brand keywords ("clinical and cold" for a product described as "warm and approachable")
- A color palette fails WCAG AA contrast and the user wants to keep it
- The user is choosing based on personal preference rather than product fit ("I like blue" vs "blue communicates trust, which matters for a financial product")
- A dense type scale is chosen for a mobile-first product where readability is critical

Advocate for:
- Accessibility — never compromise on contrast ratios
- Typography quality — good fonts are worth the effort
- Internal consistency — every token should feel like part of one system
- Personality — generic is worse than polarizing; a strong direction can always be softened

When the user wants to mix directions:
- Help them identify which dimensions to take from each (e.g., "palette from A, type from C, spacing from B")
- Check that the mixed result is still internally consistent
- Generate a new unified token set from the mix rather than just merging JSONs
</opinion_guidelines>

Remember: this is the moment where the product gets its visual identity. The user is making the single most impactful taste decision in the entire process. Take it seriously. Present each direction as a genuine design philosophy, not a color swatch."""


# ═══════════════════════════════════════════════════════════
# CONTEXT BUILDER
# ═══════════════════════════════════════════════════════════

def load_style_context(project_path: str) -> str:
    """Load brief + IA context for the style direction agent."""
    parts = []
    
    # Brief
    brief_path = os.path.join(project_path, "brief", "brief.json")
    if os.path.exists(brief_path):
        with open(brief_path) as f:
            brief = json.load(f)
        parts.append("<approved_brief>")
        parts.append(json.dumps(brief, indent=2))
        parts.append("</approved_brief>")
    
    # Features (for understanding content types)
    features_path = os.path.join(project_path, "brief", "features.json")
    if os.path.exists(features_path):
        with open(features_path) as f:
            features = json.load(f)
        parts.append("<approved_features>")
        parts.append(json.dumps(features, indent=2))
        parts.append("</approved_features>")
    
    # Responsive strategy
    strategy_path = os.path.join(project_path, "architecture", "responsive_strategy.json")
    if os.path.exists(strategy_path):
        with open(strategy_path) as f:
            strategy = json.load(f)
        parts.append("<responsive_strategy>")
        parts.append(json.dumps(strategy, indent=2))
        parts.append("</responsive_strategy>")
    
    # Approved screen list (names only, for preview component context)
    screens_dir = os.path.join(project_path, "architecture", "screens")
    if os.path.exists(screens_dir):
        screen_dirs = [d for d in sorted(os.listdir(screens_dir))
                       if os.path.isdir(os.path.join(screens_dir, d)) and not d.startswith('_')]
        if screen_dirs:
            parts.append(f"<approved_screens>")
            parts.append(f"Screens with approved IA: {', '.join(screen_dirs)}")
            parts.append(f"</approved_screens>")
    
    # Existing style work
    tokens_dir = os.path.join(project_path, "tokens")
    selected_path = os.path.join(tokens_dir, "selected.json")
    if os.path.exists(selected_path):
        with open(selected_path) as f:
            selected = json.load(f)
        parts.append("<selected_tokens>")
        parts.append(json.dumps(selected, indent=2))
        parts.append("</selected_tokens>")
    else:
        # Check for candidates
        rejected_dir = os.path.join(tokens_dir, "_rejected")
        if os.path.exists(rejected_dir):
            rejected = [f for f in os.listdir(rejected_dir) if f.endswith('.json')]
            if rejected:
                parts.append(f"<previous_candidates>")
                parts.append(f"Previously rejected directions: {', '.join(f.replace('.json','') for f in rejected)}")
                parts.append(f"</previous_candidates>")
    
    return "\n".join(parts)


# ═══════════════════════════════════════════════════════════
# TOKEN SET MANAGEMENT
# ═══════════════════════════════════════════════════════════

def save_token_candidate(project_path: str, name: str, token_set: dict):
    """Save a candidate token set."""
    tokens_dir = os.path.join(project_path, "tokens")
    os.makedirs(tokens_dir, exist_ok=True)
    
    safe_name = name.lower().replace(" ", "_").replace("&", "and")
    path = os.path.join(tokens_dir, f"candidate_{safe_name}.json")
    
    with open(path, 'w') as f:
        json.dump(token_set, f, indent=2)
    print(f"Saved token candidate: {name}")
    return path


def select_token_set(project_path: str, token_set: dict, name: str):
    """Save the selected token set and move others to _rejected."""
    tokens_dir = os.path.join(project_path, "tokens")
    rejected_dir = os.path.join(tokens_dir, "_rejected")
    os.makedirs(rejected_dir, exist_ok=True)
    
    # Move all existing candidates to _rejected
    for f in os.listdir(tokens_dir):
        if f.startswith("candidate_") and f.endswith(".json"):
            src = os.path.join(tokens_dir, f)
            dst = os.path.join(rejected_dir, f)
            os.rename(src, dst)
    
    # Save selected
    selected_path = os.path.join(tokens_dir, "selected.json")
    token_set["_selected_at"] = datetime.now().isoformat()
    token_set["_direction_name"] = name
    with open(selected_path, 'w') as f:
        json.dump(token_set, f, indent=2)
    
    print(f"✓ Selected style direction: {name}")
    return selected_path


def save_preview(project_path: str, direction_name: str, 
                 preview_jsx: str, preview_type: str = "jsx"):
    """Save a preview component for a style direction."""
    previews_dir = os.path.join(project_path, "tokens", "previews")
    os.makedirs(previews_dir, exist_ok=True)
    
    safe_name = direction_name.lower().replace(" ", "_").replace("&", "and")
    ext = "jsx" if preview_type == "jsx" else "html"
    path = os.path.join(previews_dir, f"{safe_name}_preview.{ext}")
    
    with open(path, 'w') as f:
        f.write(preview_jsx)
    
    print(f"Saved preview: {path}")
    return path


def generate_tailwind_config(token_set: dict) -> str:
    """Generate a tailwind.config.js from the selected token set."""
    colors = token_set.get("color", {})
    typography = token_set.get("typography", {})
    spacing = token_set.get("spacing", {})
    borders = token_set.get("borders", {})
    shadows = token_set.get("shadows", {})
    
    # Build color entries
    color_lines = []
    primitives = colors.get("primitive", {})
    
    # Group primitives by prefix
    color_groups = {}
    for key, value in primitives.items():
        parts = key.rsplit("-", 1)
        if len(parts) == 2 and parts[1].isdigit():
            group = parts[0]
            shade = parts[1]
            color_groups.setdefault(group, {})[shade] = value
        else:
            color_groups[key] = value
    
    for group, shades in color_groups.items():
        if isinstance(shades, dict):
            shade_str = ", ".join(f"'{k}': '{v}'" for k, v in sorted(shades.items(), key=lambda x: int(x[0]) if x[0].isdigit() else 0))
            color_lines.append(f"        '{group}': {{ {shade_str} }},")
        else:
            color_lines.append(f"        '{group}': '{shades}',")
    
    # Add semantic colors
    semantic = colors.get("semantic", {})
    for key, value in semantic.items():
        if not value.startswith("{"):
            color_lines.append(f"        '{key}': '{value}',")
    
    # Build spacing
    spacing_entries = []
    scale = spacing.get("scale", [])
    base = spacing.get("base_unit", 8)
    for val in scale:
        spacing_entries.append(f"        '{val}': '{val}px',")
    
    # Build radii
    radius_entries = []
    for name, value in borders.get("radius", {}).items():
        val = value if isinstance(value, str) else f"{value}px"
        radius_entries.append(f"        '{name}': '{val}',")
    
    # Build shadows
    shadow_entries = []
    for name, value in shadows.items():
        shadow_entries.append(f"        '{name}': '{value}',")
    
    # Font families
    fonts = typography.get("font_family", {})
    heading_font = fonts.get("heading", "Inter")
    body_font = fonts.get("body", "Inter")
    mono_font = fonts.get("mono", "JetBrains Mono")
    
    config = f"""/** @type {{import('tailwindcss').Config}} */
module.exports = {{
  theme: {{
    extend: {{
      colors: {{
{chr(10).join(color_lines)}
      }},
      fontFamily: {{
        heading: ['{heading_font}', 'sans-serif'],
        body: ['{body_font}', 'sans-serif'],
        mono: ['{mono_font}', 'monospace'],
      }},
      spacing: {{
{chr(10).join(spacing_entries)}
      }},
      borderRadius: {{
{chr(10).join(radius_entries)}
      }},
      boxShadow: {{
{chr(10).join(shadow_entries)}
      }},
    }},
  }},
  plugins: [],
}};
"""
    return config


def save_tailwind_config(project_path: str, token_set: dict):
    """Generate and save tailwind.config.js from selected tokens."""
    config = generate_tailwind_config(token_set)
    path = os.path.join(project_path, "tokens", "tailwind.config.js")
    with open(path, 'w') as f:
        f.write(config)
    print(f"Generated: {path}")
    return path


# ═══════════════════════════════════════════════════════════
# STAGE COMPLETION
# ═══════════════════════════════════════════════════════════

def complete_style(project_path: str):
    """Verify token set is selected and advance to Stage 4."""
    selected_path = os.path.join(project_path, "tokens", "selected.json")
    
    if not os.path.exists(selected_path):
        print("Cannot complete Stage 3 — no token set has been selected.")
        print("Use select_token_set() to approve a direction.")
        return False
    
    # Generate tailwind config if not already done
    config_path = os.path.join(project_path, "tokens", "tailwind.config.js")
    if not os.path.exists(config_path):
        with open(selected_path) as f:
            tokens = json.load(f)
        save_tailwind_config(project_path, tokens)
    
    import sys
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from scaffold import advance_stage, log_decision
    
    with open(selected_path) as f:
        tokens = json.load(f)
    
    direction_name = tokens.get("_direction_name", "unnamed")
    meta = tokens.get("meta", {})
    
    log_decision(project_path, "Style Direction Selected",
        f"## Selected Direction: {direction_name}\n\n"
        f"**Personality:** {meta.get('personality', 'N/A')}\n\n"
        f"**Optimizes for:** {meta.get('optimizes_for', 'N/A')}\n\n"
        f"**Tradeoff:** {meta.get('tradeoff', 'N/A')}\n\n"
        f"**Typography:** {tokens.get('typography', {}).get('font_family', {}).get('heading', '?')} / "
        f"{tokens.get('typography', {}).get('font_family', {}).get('body', '?')}\n\n"
        f"**Primary color:** {tokens.get('color', {}).get('primitive', {}).get('primary-600', '?')}\n\n"
        f"**Spacing base:** {tokens.get('spacing', {}).get('base_unit', '?')}px\n\n"
        f"Tailwind config generated."
    )
    
    advance_stage(project_path, "component_generation",
                  f"Style direction '{direction_name}' selected. Tailwind config generated.")
    
    print(f"\n✓ Stage 3 complete. Project advanced to: component_generation")
    return True


# ═══════════════════════════════════════════════════════════
# FULL PROMPT ASSEMBLY
# ═══════════════════════════════════════════════════════════

def get_full_prompt(project_path: str) -> str:
    """Assemble complete system prompt with project context."""
    prompt = STYLE_SYSTEM_PROMPT
    
    context = load_style_context(project_path)
    if context:
        prompt += f"\n\n<project_context>\n{context}\n</project_context>"
    
    return prompt


if __name__ == "__main__":
    import sys as _sys
    if len(_sys.argv) > 1:
        project_path = _sys.argv[1]
        print(f"Style direction prompt: {len(get_full_prompt(project_path))} chars")
        
        # Check if tokens already selected
        selected = os.path.join(project_path, "tokens", "selected.json")
        if os.path.exists(selected):
            with open(selected) as f:
                tokens = json.load(f)
            print(f"Selected direction: {tokens.get('_direction_name', '?')}")
        else:
            print("No direction selected yet.")
    else:
        print("Usage: python style.py <project_path>")