"""
Stage 1 Orchestrator
=====================
Runs the discovery conversation and produces brief artifacts.

Usage:
  # Start a new project
  python orchestrator.py new "My Product Name" /path/to/project

  # Resume discovery conversation
  python orchestrator.py resume /path/to/project

  # Synthesize brief from conversation (provide artifacts as JSON)
  python orchestrator.py synthesize /path/to/project brief.json features.json flows.json
"""

import json
import os
import sys

# Add src directory to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from scaffold import create_project, log_decision
from discovery import (
    DISCOVERY_SYSTEM_PROMPT,
    create_discovery_state,
    get_discovery_state,
    update_discovery_state,
    build_session_context,
    save_brief,
    save_features,
    save_flow,
    complete_discovery,
    get_full_prompt,
)


def cmd_new(project_name: str, project_path: str):
    """Create a new project and initialize discovery."""
    create_project(project_path, project_name)
    create_discovery_state(project_path)
    
    print(f"\n{'='*60}")
    print(f"Project '{project_name}' created at: {project_path}")
    print(f"{'='*60}")
    print(f"\nStage 1: Discovery is ready to begin.")
    print(f"\nTo use the discovery agent, start a conversation with this system prompt:")
    print(f"\n  python -c \"from discovery import get_full_prompt; print(get_full_prompt('{project_path}'))\"")
    print(f"\nOr use the system prompt directly in Claude with the context below:")
    print(f"\n{'─'*60}")
    print(build_session_context(project_path))
    print(f"{'─'*60}")


def cmd_resume(project_path: str):
    """Show current discovery state and context for resuming."""
    if not os.path.exists(os.path.join(project_path, "project.json")):
        print(f"Error: No project found at {project_path}")
        return
    
    with open(os.path.join(project_path, "project.json")) as f:
        manifest = json.load(f)
    
    print(f"\n{'='*60}")
    print(f"Project: {manifest['name']}")
    print(f"Stage: {manifest['current_stage']}")
    print(f"{'='*60}")
    
    context = build_session_context(project_path)
    print(f"\n{context}")
    
    if manifest['current_stage'] != 'discovery':
        print(f"\nDiscovery is already complete. Current stage: {manifest['current_stage']}")
        return
    
    print(f"\n{'─'*60}")
    print(f"To resume, inject this context into your conversation with the discovery agent.")


def cmd_synthesize(project_path: str, brief_path: str, features_path: str, *flow_paths):
    """Save brief artifacts from JSON files and advance the project."""
    with open(brief_path) as f:
        brief = json.load(f)
    with open(features_path) as f:
        features = json.load(f)
    
    flows = []
    for fp in flow_paths:
        with open(fp) as f:
            flows.append(json.load(f))
    
    complete_discovery(project_path, brief, features, flows)


def cmd_save_brief(project_path: str, brief_json_str: str):
    """Save a brief from a JSON string (for piping from conversation)."""
    brief = json.loads(brief_json_str)
    save_brief(project_path, brief)
    print("Brief saved.")


def cmd_save_features(project_path: str, features_json_str: str):
    """Save features from a JSON string."""
    features = json.loads(features_json_str)
    save_features(project_path, features)
    print("Features saved.")


def cmd_save_flow(project_path: str, flow_json_str: str):
    """Save a flow from a JSON string."""
    flow = json.loads(flow_json_str)
    flow_name = flow.get("flow_name", "unnamed_flow")
    save_flow(project_path, flow, flow_name)
    print(f"Flow '{flow_name}' saved.")


# ═══════════════════════════════════════════════════════════
# STAGE 2: INFORMATION ARCHITECTURE COMMANDS
# ═══════════════════════════════════════════════════════════

def cmd_ia_start(project_path: str):
    """Initialize Stage 2 and show screen inventory."""
    from ia import create_ia_state, extract_screens_from_flows, get_full_prompt
    
    state = create_ia_state(project_path)
    screens = extract_screens_from_flows(project_path)
    
    print(f"\n{'='*60}")
    print(f"Stage 2: Information Architecture")
    print(f"{'='*60}")
    print(f"\nScreens identified from flows: {len(screens)}")
    for s in screens:
        print(f"  ○ {s['name']}")
        print(f"    Flows: {', '.join(s['appears_in_flows'])}")
        for a in s['user_actions'][:2]:
            print(f"    Action: {a}")
    
    print(f"\nUse the IA agent system prompt in your Claude conversation.")
    print(f"The prompt includes the full brief context and screen inventory.")
    print(f"\nPrompt length: {len(get_full_prompt(project_path))} characters")


def cmd_ia_save_screen(project_path: str, screen_name: str, 
                       desktop_json_str: str, mobile_json_str: str,
                       strategy: str = "reflow"):
    """Save approved IA for a screen from JSON strings."""
    from ia import save_screen_ia
    
    desktop_map = json.loads(desktop_json_str)
    mobile_map = json.loads(mobile_json_str)
    save_screen_ia(project_path, screen_name, desktop_map, mobile_map, strategy)


def cmd_ia_save_screen_file(project_path: str, screen_name: str,
                            desktop_json_path: str, mobile_json_path: str,
                            strategy: str = "reflow"):
    """Save approved IA for a screen from JSON files."""
    from ia import save_screen_ia
    
    with open(desktop_json_path) as f:
        desktop_map = json.load(f)
    with open(mobile_json_path) as f:
        mobile_map = json.load(f)
    
    save_screen_ia(project_path, screen_name, desktop_map, mobile_map, strategy)


def cmd_ia_reject(project_path: str, screen_name: str, 
                  variant_name: str, variant_json_str: str, reason: str = ""):
    """Archive a rejected IA variant."""
    from ia import reject_screen_variant
    variant_map = json.loads(variant_json_str)
    reject_screen_variant(project_path, screen_name, variant_name, variant_map, reason)


def cmd_ia_save_sitemap(project_path: str, sitemap_json_str: str):
    """Save the navigation sitemap."""
    from ia import save_sitemap
    sitemap = json.loads(sitemap_json_str)
    save_sitemap(project_path, sitemap)


def cmd_ia_save_responsive(project_path: str, strategy_json_str: str):
    """Save the responsive strategy document."""
    from ia import save_responsive_strategy
    strategy = json.loads(strategy_json_str)
    save_responsive_strategy(project_path, strategy)


def cmd_ia_complete(project_path: str):
    """Verify all screens approved and advance to Stage 3."""
    from ia import complete_ia
    complete_ia(project_path)


# ═══════════════════════════════════════════════════════════
# STAGE 3: STYLE DIRECTION COMMANDS
# ═══════════════════════════════════════════════════════════

def cmd_style_start(project_path: str):
    """Initialize Stage 3 and show context."""
    from style import get_full_prompt, load_style_context
    
    print(f"\n{'='*60}")
    print(f"Stage 3: Style Direction")
    print(f"{'='*60}")
    
    # Show relevant brief 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)
        brand = brief.get("constraints", {}).get("brand", "not specified")
        comp = brief.get("constraints", {}).get("competitive_context", "not specified")
        print(f"\n  Brand direction: {brand}")
        print(f"  Competitive context: {comp}")
        if brief.get("personas"):
            p = brief["personas"][0]
            print(f"  Primary persona: {p.get('name', '?')} — {p.get('context', '?')}")
    
    prompt_len = len(get_full_prompt(project_path))
    print(f"\n  Style agent prompt: {prompt_len} chars")
    print(f"\n  Use the style direction system prompt in your Claude conversation.")
    print(f"  The agent will generate 3-4 distinct style directions with previews.")


def cmd_style_save_candidate(project_path: str, name: str, token_json_str: str):
    """Save a candidate token set."""
    from style import save_token_candidate
    tokens = json.loads(token_json_str)
    save_token_candidate(project_path, name, tokens)


def cmd_style_save_candidate_file(project_path: str, name: str, token_json_path: str):
    """Save a candidate token set from a file."""
    from style import save_token_candidate
    with open(token_json_path) as f:
        tokens = json.load(f)
    save_token_candidate(project_path, name, tokens)


def cmd_style_select(project_path: str, name: str, token_json_str: str):
    """Select the final token set."""
    from style import select_token_set, save_tailwind_config
    tokens = json.loads(token_json_str)
    select_token_set(project_path, tokens, name)
    save_tailwind_config(project_path, tokens)


def cmd_style_select_file(project_path: str, name: str, token_json_path: str):
    """Select the final token set from a file."""
    from style import select_token_set, save_tailwind_config
    with open(token_json_path) as f:
        tokens = json.load(f)
    select_token_set(project_path, tokens, name)
    save_tailwind_config(project_path, tokens)


def cmd_style_save_preview(project_path: str, direction_name: str, preview_path: str):
    """Save a preview component file."""
    from style import save_preview
    with open(preview_path) as f:
        content = f.read()
    ext = "jsx" if preview_path.endswith(".jsx") else "html"
    save_preview(project_path, direction_name, content, ext)


def cmd_style_complete(project_path: str):
    """Verify token set selected and advance to Stage 4."""
    from style import complete_style
    complete_style(project_path)


# ═══════════════════════════════════════════════════════════
# STAGE 4: COMPONENT GENERATION COMMANDS
# ═══════════════════════════════════════════════════════════

def cmd_comp_start(project_path: str):
    """Initialize Stage 4 and show what components are needed."""
    from components import get_full_prompt, get_registry
    
    print(f"\n{'='*60}")
    print(f"Stage 4: Component Generation")
    print(f"{'='*60}")
    
    # Show token summary
    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)
        meta = tokens.get("meta", {})
        print(f"\n  Style: {meta.get('name', '?')}")
        print(f"  Personality: {meta.get('personality', '?')[:80]}")
    
    # Show screens that need components
    screens_dir = os.path.join(project_path, "architecture", "screens")
    if os.path.exists(screens_dir):
        screens = [d for d in sorted(os.listdir(screens_dir)) if not d.startswith('_')]
        print(f"\n  Screens needing components: {', '.join(screens)}")
    
    registry = get_registry(project_path)
    print(f"  Components in registry: {len(registry.get('components', []))}")
    print(f"\n  Agent prompt: {len(get_full_prompt(project_path))} chars")


def cmd_comp_save(project_path: str, component_name: str, page_name: str,
                  code_path: str, registry_json_str: str):
    """Save an approved component from a file + registry entry."""
    from components import save_component
    with open(code_path) as f:
        code = f.read()
    registry_entry = json.loads(registry_json_str)
    save_component(project_path, component_name, page_name, code, registry_entry)


def cmd_comp_reject(project_path: str, component_name: str, page_name: str,
                    variant_name: str, code_path: str, reason: str = ""):
    """Archive a rejected component variant."""
    from components import reject_component_variant
    with open(code_path) as f:
        code = f.read()
    reject_component_variant(project_path, component_name, page_name, variant_name, code, reason)


def cmd_comp_promote(project_path: str, component_id: str, feature_name: str):
    """Promote a page-specific component to shared."""
    from components import promote_component
    promote_component(project_path, component_id, feature_name)


def cmd_comp_registry(project_path: str):
    """Show the component registry."""
    from components import get_registry
    registry = get_registry(project_path)
    components = registry.get("components", [])
    
    if not components:
        print("No components registered yet.")
        return
    
    print(f"\nComponent Registry ({len(components)} components):")
    by_type = {}
    for c in components:
        by_type.setdefault(c.get("type", "?"), []).append(c)
    
    for t, comps in sorted(by_type.items()):
        print(f"\n  {t.upper()}:")
        for c in comps:
            loc = "shared" if c.get("location") == "shared" else c.get("used_in", ["?"])[0] if c.get("used_in") else "?"
            print(f"    {c['id']:<30} [{c.get('status', '?')}] loc={loc}")


def cmd_comp_complete(project_path: str):
    """Advance to Stage 5."""
    from components import complete_components
    complete_components(project_path)


# ═══════════════════════════════════════════════════════════
# STAGE 5: PAGE ASSEMBLY COMMANDS
# ═══════════════════════════════════════════════════════════

def cmd_page_start(project_path: str):
    """Initialize Stage 5 and show assembly status."""
    from assembly import get_full_prompt
    
    print(f"\n{'='*60}")
    print(f"Stage 5: Page Assembly")
    print(f"{'='*60}")
    
    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 sorted(os.listdir(screens_dir)) if not d.startswith('_')]
        print(f"\n  Screens to assemble: {len(screens)}")
        for s in screens:
            page_dir = os.path.join(pages_dir, s)
            has_page = False
            if os.path.exists(page_dir):
                has_page = any(f.endswith('.jsx') for f in os.listdir(page_dir))
            status = "✓" if has_page else "○"
            print(f"    {status} {s}")
    
    from components import get_registry
    registry = get_registry(project_path)
    comp_count = len([c for c in registry.get("components", []) if c.get("status") == "approved"])
    print(f"\n  Available components: {comp_count}")
    print(f"  Agent prompt: {len(get_full_prompt(project_path))} chars")


def cmd_page_save(project_path: str, page_name: str, code_path: str):
    """Save an assembled page from a file."""
    from assembly import save_page
    with open(code_path) as f:
        code = f.read()
    save_page(project_path, page_name, code)


def cmd_page_revise(project_path: str, page_name: str, code_path: str, notes: str = ""):
    """Save a revised page, archiving the previous version."""
    from assembly import save_page_revision
    with open(code_path) as f:
        code = f.read()
    save_page_revision(project_path, page_name, code, notes)


def cmd_page_complete(project_path: str):
    """Verify all pages assembled and mark project complete."""
    from assembly import complete_assembly
    complete_assembly(project_path)


def cmd_status(project_path: str):
    """Show full project status."""
    if not os.path.exists(os.path.join(project_path, "project.json")):
        print(f"Error: No project found at {project_path}")
        return
    
    with open(os.path.join(project_path, "project.json")) as f:
        manifest = json.load(f)
    
    print(f"\nProject: {manifest['name']}")
    print(f"Created: {manifest['created']}")
    print(f"Current stage: {manifest['current_stage']}")
    
    if manifest['stages_completed']:
        print(f"\nCompleted stages:")
        for s in manifest['stages_completed']:
            print(f"  ✓ {s['stage']} — {s['completed'][:10]}")
    
    # Check what artifacts exist
    brief_dir = os.path.join(project_path, "brief")
    arch_dir = os.path.join(project_path, "architecture")
    tokens_dir = os.path.join(project_path, "tokens")
    comp_dir = os.path.join(project_path, "components")
    pages_dir = os.path.join(project_path, "pages")
    
    print(f"\nArtifacts:")
    
    brief_exists = os.path.exists(os.path.join(brief_dir, "brief.json"))
    features_exists = os.path.exists(os.path.join(brief_dir, "features.json"))
    flows_dir = os.path.join(brief_dir, "flows")
    flow_count = len([f for f in os.listdir(flows_dir) if f.endswith('.json')]) if os.path.exists(flows_dir) else 0
    print(f"  Brief: {'✓' if brief_exists else '○'}")
    print(f"  Features: {'✓' if features_exists else '○'}")
    print(f"  Flows: {flow_count} mapped")
    
    screens_dir = os.path.join(arch_dir, "screens")
    screen_count = len([d for d in os.listdir(screens_dir) if os.path.isdir(os.path.join(screens_dir, d))]) if os.path.exists(screens_dir) else 0
    print(f"  Screen IAs: {screen_count}")
    
    tokens_exists = os.path.exists(os.path.join(tokens_dir, "selected.json"))
    print(f"  Token set: {'✓' if tokens_exists else '○'}")
    
    registry_path = os.path.join(comp_dir, "registry.json")
    if os.path.exists(registry_path):
        with open(registry_path) as f:
            registry = json.load(f)
        comp_count = len(registry.get("components", []))
        print(f"  Components: {comp_count}")
    
    page_dirs = [d for d in os.listdir(pages_dir) if os.path.isdir(os.path.join(pages_dir, d))] if os.path.exists(pages_dir) else []
    print(f"  Pages: {len(page_dirs)}")
    
    changelog_dir = os.path.join(project_path, "changelog")
    changelog_count = len([f for f in os.listdir(changelog_dir) if f.endswith('.md')]) if os.path.exists(changelog_dir) else 0
    print(f"  Changelog entries: {changelog_count}")


COMMANDS = {
    # General
    "new": (cmd_new, "Create a new project", "<name> <path>"),
    "resume": (cmd_resume, "Show context for resuming", "<path>"),
    "status": (cmd_status, "Show project status", "<path>"),
    # Stage 1: Discovery
    "synthesize": (cmd_synthesize, "Save Stage 1 artifacts, advance", "<path> <brief> <features> [flows...]"),
    "save-brief": (cmd_save_brief, "Save brief JSON", "<path> '<json>'"),
    "save-features": (cmd_save_features, "Save features JSON", "<path> '<json>'"),
    "save-flow": (cmd_save_flow, "Save flow JSON", "<path> '<json>'"),
    # Stage 2: Information Architecture
    "ia-start": (cmd_ia_start, "Initialize Stage 2", "<path>"),
    "ia-save-screen": (cmd_ia_save_screen, "Save screen IA (strings)", "<path> <n> '<desktop>' '<mobile>' [strategy]"),
    "ia-save-screen-file": (cmd_ia_save_screen_file, "Save screen IA (files)", "<path> <n> <desk.json> <mob.json> [strategy]"),
    "ia-reject": (cmd_ia_reject, "Archive rejected variant", "<path> <screen> <variant> '<json>' [reason]"),
    "ia-save-sitemap": (cmd_ia_save_sitemap, "Save sitemap", "<path> '<json>'"),
    "ia-save-responsive": (cmd_ia_save_responsive, "Save responsive strategy", "<path> '<json>'"),
    "ia-complete": (cmd_ia_complete, "Advance to Stage 3", "<path>"),
    # Stage 3: Style Direction
    "style-start": (cmd_style_start, "Initialize Stage 3", "<path>"),
    "style-save-candidate": (cmd_style_save_candidate, "Save token candidate", "<path> <name> '<json>'"),
    "style-save-candidate-file": (cmd_style_save_candidate_file, "Save token candidate (file)", "<path> <name> <f.json>"),
    "style-select": (cmd_style_select, "Select final tokens", "<path> <name> '<json>'"),
    "style-select-file": (cmd_style_select_file, "Select final tokens (file)", "<path> <name> <f.json>"),
    "style-save-preview": (cmd_style_save_preview, "Save preview component", "<path> <direction> <f.jsx>"),
    "style-complete": (cmd_style_complete, "Advance to Stage 4", "<path>"),
    # Stage 4: Component Generation
    "comp-start": (cmd_comp_start, "Initialize Stage 4", "<path>"),
    "comp-save": (cmd_comp_save, "Save component", "<path> <name> <page> <code.jsx> '<registry_json>'"),
    "comp-reject": (cmd_comp_reject, "Archive rejected variant", "<path> <name> <page> <variant> <code.jsx> [reason]"),
    "comp-promote": (cmd_comp_promote, "Promote to shared", "<path> <component_id> <feature>"),
    "comp-registry": (cmd_comp_registry, "Show component registry", "<path>"),
    "comp-complete": (cmd_comp_complete, "Advance to Stage 5", "<path>"),
    # Stage 5: Page Assembly
    "page-start": (cmd_page_start, "Initialize Stage 5", "<path>"),
    "page-save": (cmd_page_save, "Save assembled page", "<path> <page_name> <code.jsx>"),
    "page-revise": (cmd_page_revise, "Save revised page", "<path> <page_name> <code.jsx> [notes]"),
    "page-complete": (cmd_page_complete, "Mark project complete", "<path>"),
}


if __name__ == "__main__":
    if len(sys.argv) < 2 or sys.argv[1] not in COMMANDS:
        print("Design Agent — Stage 1 Orchestrator")
        print(f"\nCommands:")
        for name, (fn, desc, args) in COMMANDS.items():
            print(f"  {name:<16} {desc}")
            print(f"  {'':16} Usage: python orchestrator.py {name} {args}")
        sys.exit(1)
    
    cmd_name = sys.argv[1]
    cmd_fn = COMMANDS[cmd_name][0]
    cmd_fn(*sys.argv[2:])