"""
Stage 2: Information Architecture Agent
=========================================
Reads the approved brief from Stage 1 and proposes screen-by-screen
information architecture as element maps (desktop + mobile pairs).

The agent:
- Extracts screens from the flow definitions
- Generates 2-3 IA variants per screen as element maps
- Renders blocked views for user evaluation
- Manages responsive strategy (reflow, divergent, split)
- Requires user approval for each screen before advancing
"""

import json
import os
from datetime import datetime


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

IA_SYSTEM_PROMPT = """You are a senior UX architect working as a design partner. You are in Stage 2: Information Architecture — your job is to define what goes on each screen, how it's structured, and how screens connect to each other.

You have the approved brief from Stage 1 injected below. Use it as your source of truth for personas, features, and user flows.

<role>
You are not a wireframe-generating machine. You are an opinionated architect who:
- Understands that IA decisions are product decisions disguised as layout choices
- Pushes back when a proposed structure won't serve the user's primary task
- Thinks in terms of user attention, cognitive load, and task completion speed
- Considers both desktop and mobile as first-class targets from the start
- Explains WHY a structure works, not just WHAT it contains
</role>

<process>
For each screen identified in the user flows:

1. PROPOSE 2-3 IA VARIANTS as element maps — each with a different structural approach:
   - Variant A might use a scrolling single page with clear sections
   - Variant B might use tabs or segmented controls to separate concerns  
   - Variant C might use progressive disclosure (show less, reveal on demand)
   
2. For EACH variant, produce BOTH a desktop and mobile element map.
   Consider how the structure adapts:
   
   | Desktop pattern         | Common mobile adaptation              |
   |------------------------|---------------------------------------|
   | Multi-column grid       | Single column stack or horizontal scroll |
   | Sidebar navigation      | Bottom tab bar + hamburger drawer      |
   | Data table              | Stacked cards or accordion             |
   | Comparison layout       | Tabbed single-item view with swipe     |
   | Persistent filters      | Bottom sheet or modal filters          |
   | Side-by-side panels     | Full-screen with back navigation       |

3. RECOMMEND one variant with clear reasoning tied to the personas and flows.

4. CLASSIFY the responsive strategy for each screen:
   - "reflow": Same IA, different layout (grid becomes stack)
   - "divergent": Different IA per breakpoint (desktop table → mobile cards)  
   - "split": One desktop screen becomes multiple mobile steps
</process>

<element_map_format>
Produce element maps in this exact JSON format. This format feeds directly into our blocked view rendering pipeline.

CRITICAL RULES:
- Estimate coordinates for a 375×812 viewport (mobile) or 1280×800 viewport (desktop)
- Every element MUST have x, y, w, h as integers
- Types must be one of: nav, container, card, button, text_heading, text_body, text_label, image, icon, input, divider, toggle, badge, chip, tab, list_item
- Use depth and parent to express nesting (0-indexed parent references)
- Labels should be descriptive enough that someone can understand the screen without seeing it rendered

```json
{
  "screen_name": "home",
  "viewport": "mobile",
  "image_width": 375,
  "image_height": 812,
  "elements": [
    {
      "type": "nav",
      "x": 0, "y": 0, "w": 375, "h": 56,
      "label": "Top navigation bar with search and profile",
      "depth": 0,
      "parent": null
    },
    {
      "type": "text_heading",
      "x": 16, "y": 12, "w": 200, "h": 28,
      "label": "App title / section heading",
      "depth": 1,
      "parent": 0
    }
  ]
}
```
</element_map_format>

<opinion_guidelines>
Push back when:
- A screen tries to do too much (more than one primary action)
- Critical content is below the fold on mobile
- Navigation patterns are inconsistent across screens
- The structure forces unnecessary taps to reach key content
- Filter/sort is buried when the content list is large
- There's no clear visual hierarchy (everything feels equally important)

Advocate for:
- The user's primary task being completable with minimum friction
- Consistent navigation patterns across all screens
- Mobile users who are often in interrupted, one-handed contexts
- Progressive disclosure over overwhelming upfront density
- Clear "what do I do here?" signals on every screen

When the user disagrees:
1. State your concern with specifics: "Putting filters behind a hamburger menu means users need 2 taps to narrow results. On a screen with potentially hundreds of items, that friction adds up."
2. Offer a compromise: "What if we show the top 3 filters inline and put the rest in an expandable panel?"
3. If they insist, commit fully: "Got it — I'll make the hamburger approach work well by [specific accommodation]."
4. Log the disagreement for the changelog.
</opinion_guidelines>

<output_expectations>
When presenting IA variants:

1. Start with a brief analysis of the screen's purpose and key decisions
2. Present each variant with:
   - A name and 1-sentence description
   - The desktop element map (JSON)
   - The mobile element map (JSON)
   - The responsive strategy classification
   - Pros and cons specific to the persona and flow
3. State your recommendation and why
4. Ask for the user's choice

When the user approves, confirm what was selected and flag any implications for connected screens.
</output_expectations>"""


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

def load_brief_context(project_path: str) -> str:
    """Load all Stage 1 artifacts and format them as context for the IA 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
    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>")
    
    # Flows
    flows_dir = os.path.join(project_path, "brief", "flows")
    if os.path.exists(flows_dir):
        flow_files = [f for f in os.listdir(flows_dir) if f.endswith('.json')]
        if flow_files:
            parts.append("<approved_flows>")
            for ff in sorted(flow_files):
                with open(os.path.join(flows_dir, ff)) as f:
                    flow = json.load(f)
                parts.append(f"### Flow: {flow.get('flow_name', ff)}")
                parts.append(json.dumps(flow, indent=2))
            parts.append("</approved_flows>")
    
    # Existing IA work (for resuming)
    arch_dir = os.path.join(project_path, "architecture")
    screens_dir = os.path.join(arch_dir, "screens")
    if os.path.exists(screens_dir):
        screen_dirs = [d for d in os.listdir(screens_dir) 
                       if os.path.isdir(os.path.join(screens_dir, d))]
        if screen_dirs:
            parts.append("<existing_ia_work>")
            for sd in sorted(screen_dirs):
                screen_path = os.path.join(screens_dir, sd)
                approved_files = [f for f in os.listdir(screen_path) 
                                 if f.endswith('.json') and not f.startswith('_')]
                rejected_dir = os.path.join(screen_path, "_rejected")
                rejected_count = len([f for f in os.listdir(rejected_dir) 
                                     if f.endswith('.json')]) if os.path.exists(rejected_dir) else 0
                
                status = "approved" if any("desktop" in f for f in approved_files) else "pending"
                parts.append(f"Screen '{sd}': {status} ({len(approved_files)} approved, {rejected_count} rejected)")
            parts.append("</existing_ia_work>")
    
    # Responsive strategy if it exists
    strategy_path = os.path.join(arch_dir, "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>")
    
    return "\n".join(parts)


def extract_screens_from_flows(project_path: str) -> list[dict]:
    """Extract unique screen names from all flows, preserving order and context."""
    flows_dir = os.path.join(project_path, "brief", "flows")
    if not os.path.exists(flows_dir):
        return []
    
    screens = {}  # name → metadata
    
    for ff in sorted(os.listdir(flows_dir)):
        if not ff.endswith('.json'):
            continue
        with open(os.path.join(flows_dir, ff)) as f:
            flow = json.load(f)
        
        flow_name = flow.get("flow_name", ff.replace(".json", ""))
        
        for step in flow.get("steps", []):
            screen_name = step.get("screen", "")
            if not screen_name:
                continue
            
            safe_name = screen_name.lower().replace(" ", "_").replace("/", "_")
            
            if safe_name not in screens:
                screens[safe_name] = {
                    "name": screen_name,
                    "safe_name": safe_name,
                    "appears_in_flows": [],
                    "user_actions": [],
                    "system_responses": [],
                }
            
            entry = screens[safe_name]
            if flow_name not in entry["appears_in_flows"]:
                entry["appears_in_flows"].append(flow_name)
            
            action = step.get("user_action", "")
            if action and action not in entry["user_actions"]:
                entry["user_actions"].append(action)
            
            response = step.get("system_response", "")
            if response and response not in entry["system_responses"]:
                entry["system_responses"].append(response)
    
    return list(screens.values())


# ═══════════════════════════════════════════════════════════
# IA STATE MANAGEMENT
# ═══════════════════════════════════════════════════════════

def create_ia_state(project_path: str) -> dict:
    """Initialize the IA conversation state."""
    screens = extract_screens_from_flows(project_path)
    
    state = {
        "screens_identified": [s["safe_name"] for s in screens],
        "screen_details": {s["safe_name"]: s for s in screens},
        "screens_approved": [],
        "screens_pending": [s["safe_name"] for s in screens],
        "responsive_strategy": {},
        "started": datetime.now().isoformat(),
    }
    
    state_path = os.path.join(project_path, "architecture", "_ia_state.json")
    with open(state_path, 'w') as f:
        json.dump(state, f, indent=2)
    
    return state


def get_ia_state(project_path: str) -> dict:
    """Load current IA state."""
    state_path = os.path.join(project_path, "architecture", "_ia_state.json")
    if os.path.exists(state_path):
        with open(state_path) as f:
            return json.load(f)
    return None


def update_ia_state(project_path: str, updates: dict) -> dict:
    """Update IA state."""
    state_path = os.path.join(project_path, "architecture", "_ia_state.json")
    with open(state_path) as f:
        state = json.load(f)
    state.update(updates)
    state["last_updated"] = datetime.now().isoformat()
    with open(state_path, 'w') as f:
        json.dump(state, f, indent=2)
    return state


# ═══════════════════════════════════════════════════════════
# SCREEN IA MANAGEMENT
# ═══════════════════════════════════════════════════════════

def save_screen_ia(project_path: str, screen_name: str, 
                   desktop_map: dict, mobile_map: dict,
                   responsive_strategy: str = "reflow",
                   notes: str = ""):
    """Save approved IA for a screen (desktop + mobile element maps)."""
    safe_name = screen_name.lower().replace(" ", "_")
    screen_dir = os.path.join(project_path, "architecture", "screens", safe_name)
    os.makedirs(screen_dir, exist_ok=True)
    os.makedirs(os.path.join(screen_dir, "_rejected"), exist_ok=True)
    
    # Save element maps
    desktop_path = os.path.join(screen_dir, f"{safe_name}_desktop.json")
    mobile_path = os.path.join(screen_dir, f"{safe_name}_mobile.json")
    
    with open(desktop_path, 'w') as f:
        json.dump(desktop_map, f, indent=2)
    with open(mobile_path, 'w') as f:
        json.dump(mobile_map, f, indent=2)
    
    # Render blocked views
    try:
        import sys
        eval_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'eval')
        sys.path.insert(0, eval_dir)
        from two_pass import render_blocked_from_map, render_blocked_with_spacing
        
        desktop_blocked = os.path.join(screen_dir, f"{safe_name}_desktop_blocked.png")
        mobile_blocked = os.path.join(screen_dir, f"{safe_name}_mobile_blocked.png")
        desktop_spaced = os.path.join(screen_dir, f"{safe_name}_desktop_spaced.png")
        mobile_spaced = os.path.join(screen_dir, f"{safe_name}_mobile_spaced.png")
        
        render_blocked_from_map(desktop_map, desktop_blocked)
        render_blocked_with_spacing(desktop_map, desktop_spaced)
        render_blocked_from_map(mobile_map, mobile_blocked)
        render_blocked_with_spacing(mobile_map, mobile_spaced)
        
        print(f"  Rendered blocked views for {screen_name}")
    except ImportError:
        print(f"  Warning: Could not import renderer. Blocked views not generated.")
    
    # Update IA state
    state = get_ia_state(project_path)
    if state:
        if safe_name in state.get("screens_pending", []):
            state["screens_pending"].remove(safe_name)
        if safe_name not in state.get("screens_approved", []):
            state["screens_approved"].append(safe_name)
        state["responsive_strategy"][safe_name] = {
            "strategy": responsive_strategy,
            "notes": notes,
        }
        update_ia_state(project_path, state)
    
    print(f"  ✓ Saved IA for screen: {screen_name} ({responsive_strategy})")


def reject_screen_variant(project_path: str, screen_name: str, 
                          variant_name: str, variant_map: dict, reason: str = ""):
    """Save a rejected IA variant for reference."""
    safe_name = screen_name.lower().replace(" ", "_")
    rejected_dir = os.path.join(project_path, "architecture", "screens", safe_name, "_rejected")
    os.makedirs(rejected_dir, exist_ok=True)
    
    safe_variant = variant_name.lower().replace(" ", "_")
    path = os.path.join(rejected_dir, f"{safe_variant}.json")
    
    variant_map["_rejection_reason"] = reason
    variant_map["_rejected_at"] = datetime.now().isoformat()
    
    with open(path, 'w') as f:
        json.dump(variant_map, f, indent=2)
    
    print(f"  Archived rejected variant: {variant_name} for {screen_name}")


def save_responsive_strategy(project_path: str, strategy: dict):
    """Save the overall responsive strategy document."""
    path = os.path.join(project_path, "architecture", "responsive_strategy.json")
    with open(path, 'w') as f:
        json.dump(strategy, f, indent=2)
    print(f"Saved responsive strategy")


def save_sitemap(project_path: str, sitemap: dict):
    """Save the screen navigation graph."""
    path = os.path.join(project_path, "architecture", "sitemap.json")
    with open(path, 'w') as f:
        json.dump(sitemap, f, indent=2)
    print(f"Saved sitemap")


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

def complete_ia(project_path: str):
    """Verify all screens are approved and advance to Stage 3."""
    state = get_ia_state(project_path)
    
    if state and state.get("screens_pending"):
        pending = state["screens_pending"]
        print(f"Cannot complete Stage 2 — {len(pending)} screens still pending:")
        for s in pending:
            print(f"  ○ {s}")
        return False
    
    approved = state.get("screens_approved", []) if state else []
    
    import sys
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from scaffold import advance_stage, log_decision
    
    strategies = state.get("responsive_strategy", {}) if state else {}
    strategy_summary = []
    for screen, info in strategies.items():
        s = info.get("strategy", "?") if isinstance(info, dict) else info
        strategy_summary.append(f"- {screen}: {s}")
    
    log_decision(project_path, "IA Approved",
        f"## Information Architecture Complete\n\n"
        f"### Approved Screens ({len(approved)})\n" +
        "\n".join(f"- {s}" for s in approved) +
        f"\n\n### Responsive Strategies\n" +
        "\n".join(strategy_summary)
    )
    
    advance_stage(project_path, "style_direction",
                  f"IA complete. {len(approved)} screens approved with responsive strategies defined.")
    
    print(f"\n✓ Stage 2 complete. Project advanced to: style_direction")
    return True


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

def get_full_prompt(project_path: str) -> str:
    """Assemble the complete system prompt for the IA agent with project context."""
    prompt = IA_SYSTEM_PROMPT
    
    # Inject brief context
    brief_context = load_brief_context(project_path)
    if brief_context:
        prompt += f"\n\n<project_context>\n{brief_context}\n</project_context>"
    
    # Inject screen inventory
    screens = extract_screens_from_flows(project_path)
    if screens:
        prompt += "\n\n<screens_to_design>"
        prompt += f"\nThe user flows reference {len(screens)} unique screens:\n"
        for s in screens:
            prompt += f"\n### {s['name']}"
            prompt += f"\n  Appears in: {', '.join(s['appears_in_flows'])}"
            if s['user_actions']:
                prompt += f"\n  User actions: {'; '.join(s['user_actions'])}"
            if s['system_responses']:
                prompt += f"\n  System shows: {'; '.join(s['system_responses'])}"
        prompt += "\n</screens_to_design>"
    
    # Inject IA progress
    state = get_ia_state(project_path)
    if state:
        approved = state.get("screens_approved", [])
        pending = state.get("screens_pending", [])
        if approved or pending:
            prompt += "\n\n<ia_progress>"
            if approved:
                prompt += f"\nApproved screens: {', '.join(approved)}"
            if pending:
                prompt += f"\nPending screens: {', '.join(pending)}"
                prompt += f"\nNext screen to work on: {pending[0]}"
            prompt += "\n</ia_progress>"
    
    return prompt


if __name__ == "__main__":
    import sys as _sys
    if len(_sys.argv) > 1:
        project_path = _sys.argv[1]
        
        # Show screen inventory
        screens = extract_screens_from_flows(project_path)
        print(f"Screens identified from flows: {len(screens)}")
        for s in screens:
            print(f"  {s['name']} — flows: {', '.join(s['appears_in_flows'])}")
            for a in s['user_actions'][:3]:
                print(f"    action: {a}")
        
        print(f"\nFull prompt length: {len(get_full_prompt(project_path))} chars")
    else:
        print("Usage: python ia.py <project_path>")