"""
Stage 5: Page Assembly & Refinement Agent
===========================================
Composes approved components into full pages following the IA structure.
Runs evaluation before presenting. Handles refinement requests.

The agent:
- Selects components from the registry for each screen
- Generates page-level layout code (CSS Grid/Flexbox)
- Renders at all breakpoints and runs the evaluation pipeline
- Presents pages with evaluation reports
- Handles iterative refinement
"""

import json
import os
from datetime import datetime


ASSEMBLY_SYSTEM_PROMPT = """You are a senior frontend engineer and production designer. You are in Stage 5: Page Assembly — your job is to compose the approved components into complete, production-ready React pages.

You have the full project context injected below: brief, IA structure, token set, Tailwind config, and the component registry.

<role>
You are assembling a real product, not creating mockups. Every page must:
- Use ONLY approved components from the registry (no creating new components here)
- Follow the approved IA structure exactly (element positions, hierarchy, nesting)
- Apply the approved token set via Tailwind classes for all page-level spacing and layout
- Work at both mobile and desktop breakpoints
- Handle real-world content scenarios (empty states, loading, overflow)

If a needed component doesn't exist in the registry, flag it — don't create one inline. It should go back to Stage 4 for proper generation and approval.
</role>

<process>
For each screen in the approved IA:

1. IDENTIFY which components from the registry map to which IA elements:
   - Match by type, tags, and "used_in" metadata
   - If multiple components of the same type exist, select the best fit and explain why
   - Flag any gaps (IA elements with no matching component)

2. GENERATE the page as a React component:
   - Import approved components
   - Use CSS Grid or Flexbox for page-level layout
   - Apply Tailwind spacing classes from the token set for margins between sections
   - Handle responsive layout using Tailwind breakpoint prefixes
   - Include page-level state management if needed (active filters, scroll position)

3. GENERATE both mobile and desktop versions:
   - If responsive strategy is "reflow": single component with responsive classes
   - If "divergent": separate mobile/desktop components with shared state
   - If "split": separate components for each step in the mobile flow

4. SELF-EVALUATE before presenting:
   - Does the page follow the IA structure?
   - Is every spacing value from the token scale?
   - Are all components imported from approved paths?
   - Does the layout work at 375px and 1280px?
   - Is there a clear visual hierarchy?
   - Can the user complete the critical flow from this screen?

5. PRESENT the page with:
   - The complete React code
   - Component mapping (which registry components are used where)
   - Any concerns or compromises made
   - Suggestions for refinement
</process>

<page_format>
```jsx
// pages/home/Home.jsx
// Screen: Home
// Responsive strategy: reflow
// Components used: ProductCard, SearchInput, CategoryChips, NavBar, TabBar

import React, { useState } from 'react';
import NavBar from './components/NavBar/NavBar';
import SearchInput from './components/SearchInput/SearchInput';
import CategoryChips from './components/CategoryChips/CategoryChips';
import ProductCard from './components/ProductCard/ProductCard';
import TabBar from './components/TabBar/TabBar';

export default function Home() {
  const [activeCategory, setActiveCategory] = useState('all');
  
  return (
    <div className="min-h-screen bg-background">
      <NavBar />
      
      {/* Search Section */}
      <div className="px-16 pt-12 pb-8">
        <SearchInput placeholder="Search recipes..." />
      </div>
      
      {/* Category Filters */}
      <div className="px-16 pb-16">
        <CategoryChips 
          active={activeCategory} 
          onChange={setActiveCategory} 
        />
      </div>
      
      {/* Product Grid */}
      <div className="px-16 pb-24">
        <div className="grid grid-cols-2 md:grid-cols-3 lg:grid-cols-4 gap-16">
          {/* ProductCard instances */}
        </div>
      </div>
      
      {/* Mobile Tab Bar */}
      <div className="fixed bottom-0 left-0 right-0 md:hidden">
        <TabBar />
      </div>
    </div>
  );
}
```
</page_format>

<evaluation_checklist>
Before presenting any page, verify:

STRUCTURE:
□ Page follows the approved IA element map
□ All IA elements are represented
□ Nesting matches the IA hierarchy (containers hold their children)
□ Navigation elements are present and correct

TOKENS:
□ No hardcoded colors — all from semantic color tokens
□ No hardcoded font sizes — all from type scale
□ No hardcoded spacing — all from spacing scale
□ Border radii from token set
□ Shadows from token set

COMPONENTS:
□ All components imported from approved registry paths
□ No inline component creation (flag gaps for Stage 4)
□ Components receive appropriate props
□ Default/fallback props handle missing data

RESPONSIVE:
□ Layout works at 375px mobile viewport
□ Layout works at 1280px desktop viewport
□ Grid columns adjust at breakpoints
□ Touch targets are ≥44px on mobile
□ Bottom tab bar appears on mobile only (if applicable)
□ Content doesn't overflow viewport

ACCESSIBILITY:
□ Semantic HTML structure (main, nav, section, article, aside)
□ Heading hierarchy (h1 > h2 > h3, no skipped levels)
□ All images have alt text via component props
□ Interactive elements are keyboard accessible
□ Focus states are visible

FLOW:
□ User can reach the next screen in the critical flow
□ Primary action is visually prominent
□ Error/empty states are handled
</evaluation_checklist>

<refinement_protocol>
When the user requests changes:
1. Clarify whether the change is to the page layout or to a component
2. If it's a component change → flag for Stage 4 revision, don't modify inline
3. If it's a layout change → implement directly if it stays within IA bounds
4. If it's an IA change → flag that this requires Stage 2 revision
5. After any change, re-run the self-evaluation checklist
6. Present the updated page with a diff summary
</refinement_protocol>

<opinion_guidelines>
Push back when:
- The user wants to add an element not in the IA (that's a Stage 2 change)
- A layout change would break the flow for another screen
- Spacing adjustments go outside the token scale
- A component is being stretched beyond its intended use case

Advocate for:
- Consistent page structure across screens
- Generous breathing room between sections (pages often feel better with more space than designers initially choose)
- Mobile experience as the primary design target
- Skeleton loading states for async content
- Meaningful empty states ("No recipes found. Try broadening your filters.")
</opinion_guidelines>"""


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

def load_assembly_context(project_path: str) -> str:
    """Load all context needed for page assembly."""
    parts = []
    
    # Brief (light — just product context)
    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("<brief_summary>")
        parts.append(f"Product: {brief.get('product_name', '?')}")
        parts.append(f"One-liner: {brief.get('one_liner', '?')}")
        if brief.get("personas"):
            p = brief["personas"][0]
            parts.append(f"Primary persona: {p.get('name', '?')} — {p.get('context', '?')}")
        parts.append("</brief_summary>")
    
    # Flows (for understanding navigation)
    flows_dir = os.path.join(project_path, "brief", "flows")
    if os.path.exists(flows_dir):
        parts.append("<user_flows>")
        for ff in sorted(os.listdir(flows_dir)):
            if ff.endswith('.json'):
                with open(os.path.join(flows_dir, ff)) as f:
                    flow = json.load(f)
                parts.append(json.dumps(flow, indent=2))
        parts.append("</user_flows>")
    
    # IA (full — needed for structure)
    screens_dir = os.path.join(project_path, "architecture", "screens")
    if os.path.exists(screens_dir):
        parts.append("<approved_ia>")
        for screen_name in sorted(os.listdir(screens_dir)):
            screen_dir = os.path.join(screens_dir, screen_name)
            if not os.path.isdir(screen_dir) or screen_name.startswith('_'):
                continue
            for f_name in sorted(os.listdir(screen_dir)):
                if f_name.endswith('.json') and not f_name.startswith('_'):
                    with open(os.path.join(screen_dir, f_name)) as f:
                        ia = json.load(f)
                    parts.append(f"\n### {f_name}")
                    parts.append(json.dumps(ia, indent=2))
        parts.append("</approved_ia>")
    
    # 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>")
    
    # Token set
    tokens_path = os.path.join(project_path, "tokens", "selected.json")
    if os.path.exists(tokens_path):
        with open(tokens_path) as f:
            tokens = json.load(f)
        parts.append("<approved_tokens>")
        parts.append(json.dumps(tokens, indent=2))
        parts.append("</approved_tokens>")
    
    # Tailwind config
    config_path = os.path.join(project_path, "tokens", "tailwind.config.js")
    if os.path.exists(config_path):
        with open(config_path) as f:
            config = f.read()
        parts.append("<tailwind_config>")
        parts.append(config)
        parts.append("</tailwind_config>")
    
    # Component registry (critical — this is what's available)
    registry_path = os.path.join(project_path, "components", "registry.json")
    if os.path.exists(registry_path):
        with open(registry_path) as f:
            registry = json.load(f)
        parts.append("<component_registry>")
        parts.append(json.dumps(registry, indent=2))
        parts.append("</component_registry>")
    
    # Existing pages (for consistency and to know what's done)
    pages_dir = os.path.join(project_path, "pages")
    if os.path.exists(pages_dir):
        existing = []
        for page_name in sorted(os.listdir(pages_dir)):
            page_dir = os.path.join(pages_dir, page_name)
            if os.path.isdir(page_dir):
                jsx_files = [f for f in os.listdir(page_dir) if f.endswith('.jsx')]
                if jsx_files:
                    existing.append(f"{page_name}: {', '.join(jsx_files)}")
        if existing:
            parts.append("<existing_pages>")
            parts.append("\n".join(existing))
            parts.append("</existing_pages>")
    
    return "\n".join(parts)


# ═══════════════════════════════════════════════════════════
# PAGE MANAGEMENT
# ═══════════════════════════════════════════════════════════

def save_page(project_path: str, page_name: str, code: str,
              component_mapping: dict = None, eval_report: dict = None):
    """Save an assembled page."""
    page_dir = os.path.join(project_path, "pages", page_name)
    os.makedirs(page_dir, exist_ok=True)
    os.makedirs(os.path.join(page_dir, "components"), exist_ok=True)
    os.makedirs(os.path.join(page_dir, "screenshots"), exist_ok=True)
    
    # Determine filename
    # Check responsive strategy for this screen
    strategy_path = os.path.join(project_path, "architecture", "responsive_strategy.json")
    is_mobile_variant = page_name.endswith("_mobile")
    
    if is_mobile_variant:
        filename = f"{page_name.replace('_mobile', '').title()}.mobile.jsx"
    else:
        filename = f"{page_name.title().replace('_', '')}.jsx"
    
    code_path = os.path.join(page_dir, filename)
    with open(code_path, 'w') as f:
        f.write(code)
    
    # Save component mapping
    if component_mapping:
        mapping_path = os.path.join(page_dir, f"{page_name}_components.json")
        with open(mapping_path, 'w') as f:
            json.dump(component_mapping, f, indent=2)
    
    # Save evaluation report
    if eval_report:
        eval_path = os.path.join(page_dir, f"{page_name}.eval.json")
        with open(eval_path, 'w') as f:
            json.dump(eval_report, f, indent=2)
    
    print(f"  ✓ Saved page: {filename}")
    return code_path


def save_page_revision(project_path: str, page_name: str, code: str,
                       revision_notes: str = ""):
    """Save a revised version of a page, archiving the previous."""
    page_dir = os.path.join(project_path, "pages", page_name)
    
    # Find current file
    jsx_files = [f for f in os.listdir(page_dir) if f.endswith('.jsx') and not f.startswith('_')]
    
    if jsx_files:
        # Archive current version
        rejected_dir = os.path.join(page_dir, "_rejected")
        os.makedirs(rejected_dir, exist_ok=True)
        
        for jf in jsx_files:
            current_path = os.path.join(page_dir, jf)
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            archive_name = f"{jf.replace('.jsx', '')}_{timestamp}.jsx"
            archive_path = os.path.join(rejected_dir, archive_name)
            
            header = f"// Superseded: {datetime.now().isoformat()}\n// Reason: {revision_notes}\n\n"
            with open(current_path) as f:
                old_code = f.read()
            with open(archive_path, 'w') as f:
                f.write(header + old_code)
    
    # Save new version
    return save_page(project_path, page_name, code)


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

def complete_assembly(project_path: str):
    """Verify all screens have pages and mark project as complete."""
    # Check which screens have pages
    screens_dir = os.path.join(project_path, "architecture", "screens")
    pages_dir = os.path.join(project_path, "pages")
    
    if not os.path.exists(screens_dir):
        print("No approved screens found.")
        return False
    
    screen_names = [d for d in os.listdir(screens_dir) 
                    if os.path.isdir(os.path.join(screens_dir, d)) and not d.startswith('_')]
    
    pages_with_code = set()
    if os.path.exists(pages_dir):
        for page_name in os.listdir(pages_dir):
            page_dir = os.path.join(pages_dir, page_name)
            if os.path.isdir(page_dir):
                has_jsx = any(f.endswith('.jsx') for f in os.listdir(page_dir))
                if has_jsx:
                    pages_with_code.add(page_name)
    
    missing = [s for s in screen_names if s not in pages_with_code]
    
    if missing:
        print(f"Cannot complete Stage 5 — {len(missing)} screens still need pages:")
        for s in missing:
            print(f"  ○ {s}")
        return False
    
    import sys
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from scaffold import advance_stage, log_decision
    
    log_decision(project_path, "Pages Complete",
        f"## All Pages Assembled\n\n"
        f"**{len(pages_with_code)} pages** assembled from approved components:\n\n" +
        "\n".join(f"- {p}" for p in sorted(pages_with_code))
    )
    
    advance_stage(project_path, "complete",
                  f"All {len(pages_with_code)} pages assembled. Project design complete.")
    
    print(f"\n✓ Stage 5 complete. Project design is finished.")
    print(f"  {len(pages_with_code)} pages assembled across {len(screen_names)} screens.")
    return True


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

def get_full_prompt(project_path: str) -> str:
    """Assemble complete system prompt with all project context."""
    prompt = ASSEMBLY_SYSTEM_PROMPT
    context = load_assembly_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]
        prompt = get_full_prompt(project_path)
        print(f"Assembly agent prompt: {len(prompt)} chars")
        
        # Show page coverage
        screens_dir = os.path.join(project_path, "architecture", "screens")
        pages_dir = os.path.join(project_path, "pages")
        
        if os.path.exists(screens_dir):
            screens = [d for d in os.listdir(screens_dir) if not d.startswith('_')]
            print(f"Screens to assemble: {len(screens)}")
            for s in screens:
                page_dir = os.path.join(pages_dir, s)
                has_page = os.path.exists(page_dir) and any(
                    f.endswith('.jsx') for f in os.listdir(page_dir)) if os.path.exists(page_dir) else False
                status = "✓" if has_page else "○"
                print(f"  {status} {s}")
    else:
        print("Usage: python assembly.py <project_path>")