app.py

"""
Design System Extractor v2 — Main Application
==============================================

Flow:
1. User enters URL
2. Agent 1 discovers pages → User confirms
3. Agent 1 extracts tokens (Desktop + Mobile)
4. Agent 2 normalizes tokens
5. Stage 1 UI: User reviews tokens (accept/reject, Desktop↔Mobile toggle)
6. Agent 3 proposes upgrades
7. Stage 2 UI: User selects options with live preview
8. Agent 4 generates JSON
9. Stage 3 UI: User exports
"""

import os
import asyncio
import json
import gradio as gr
from datetime import datetime
from typing import Optional

# Get HF token from environment
HF_TOKEN_FROM_ENV = os.getenv("HF_TOKEN", "")

# =============================================================================
# GLOBAL STATE
# =============================================================================

class AppState:
    """Global application state."""
    def __init__(self):
        self.reset()
    
    def reset(self):
        self.discovered_pages = []
        self.base_url = ""
        self.desktop_raw = None  # ExtractedTokens
        self.mobile_raw = None   # ExtractedTokens
        self.desktop_normalized = None  # NormalizedTokens
        self.mobile_normalized = None   # NormalizedTokens
        self.upgrade_recommendations = None  # UpgradeRecommendations
        self.selected_upgrades = {}  # User selections
        self.logs = []
    
    def log(self, message: str):
        timestamp = datetime.now().strftime("%H:%M:%S")
        self.logs.append(f"[{timestamp}] {message}")
        if len(self.logs) > 100:
            self.logs.pop(0)
    
    def get_logs(self) -> str:
        return "\n".join(self.logs)

state = AppState()


# =============================================================================
# LAZY IMPORTS
# =============================================================================

def get_crawler():
    import agents.crawler
    return agents.crawler

def get_extractor():
    import agents.extractor
    return agents.extractor

def get_normalizer():
    import agents.normalizer
    return agents.normalizer

def get_advisor():
    import agents.advisor
    return agents.advisor

def get_schema():
    import core.token_schema
    return core.token_schema


# =============================================================================
# PHASE 1: DISCOVER PAGES
# =============================================================================

async def discover_pages(url: str, progress=gr.Progress()):
    """Discover pages from URL."""
    state.reset()
    
    if not url or not url.startswith(("http://", "https://")):
        return "❌ Please enter a valid URL", "", None
    
    state.log(f"🚀 Starting discovery for: {url}")
    progress(0.1, desc="🔍 Discovering pages...")
    
    try:
        crawler = get_crawler()
        discoverer = crawler.PageDiscoverer()
        
        pages = await discoverer.discover(url)
        
        state.discovered_pages = pages
        state.base_url = url
        
        state.log(f"✅ Found {len(pages)} pages")
        
        # Format for display
        pages_data = []
        for page in pages:
            pages_data.append([
                True,  # Selected by default
                page.url,
                page.title if page.title else "(No title)",
                page.page_type.value,
                "✓" if not page.error else f"⚠ {page.error}"
            ])
        
        progress(1.0, desc="✅ Discovery complete!")
        
        status = f"✅ Found {len(pages)} pages. Review and click 'Extract Tokens' to continue."
        
        return status, state.get_logs(), pages_data
        
    except Exception as e:
        import traceback
        state.log(f"❌ Error: {str(e)}")
        return f"❌ Error: {str(e)}", state.get_logs(), None


# =============================================================================
# PHASE 2: EXTRACT TOKENS
# =============================================================================

async def extract_tokens(pages_data, progress=gr.Progress()):
    """Extract tokens from selected pages (both viewports)."""
    
    state.log(f"📥 Received pages_data type: {type(pages_data)}")
    
    if pages_data is None:
        return "❌ Please discover pages first", state.get_logs(), None, None
    
    # Get selected URLs - handle pandas DataFrame
    selected_urls = []
    
    try:
        # Check if it's a pandas DataFrame
        if hasattr(pages_data, 'iterrows'):
            state.log(f"📥 DataFrame with {len(pages_data)} rows, columns: {list(pages_data.columns)}")
            
            for idx, row in pages_data.iterrows():
                # Get values by column name or position
                try:
                    # Try column names first
                    is_selected = row.get('Select', row.iloc[0] if len(row) > 0 else False)
                    url = row.get('URL', row.iloc[1] if len(row) > 1 else '')
                except:
                    # Fallback to positional
                    is_selected = row.iloc[0] if len(row) > 0 else False
                    url = row.iloc[1] if len(row) > 1 else ''
                
                if is_selected and url:
                    selected_urls.append(url)
        
        # If it's a dict (Gradio sometimes sends this)
        elif isinstance(pages_data, dict):
            state.log(f"📥 Dict with keys: {list(pages_data.keys())}")
            data = pages_data.get('data', [])
            for row in data:
                if isinstance(row, (list, tuple)) and len(row) >= 2 and row[0]:
                    selected_urls.append(row[1])
        
        # If it's a list
        elif isinstance(pages_data, (list, tuple)):
            state.log(f"📥 List with {len(pages_data)} items")
            for row in pages_data:
                if isinstance(row, (list, tuple)) and len(row) >= 2 and row[0]:
                    selected_urls.append(row[1])
                    
    except Exception as e:
        state.log(f"❌ Error parsing pages_data: {str(e)}")
        import traceback
        state.log(traceback.format_exc())
    
    state.log(f"📋 Found {len(selected_urls)} selected URLs")
    
    # If still no URLs, try using stored discovered pages
    if not selected_urls and state.discovered_pages:
        state.log("⚠️ No URLs from table, using all discovered pages")
        selected_urls = [p.url for p in state.discovered_pages if not p.error][:10]
    
    if not selected_urls:
        return "❌ No pages selected. Please select pages or rediscover.", state.get_logs(), None, None
    
    # Limit to 10 pages for performance
    selected_urls = selected_urls[:10]
    
    state.log(f"📋 Extracting from {len(selected_urls)} pages:")
    for url in selected_urls[:3]:
        state.log(f"   • {url}")
    if len(selected_urls) > 3:
        state.log(f"   ... and {len(selected_urls) - 3} more")
    
    progress(0.05, desc="🚀 Starting extraction...")
    
    try:
        schema = get_schema()
        extractor_mod = get_extractor()
        normalizer_mod = get_normalizer()
        
        # === DESKTOP EXTRACTION ===
        state.log("")
        state.log("🖥️ DESKTOP EXTRACTION (1440px)")
        progress(0.1, desc="🖥️ Extracting desktop tokens...")
        
        desktop_extractor = extractor_mod.TokenExtractor(viewport=schema.Viewport.DESKTOP)
        
        def desktop_progress(p):
            progress(0.1 + (p * 0.35), desc=f"🖥️ Desktop... {int(p*100)}%")
        
        state.desktop_raw = await desktop_extractor.extract(selected_urls, progress_callback=desktop_progress)
        
        state.log(f"   Raw: {len(state.desktop_raw.colors)} colors, {len(state.desktop_raw.typography)} typography, {len(state.desktop_raw.spacing)} spacing")
        
        # Normalize desktop
        state.log("   Normalizing...")
        state.desktop_normalized = normalizer_mod.normalize_tokens(state.desktop_raw)
        state.log(f"   Normalized: {len(state.desktop_normalized.colors)} colors, {len(state.desktop_normalized.typography)} typography, {len(state.desktop_normalized.spacing)} spacing")
        
        # === MOBILE EXTRACTION ===
        state.log("")
        state.log("📱 MOBILE EXTRACTION (375px)")
        progress(0.5, desc="📱 Extracting mobile tokens...")
        
        mobile_extractor = extractor_mod.TokenExtractor(viewport=schema.Viewport.MOBILE)
        
        def mobile_progress(p):
            progress(0.5 + (p * 0.35), desc=f"📱 Mobile... {int(p*100)}%")
        
        state.mobile_raw = await mobile_extractor.extract(selected_urls, progress_callback=mobile_progress)
        
        state.log(f"   Raw: {len(state.mobile_raw.colors)} colors, {len(state.mobile_raw.typography)} typography, {len(state.mobile_raw.spacing)} spacing")
        
        # Normalize mobile
        state.log("   Normalizing...")
        state.mobile_normalized = normalizer_mod.normalize_tokens(state.mobile_raw)
        state.log(f"   Normalized: {len(state.mobile_normalized.colors)} colors, {len(state.mobile_normalized.typography)} typography, {len(state.mobile_normalized.spacing)} spacing")
        
        progress(0.95, desc="📊 Preparing results...")
        
        # Format results for Stage 1 UI
        desktop_data = format_tokens_for_display(state.desktop_normalized)
        mobile_data = format_tokens_for_display(state.mobile_normalized)
        
        state.log("")
        state.log("=" * 50)
        state.log("✅ EXTRACTION COMPLETE!")
        state.log("=" * 50)
        
        progress(1.0, desc="✅ Complete!")
        
        status = f"""## ✅ Extraction Complete!

| Viewport | Colors | Typography | Spacing |
|----------|--------|------------|---------|
| Desktop | {len(state.desktop_normalized.colors)} | {len(state.desktop_normalized.typography)} | {len(state.desktop_normalized.spacing)} |
| Mobile | {len(state.mobile_normalized.colors)} | {len(state.mobile_normalized.typography)} | {len(state.mobile_normalized.spacing)} |

**Next:** Review the tokens below. Accept or reject, then proceed to Stage 2.
"""
        
        return status, state.get_logs(), desktop_data, mobile_data
        
    except Exception as e:
        import traceback
        state.log(f"❌ Error: {str(e)}")
        state.log(traceback.format_exc())
        return f"❌ Error: {str(e)}", state.get_logs(), None, None


def format_tokens_for_display(normalized) -> dict:
    """Format normalized tokens for Gradio display."""
    if normalized is None:
        return {"colors": [], "typography": [], "spacing": []}
    
    # Colors are now a dict
    colors = []
    color_items = list(normalized.colors.values()) if isinstance(normalized.colors, dict) else normalized.colors
    for c in sorted(color_items, key=lambda x: -x.frequency)[:50]:
        colors.append([
            True,  # Accept checkbox
            c.value,
            c.suggested_name or "",
            c.frequency,
            c.confidence.value if c.confidence else "medium",
            f"{c.contrast_white:.1f}:1" if c.contrast_white else "N/A",
            "✓" if c.wcag_aa_small_text else "✗",
            ", ".join(c.contexts[:2]) if c.contexts else "",
        ])
    
    # Typography
    typography = []
    typo_items = list(normalized.typography.values()) if isinstance(normalized.typography, dict) else normalized.typography
    for t in sorted(typo_items, key=lambda x: -x.frequency)[:30]:
        typography.append([
            True,  # Accept checkbox
            t.font_family,
            t.font_size,
            str(t.font_weight),
            t.line_height or "",
            t.suggested_name or "",
            t.frequency,
            t.confidence.value if t.confidence else "medium",
        ])
    
    # Spacing
    spacing = []
    spacing_items = list(normalized.spacing.values()) if isinstance(normalized.spacing, dict) else normalized.spacing
    for s in sorted(spacing_items, key=lambda x: x.value_px)[:20]:
        spacing.append([
            True,  # Accept checkbox
            s.value,
            f"{s.value_px}px",
            s.suggested_name or "",
            s.frequency,
            "✓" if s.fits_base_8 else "",
            s.confidence.value if s.confidence else "medium",
        ])
    
    return {
        "colors": colors,
        "typography": typography,
        "spacing": spacing,
    }


def switch_viewport(viewport: str):
    """Switch between desktop and mobile view."""
    if viewport == "Desktop (1440px)":
        data = format_tokens_for_display(state.desktop_normalized)
    else:
        data = format_tokens_for_display(state.mobile_normalized)
    
    return data["colors"], data["typography"], data["spacing"]


# =============================================================================
# STAGE 2: AI ANALYSIS
# =============================================================================

async def run_stage2_analysis(progress=gr.Progress()):
    """Run Agent 3 analysis on extracted tokens."""
    
    if not state.desktop_normalized or not state.mobile_normalized:
        return ("❌ Please complete Stage 1 first", "", "", None, None, None, "", "", "")
    
    state.log("")
    state.log("=" * 60)
    state.log("🧠 STAGE 2: AI-POWERED ANALYSIS")
    state.log("=" * 60)
    state.log("")
    
    # Log model info
    model_name = os.getenv("AGENT3_MODEL", "meta-llama/Llama-3.1-70B-Instruct")
    state.log("📦 LLM CONFIGURATION:")
    state.log(f"   Model: {model_name}")
    state.log(f"   Expertise: Design system reasoning, best practices comparison")
    state.log(f"   Task: Analyze tokens against Material, Apple, Polaris, Carbon, Atlassian")
    state.log("")
    
    progress(0.1, desc="🤖 Starting AI analysis...")
    
    try:
        advisor_mod = get_advisor()
        
        # Log what we're analyzing
        desktop_colors = len(state.desktop_normalized.colors)
        desktop_typo = len(state.desktop_normalized.typography)
        mobile_typo = len(state.mobile_normalized.typography)
        
        state.log("📊 INPUT DATA:")
        state.log(f"   Colors: {desktop_colors} (viewport-agnostic)")
        state.log(f"   Typography: {desktop_typo} desktop, {mobile_typo} mobile")
        state.log(f"   Spacing: {len(state.desktop_normalized.spacing)} values")
        state.log("")
        
        # Get detected font info
        fonts = get_detected_fonts()
        base_size = get_base_font_size()
        
        state.log(f"🔤 DETECTED FONT: {fonts.get('primary', 'Unknown')}")
        state.log(f"   Weights: {', '.join(map(str, fonts.get('weights', [])))}")
        state.log(f"   Base size: {base_size}px")
        state.log("")
        
        state.log("🔍 RESEARCHING TOP DESIGN SYSTEMS...")
        progress(0.2, desc="🔍 Researching brands...")
        
        recommendations = await advisor_mod.analyze_design_system(
            desktop_tokens=state.desktop_normalized,
            mobile_tokens=state.mobile_normalized,
            log_callback=state.log,
        )
        
        state.upgrade_recommendations = recommendations
        
        # Log brand analysis
        state.log("")
        state.log("📊 BRAND COMPARISON RESULTS:")
        for brand in recommendations.brand_analysis:
            state.log(f"   • {brand.get('brand', 'Unknown')}:")
            state.log(f"     Ratio: {brand.get('ratio', '?')}, Base: {brand.get('base', '?')}px, Grid: {brand.get('spacing', '?')}")
            if brand.get('notes'):
                state.log(f"     Notes: {brand.get('notes', '')[:100]}")
        
        state.log("")
        state.log("💡 LLM RECOMMENDATION:")
        if recommendations.llm_rationale:
            # Split into sentences for readability
            sentences = recommendations.llm_rationale.split('. ')
            for s in sentences[:5]:
                if s.strip():
                    state.log(f"   {s.strip()}.")
        
        if recommendations.color_observations:
            state.log("")
            state.log("🎨 COLOR ANALYSIS:")
            state.log(f"   {recommendations.color_observations[:200]}")
        
        if recommendations.accessibility_issues:
            state.log("")
            state.log("⚠️ ACCESSIBILITY CONCERNS:")
            for issue in recommendations.accessibility_issues[:3]:
                state.log(f"   • {issue}")
        
        progress(0.9, desc="📊 Preparing recommendations...")
        
        # Format brand comparison markdown
        brand_md = format_brand_comparison(recommendations)
        
        # Format typography with BOTH desktop and mobile
        typography_desktop_data = format_typography_comparison_viewport(
            state.desktop_normalized, base_size, "desktop"
        )
        typography_mobile_data = format_typography_comparison_viewport(
            state.mobile_normalized, base_size, "mobile"
        )
        
        # Format spacing comparison table
        spacing_data = format_spacing_comparison(recommendations)
        
        # Format color display: BASE colors + ramps separately
        base_colors_md = format_base_colors()
        color_ramps_md = format_color_ramps_visual(recommendations)
        
        # Format radius display (with token suggestions)
        radius_md = format_radius_with_tokens()
        
        # Format shadows display (with token suggestions)
        shadows_md = format_shadows_with_tokens()
        
        state.log("")
        state.log("=" * 60)
        state.log("✅ ANALYSIS COMPLETE!")
        state.log("=" * 60)
        progress(1.0, desc="✅ Complete!")
        
        # Build status with font info
        status = f"""## 🧠 AI Analysis Complete!

### Detected Font
**{fonts.get('primary', 'Unknown')}** — Weights: {', '.join(map(str, fonts.get('weights', [])))}

**Base Size:** {base_size}px (detected from body text)

### LLM Recommendation
{recommendations.llm_rationale if recommendations.llm_rationale else "Analysis based on rule-based comparison with industry design systems."}

{f"### Accessibility Notes{chr(10)}" + chr(10).join(['• ' + a for a in recommendations.accessibility_issues]) if recommendations.accessibility_issues else ""}
"""
        
        return (status, state.get_logs(), brand_md, 
                typography_desktop_data, typography_mobile_data, spacing_data, 
                base_colors_md, color_ramps_md, radius_md, shadows_md)
        
    except Exception as e:
        import traceback
        state.log(f"❌ Error: {str(e)}")
        state.log(traceback.format_exc())
        return (f"❌ Analysis failed: {str(e)}", state.get_logs(), "", None, None, None, "", "", "", "")


def get_detected_fonts() -> dict:
    """Get detected font information."""
    if not state.desktop_normalized:
        return {"primary": "Unknown", "weights": []}
    
    fonts = {}
    weights = set()
    
    for t in state.desktop_normalized.typography.values():
        family = t.font_family
        weight = t.font_weight
        
        if family not in fonts:
            fonts[family] = 0
        fonts[family] += t.frequency
        
        if weight:
            try:
                weights.add(int(weight))
            except:
                pass
    
    primary = max(fonts.items(), key=lambda x: x[1])[0] if fonts else "Unknown"
    
    return {
        "primary": primary,
        "weights": sorted(weights) if weights else [400],
        "all_fonts": fonts,
    }


def get_base_font_size() -> int:
    """Detect base font size from typography."""
    if not state.desktop_normalized:
        return 16
    
    # Find most common size in body range (14-18px)
    sizes = {}
    for t in state.desktop_normalized.typography.values():
        size_str = str(t.font_size).replace('px', '').replace('rem', '').replace('em', '')
        try:
            size = float(size_str)
            if 14 <= size <= 18:
                sizes[size] = sizes.get(size, 0) + t.frequency
        except:
            pass
    
    if sizes:
        return int(max(sizes.items(), key=lambda x: x[1])[0])
    return 16


def format_brand_comparison(recommendations) -> str:
    """Format brand comparison as markdown table."""
    if not recommendations.brand_analysis:
        return "*Brand analysis not available*"
    
    lines = [
        "### 📊 Design System Comparison (5 Top Brands)",
        "",
        "| Brand | Type Ratio | Base Size | Spacing | Notes |",
        "|-------|------------|-----------|---------|-------|",
    ]
    
    for brand in recommendations.brand_analysis[:5]:
        name = brand.get("brand", "Unknown")
        ratio = brand.get("ratio", "?")
        base = brand.get("base", "?")
        spacing = brand.get("spacing", "?")
        notes = brand.get("notes", "")[:50] + ("..." if len(brand.get("notes", "")) > 50 else "")
        lines.append(f"| {name} | {ratio} | {base}px | {spacing} | {notes} |")
    
    return "\n".join(lines)


def format_typography_comparison_viewport(normalized_tokens, base_size: int, viewport: str) -> list:
    """Format typography comparison for a specific viewport."""
    if not normalized_tokens:
        return []
    
    # Get current typography sorted by size
    current_typo = list(normalized_tokens.typography.values())
    
    # Parse and sort sizes
    def parse_size(t):
        size_str = str(t.font_size).replace('px', '').replace('rem', '').replace('em', '')
        try:
            return float(size_str)
        except:
            return 16
    
    current_typo.sort(key=lambda t: -parse_size(t))
    sizes = [parse_size(t) for t in current_typo]
    
    # Use detected base or default
    base = base_size if base_size else 16
    
    # Scale factors for mobile (typically 0.85-0.9 of desktop)
    mobile_factor = 0.875 if viewport == "mobile" else 1.0
    
    # Token names (13 levels)
    token_names = [
        "display.2xl", "display.xl", "display.lg", "display.md",
        "heading.xl", "heading.lg", "heading.md", "heading.sm",
        "body.lg", "body.md", "body.sm",
        "caption", "overline"
    ]
    
    # Generate scales - use base size and round to sensible values
    def round_to_even(val):
        """Round to even numbers for cleaner type scales."""
        return int(round(val / 2) * 2)
    
    scales = {
        "1.2": [round_to_even(base * mobile_factor * (1.2 ** (8-i))) for i in range(13)],
        "1.25": [round_to_even(base * mobile_factor * (1.25 ** (8-i))) for i in range(13)],
        "1.333": [round_to_even(base * mobile_factor * (1.333 ** (8-i))) for i in range(13)],
    }
    
    # Build comparison table
    data = []
    for i, name in enumerate(token_names):
        current = f"{int(sizes[i])}px" if i < len(sizes) else "—"
        s12 = f"{scales['1.2'][i]}px"
        s125 = f"{scales['1.25'][i]}px"
        s133 = f"{scales['1.333'][i]}px"
        keep = current
        data.append([name, current, s12, s125, s133, keep])
    
    return data


def format_base_colors() -> str:
    """Format base colors (detected) separately from ramps."""
    if not state.desktop_normalized:
        return "*No colors detected*"
    
    colors = list(state.desktop_normalized.colors.values())
    colors.sort(key=lambda c: -c.frequency)
    
    lines = [
        "### 🎨 Base Colors (Detected)",
        "",
        "These are the primary colors extracted from your website:",
        "",
        "| Color | Hex | Role | Frequency | Contrast |",
        "|-------|-----|------|-----------|----------|",
    ]
    
    for color in colors[:10]:
        hex_val = color.value
        role = "Primary" if color.suggested_name and "primary" in color.suggested_name.lower() else \
               "Text" if color.suggested_name and "text" in color.suggested_name.lower() else \
               "Background" if color.suggested_name and "background" in color.suggested_name.lower() else \
               "Border" if color.suggested_name and "border" in color.suggested_name.lower() else \
               "Accent"
        freq = f"{color.frequency:,}"
        contrast = f"{color.contrast_white:.1f}:1" if color.contrast_white else "—"
        
        # Create a simple color indicator
        lines.append(f"| 🟦 | `{hex_val}` | {role} | {freq} | {contrast} |")
    
    return "\n".join(lines)


def format_color_ramps_visual(recommendations) -> str:
    """Format color ramps with visual display showing all shades."""
    if not state.desktop_normalized:
        return "*No colors to display*"
    
    colors = list(state.desktop_normalized.colors.values())
    colors.sort(key=lambda c: -c.frequency)
    
    lines = [
        "### 🌈 Generated Color Ramps",
        "",
        "Full ramp (50-950) generated for each base color:",
        "",
    ]
    
    from core.color_utils import generate_color_ramp
    
    for color in colors[:6]:  # Top 6 colors
        hex_val = color.value
        role = color.suggested_name.split('.')[1] if color.suggested_name and '.' in color.suggested_name else "color"
        
        # Generate ramp
        try:
            ramp = generate_color_ramp(hex_val)
            
            lines.append(f"**{role.upper()}** (base: `{hex_val}`)")
            lines.append("")
            lines.append("| 50 | 100 | 200 | 300 | 400 | 500 | 600 | 700 | 800 | 900 |")
            lines.append("|---|---|---|---|---|---|---|---|---|---|")
            
            # Create row with hex values
            row = "|"
            for i in range(10):
                if i < len(ramp):
                    row += f" `{ramp[i]}` |"
                else:
                    row += " — |"
            lines.append(row)
            lines.append("")
            
        except Exception as e:
            lines.append(f"**{role}** (`{hex_val}`) — Could not generate ramp: {str(e)}")
            lines.append("")
    
    return "\n".join(lines)


def format_radius_with_tokens() -> str:
    """Format radius with token name suggestions."""
    if not state.desktop_normalized or not state.desktop_normalized.radius:
        return "*No border radius values detected.*"
    
    radii = list(state.desktop_normalized.radius.values())
    
    lines = [
        "### 🔘 Border Radius Tokens",
        "",
        "| Detected | Suggested Token | Usage |",
        "|----------|-----------------|-------|",
    ]
    
    # Sort by pixel value
    def parse_radius(r):
        val = str(r.value).replace('px', '').replace('%', '')
        try:
            return float(val)
        except:
            return 999
    
    radii.sort(key=lambda r: parse_radius(r))
    
    token_map = {
        (0, 2): ("radius.none", "Sharp corners"),
        (2, 4): ("radius.xs", "Subtle rounding"),
        (4, 6): ("radius.sm", "Small elements"),
        (6, 10): ("radius.md", "Buttons, cards"),
        (10, 16): ("radius.lg", "Modals, panels"),
        (16, 32): ("radius.xl", "Large containers"),
        (32, 100): ("radius.2xl", "Pill shapes"),
    }
    
    for r in radii[:8]:
        val = str(r.value)
        px = parse_radius(r)
        
        if "%" in str(r.value) or px >= 50:
            token = "radius.full"
            usage = "Circles, avatars"
        else:
            token = "radius.md"
            usage = "General use"
            for (low, high), (t, u) in token_map.items():
                if low <= px < high:
                    token = t
                    usage = u
                    break
        
        lines.append(f"| {val} | `{token}` | {usage} |")
    
    return "\n".join(lines)


def format_shadows_with_tokens() -> str:
    """Format shadows with token name suggestions."""
    if not state.desktop_normalized or not state.desktop_normalized.shadows:
        return "*No shadow values detected.*"
    
    shadows = list(state.desktop_normalized.shadows.values())
    
    lines = [
        "### 🌫️ Shadow Tokens",
        "",
        "| Detected Value | Suggested Token | Use Case |",
        "|----------------|-----------------|----------|",
    ]
    
    shadow_sizes = ["shadow.xs", "shadow.sm", "shadow.md", "shadow.lg", "shadow.xl", "shadow.2xl"]
    
    for i, s in enumerate(shadows[:6]):
        val = str(s.value)[:40] + ("..." if len(str(s.value)) > 40 else "")
        token = shadow_sizes[i] if i < len(shadow_sizes) else f"shadow.custom-{i}"
        
        # Guess use case based on index
        use_cases = ["Subtle elevation", "Cards, dropdowns", "Modals, dialogs", "Popovers", "Floating elements", "Dramatic effect"]
        use = use_cases[i] if i < len(use_cases) else "Custom"
        
        lines.append(f"| `{val}` | `{token}` | {use} |")
    
    return "\n".join(lines)


def format_spacing_comparison(recommendations) -> list:
    """Format spacing comparison table."""
    if not state.desktop_normalized:
        return []
    
    # Get current spacing
    current_spacing = list(state.desktop_normalized.spacing.values())
    current_spacing.sort(key=lambda s: s.value_px)
    
    data = []
    for s in current_spacing[:10]:
        current = f"{s.value_px}px"
        grid_8 = f"{snap_to_grid(s.value_px, 8)}px"
        grid_4 = f"{snap_to_grid(s.value_px, 4)}px"
        
        # Mark if value fits
        if s.value_px == snap_to_grid(s.value_px, 8):
            grid_8 += " ✓"
        if s.value_px == snap_to_grid(s.value_px, 4):
            grid_4 += " ✓"
        
        data.append([current, grid_8, grid_4])
    
    return data


def snap_to_grid(value: float, base: int) -> int:
    """Snap value to grid."""
    return round(value / base) * base


def apply_selected_upgrades(type_choice: str, spacing_choice: str, apply_ramps: bool):
    """Apply selected upgrade options."""
    if not state.upgrade_recommendations:
        return "❌ Run analysis first", ""
    
    state.log("✨ Applying selected upgrades...")
    
    # Store selections
    state.selected_upgrades = {
        "type_scale": type_choice,
        "spacing": spacing_choice,
        "color_ramps": apply_ramps,
    }
    
    state.log(f"   Type Scale: {type_choice}")
    state.log(f"   Spacing: {spacing_choice}")
    state.log(f"   Color Ramps: {'Yes' if apply_ramps else 'No'}")
    
    state.log("✅ Upgrades applied! Proceed to Stage 3 for export.")
    
    return "✅ Upgrades applied! Proceed to Stage 3 to export.", state.get_logs()


def export_stage1_json():
    """Export Stage 1 tokens (as-is extraction) to JSON."""
    result = {
        "metadata": {
            "source_url": state.base_url,
            "extracted_at": datetime.now().isoformat(),
            "version": "v1-stage1-extracted",
            "stage": "extraction",
        },
        "colors": {},  # Viewport-agnostic
        "typography": {
            "desktop": {},
            "mobile": {},
        },
        "spacing": {
            "desktop": {},
            "mobile": {},
        },
        "radius": {},  # Viewport-agnostic
    }
    
    # Colors (no viewport prefix - same across devices)
    if state.desktop_normalized:
        for name, c in state.desktop_normalized.colors.items():
            result["colors"][c.suggested_name or c.value] = {
                "value": c.value,
                "frequency": c.frequency,
                "confidence": c.confidence.value if c.confidence else "medium",
                "contexts": c.contexts[:3],
            }
    
    # Typography (viewport-specific)
    if state.desktop_normalized:
        for name, t in state.desktop_normalized.typography.items():
            key = t.suggested_name or f"{t.font_family}-{t.font_size}"
            result["typography"]["desktop"][key] = {
                "font_family": t.font_family,
                "font_size": t.font_size,
                "font_weight": t.font_weight,
                "line_height": t.line_height,
                "frequency": t.frequency,
            }
    
    if state.mobile_normalized:
        for name, t in state.mobile_normalized.typography.items():
            key = t.suggested_name or f"{t.font_family}-{t.font_size}"
            result["typography"]["mobile"][key] = {
                "font_family": t.font_family,
                "font_size": t.font_size,
                "font_weight": t.font_weight,
                "line_height": t.line_height,
                "frequency": t.frequency,
            }
    
    # Spacing (viewport-specific if different)
    if state.desktop_normalized:
        for name, s in state.desktop_normalized.spacing.items():
            key = s.suggested_name or s.value
            result["spacing"]["desktop"][key] = {
                "value": s.value,
                "value_px": s.value_px,
                "fits_base_8": s.fits_base_8,
                "frequency": s.frequency,
            }
    
    if state.mobile_normalized:
        for name, s in state.mobile_normalized.spacing.items():
            key = s.suggested_name or s.value
            result["spacing"]["mobile"][key] = {
                "value": s.value,
                "value_px": s.value_px,
                "fits_base_8": s.fits_base_8,
                "frequency": s.frequency,
            }
    
    # Radius (no viewport prefix)
    if state.desktop_normalized:
        for name, r in state.desktop_normalized.radius.items():
            result["radius"][name] = {
                "value": r.value,
                "frequency": r.frequency,
            }
    
    return json.dumps(result, indent=2, default=str)


def export_tokens_json():
    """Export tokens to JSON."""
    result = {
        "metadata": {
            "source_url": state.base_url,
            "extracted_at": datetime.now().isoformat(),
            "version": "v1-extracted",
        },
        "desktop": None,
        "mobile": None,
    }
    
    if state.desktop_normalized:
        result["desktop"] = {
            "colors": [
                {"value": c.value, "name": c.suggested_name, "frequency": c.frequency, 
                 "confidence": c.confidence.value if c.confidence else "medium"}
                for c in state.desktop_normalized.colors
            ],
            "typography": [
                {"font_family": t.font_family, "font_size": t.font_size, 
                 "font_weight": t.font_weight, "line_height": t.line_height,
                 "name": t.suggested_name, "frequency": t.frequency}
                for t in state.desktop_normalized.typography
            ],
            "spacing": [
                {"value": s.value, "value_px": s.value_px, "name": s.suggested_name,
                 "frequency": s.frequency, "fits_base_8": s.fits_base_8}
                for s in state.desktop_normalized.spacing
            ],
        }
    
    if state.mobile_normalized:
        result["mobile"] = {
            "colors": [
                {"value": c.value, "name": c.suggested_name, "frequency": c.frequency,
                 "confidence": c.confidence.value if c.confidence else "medium"}
                for c in state.mobile_normalized.colors
            ],
            "typography": [
                {"font_family": t.font_family, "font_size": t.font_size,
                 "font_weight": t.font_weight, "line_height": t.line_height,
                 "name": t.suggested_name, "frequency": t.frequency}
                for t in state.mobile_normalized.typography
            ],
            "spacing": [
                {"value": s.value, "value_px": s.value_px, "name": s.suggested_name,
                 "frequency": s.frequency, "fits_base_8": s.fits_base_8}
                for s in state.mobile_normalized.spacing
            ],
        }
    
    return json.dumps(result, indent=2, default=str)


# =============================================================================
# UI BUILDING
# =============================================================================

def create_ui():
    """Create the Gradio interface."""
    
    with gr.Blocks(
        title="Design System Extractor v2",
        theme=gr.themes.Soft(),
        css="""
        .color-swatch { display: inline-block; width: 24px; height: 24px; border-radius: 4px; margin-right: 8px; vertical-align: middle; }
        """
    ) as app:
        
        gr.Markdown("""
        # 🎨 Design System Extractor v2
        
        **Reverse-engineer design systems from live websites.**
        
        A semi-automated, human-in-the-loop system that extracts, normalizes, and upgrades design tokens.
        
        ---
        """)
        
        # =================================================================
        # CONFIGURATION
        # =================================================================
        
        with gr.Accordion("⚙️ Configuration", open=not bool(HF_TOKEN_FROM_ENV)):
            gr.Markdown("**HuggingFace Token** — Required for Stage 2 (AI upgrades)")
            with gr.Row():
                hf_token_input = gr.Textbox(
                    label="HF Token", placeholder="hf_xxxx", type="password",
                    scale=4, value=HF_TOKEN_FROM_ENV,
                )
                save_token_btn = gr.Button("💾 Save", scale=1)
            token_status = gr.Markdown("✅ Token loaded" if HF_TOKEN_FROM_ENV else "⏳ Enter token")
            
            def save_token(token):
                if token and len(token) > 10:
                    os.environ["HF_TOKEN"] = token.strip()
                    return "✅ Token saved!"
                return "❌ Invalid token"
            
            save_token_btn.click(save_token, [hf_token_input], [token_status])
        
        # =================================================================
        # URL INPUT & PAGE DISCOVERY
        # =================================================================
        
        with gr.Accordion("🔍 Step 1: Discover Pages", open=True):
            gr.Markdown("Enter your website URL to discover pages for extraction.")
            
            with gr.Row():
                url_input = gr.Textbox(label="Website URL", placeholder="https://example.com", scale=4)
                discover_btn = gr.Button("🔍 Discover Pages", variant="primary", scale=1)
            
            discover_status = gr.Markdown("")
            
            with gr.Row():
                log_output = gr.Textbox(label="📋 Log", lines=8, interactive=False)
            
            pages_table = gr.Dataframe(
                headers=["Select", "URL", "Title", "Type", "Status"],
                datatype=["bool", "str", "str", "str", "str"],
                label="Discovered Pages",
                interactive=True,
                visible=False,
            )
            
            extract_btn = gr.Button("🚀 Extract Tokens (Desktop + Mobile)", variant="primary", visible=False)
        
        # =================================================================
        # STAGE 1: EXTRACTION REVIEW
        # =================================================================
        
        with gr.Accordion("📊 Stage 1: Review Extracted Tokens", open=False) as stage1_accordion:
            
            extraction_status = gr.Markdown("")
            
            gr.Markdown("""
            **Review the extracted tokens.** Toggle between Desktop and Mobile viewports.
            Accept or reject tokens, then proceed to Stage 2 for AI-powered upgrades.
            """)
            
            viewport_toggle = gr.Radio(
                choices=["Desktop (1440px)", "Mobile (375px)"],
                value="Desktop (1440px)",
                label="Viewport",
            )
            
            with gr.Tabs():
                with gr.Tab("🎨 Colors"):
                    colors_table = gr.Dataframe(
                        headers=["Accept", "Color", "Suggested Name", "Frequency", "Confidence", "Contrast", "AA", "Context"],
                        datatype=["bool", "str", "str", "number", "str", "str", "str", "str"],
                        label="Colors",
                        interactive=True,
                    )
                
                with gr.Tab("📝 Typography"):
                    typography_table = gr.Dataframe(
                        headers=["Accept", "Font", "Size", "Weight", "Line Height", "Suggested Name", "Frequency", "Confidence"],
                        datatype=["bool", "str", "str", "str", "str", "str", "number", "str"],
                        label="Typography",
                        interactive=True,
                    )
                
                with gr.Tab("📏 Spacing"):
                    spacing_table = gr.Dataframe(
                        headers=["Accept", "Value", "Pixels", "Suggested Name", "Frequency", "Base 8", "Confidence"],
                        datatype=["bool", "str", "str", "str", "number", "str", "str"],
                        label="Spacing",
                        interactive=True,
                    )
                
                with gr.Tab("🔘 Radius"):
                    radius_table = gr.Dataframe(
                        headers=["Accept", "Value", "Frequency", "Context"],
                        datatype=["bool", "str", "number", "str"],
                        label="Border Radius",
                        interactive=True,
                    )
            
            with gr.Row():
                proceed_stage2_btn = gr.Button("➡️ Proceed to Stage 2: AI Upgrades", variant="primary")
                download_stage1_btn = gr.Button("📥 Download Stage 1 JSON", variant="secondary")
        
        # =================================================================
        # STAGE 2: AI UPGRADES
        # =================================================================
        
        with gr.Accordion("🧠 Stage 2: AI-Powered Upgrades", open=False) as stage2_accordion:
            
            stage2_status = gr.Markdown("Click 'Analyze' to start AI-powered design system analysis.")
            
            analyze_btn = gr.Button("🤖 Analyze Design System", variant="primary")
            
            with gr.Accordion("📋 AI Analysis Log (Click to expand)", open=True):
                stage2_log = gr.Textbox(label="Log", lines=15, interactive=False)
            
            # =============================================================
            # BRAND COMPARISON (LLM Research)
            # =============================================================
            gr.Markdown("---")
            brand_comparison = gr.Markdown("*Brand comparison will appear after analysis*")
            
            # =============================================================
            # TYPOGRAPHY SECTION - Desktop & Mobile
            # =============================================================
            gr.Markdown("---")
            gr.Markdown("## 📐 Typography")
            
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown("### 🖥️ Desktop (1440px)")
                    typography_desktop = gr.Dataframe(
                        headers=["Token", "Current", "Scale 1.2", "Scale 1.25 ⭐", "Scale 1.333", "Keep"],
                        datatype=["str", "str", "str", "str", "str", "str"],
                        label="Desktop Typography",
                        interactive=False,
                    )
                
                with gr.Column(scale=2):
                    gr.Markdown("### 📱 Mobile (375px)")
                    typography_mobile = gr.Dataframe(
                        headers=["Token", "Current", "Scale 1.2", "Scale 1.25 ⭐", "Scale 1.333", "Keep"],
                        datatype=["str", "str", "str", "str", "str", "str"],
                        label="Mobile Typography",
                        interactive=False,
                    )
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Select Type Scale Option")
                    type_scale_radio = gr.Radio(
                        choices=["Keep Current", "Scale 1.2 (Minor Third)", "Scale 1.25 (Major Third) ⭐", "Scale 1.333 (Perfect Fourth)"],
                        value="Scale 1.25 (Major Third) ⭐",
                        label="Type Scale",
                        interactive=True,
                    )
                    gr.Markdown("*Font family will be preserved. Sizes rounded to even numbers.*")
            
            # =============================================================
            # COLORS SECTION - Base Colors + Ramps
            # =============================================================
            gr.Markdown("---")
            gr.Markdown("## 🎨 Colors")
            
            base_colors_display = gr.Markdown("*Base colors will appear after analysis*")
            
            gr.Markdown("---")
            
            color_ramps_display = gr.Markdown("*Color ramps will appear after analysis*")
            
            color_ramps_checkbox = gr.Checkbox(
                label="✓ Generate color ramps (keeps base colors, adds 50-950 shades)",
                value=True,
            )
            
            # =============================================================
            # SPACING SECTION
            # =============================================================
            gr.Markdown("---")
            gr.Markdown("## 📏 Spacing (Rule-Based)")
            
            with gr.Row():
                with gr.Column(scale=2):
                    spacing_comparison = gr.Dataframe(
                        headers=["Current", "8px Grid", "4px Grid"],
                        datatype=["str", "str", "str"],
                        label="Spacing Comparison",
                        interactive=False,
                    )
                
                with gr.Column(scale=1):
                    spacing_radio = gr.Radio(
                        choices=["Keep Current", "8px Base Grid ⭐", "4px Base Grid"],
                        value="8px Base Grid ⭐",
                        label="Spacing System",
                        interactive=True,
                    )
            
            # =============================================================
            # RADIUS SECTION
            # =============================================================
            gr.Markdown("---")
            gr.Markdown("## 🔘 Border Radius (Rule-Based)")
            
            radius_display = gr.Markdown("*Radius tokens will appear after analysis*")
            
            # =============================================================
            # SHADOWS SECTION
            # =============================================================
            gr.Markdown("---")
            gr.Markdown("## 🌫️ Shadows (Rule-Based)")
            
            shadows_display = gr.Markdown("*Shadow tokens will appear after analysis*")
            
            # =============================================================
            # APPLY SECTION
            # =============================================================
            gr.Markdown("---")
            
            with gr.Row():
                apply_upgrades_btn = gr.Button("✨ Apply Selected Upgrades", variant="primary", scale=2)
                reset_btn = gr.Button("↩️ Reset to Original", variant="secondary", scale=1)
            
            apply_status = gr.Markdown("")
        
        # =================================================================
        # STAGE 3: EXPORT
        # =================================================================
        
        with gr.Accordion("📦 Stage 3: Export", open=False):
            gr.Markdown("""
            Export your design tokens to JSON (compatible with Figma Tokens Studio).
            
            - **Stage 1 JSON**: Raw extracted tokens (as-is)
            - **Final JSON**: Upgraded tokens with selected improvements
            """)
            
            with gr.Row():
                export_stage1_btn = gr.Button("📥 Export Stage 1 (As-Is)", variant="secondary")
                export_final_btn = gr.Button("📥 Export Final (Upgraded)", variant="primary")
            
            export_output = gr.Code(label="Tokens JSON", language="json", lines=25)
            
            export_stage1_btn.click(export_stage1_json, outputs=[export_output])
            export_final_btn.click(export_tokens_json, outputs=[export_output])
        
        # =================================================================
        # EVENT HANDLERS
        # =================================================================
        
        # Store data for viewport toggle
        desktop_data = gr.State({})
        mobile_data = gr.State({})
        
        # Discover pages
        discover_btn.click(
            fn=discover_pages,
            inputs=[url_input],
            outputs=[discover_status, log_output, pages_table],
        ).then(
            fn=lambda: (gr.update(visible=True), gr.update(visible=True)),
            outputs=[pages_table, extract_btn],
        )
        
        # Extract tokens
        extract_btn.click(
            fn=extract_tokens,
            inputs=[pages_table],
            outputs=[extraction_status, log_output, desktop_data, mobile_data],
        ).then(
            fn=lambda d: (d.get("colors", []), d.get("typography", []), d.get("spacing", [])),
            inputs=[desktop_data],
            outputs=[colors_table, typography_table, spacing_table],
        ).then(
            fn=lambda: gr.update(open=True),
            outputs=[stage1_accordion],
        )
        
        # Viewport toggle
        viewport_toggle.change(
            fn=switch_viewport,
            inputs=[viewport_toggle],
            outputs=[colors_table, typography_table, spacing_table],
        )
        
        # Stage 2: Analyze
        analyze_btn.click(
            fn=run_stage2_analysis,
            outputs=[stage2_status, stage2_log, brand_comparison, 
                     typography_desktop, typography_mobile, spacing_comparison, 
                     base_colors_display, color_ramps_display, radius_display, shadows_display],
        )
        
        # Stage 2: Apply upgrades
        apply_upgrades_btn.click(
            fn=apply_selected_upgrades,
            inputs=[type_scale_radio, spacing_radio, color_ramps_checkbox],
            outputs=[apply_status, stage2_log],
        )
        
        # Stage 1: Download JSON
        download_stage1_btn.click(
            fn=export_stage1_json,
            outputs=[export_output],
        )
        
        # Proceed to Stage 2 button
        proceed_stage2_btn.click(
            fn=lambda: gr.update(open=True),
            outputs=[stage2_accordion],
        )
        
        # =================================================================
        # FOOTER
        # =================================================================
        
        gr.Markdown("""
        ---
        **Design System Extractor v2** | Built with Playwright + Gradio + LangGraph + HuggingFace
        
        *A semi-automated co-pilot for design system recovery and modernization.*
        """)
    
    return app


# =============================================================================
# MAIN
# =============================================================================

if __name__ == "__main__":
    app = create_ui()
    app.launch(server_name="0.0.0.0", server_port=7860)