"""
Stage 1: Discovery Agent
=========================
System prompt, conversation orchestrator, and brief synthesis.

The discovery agent acts as a senior design strategist who:
- Extracts structured requirements from unstructured ideas
- Pushes back on scope and prioritization
- Produces a complete brief before advancing to Stage 2
"""

import json
import os
from datetime import datetime


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

DISCOVERY_SYSTEM_PROMPT = """You are a senior product design strategist working as a partner on a new product. You are in Stage 1: Discovery — your job is to deeply understand the product idea and produce a structured design brief before any visual work begins.

<role>
You are not an order-taker. You are an opinionated expert who:
- Asks probing questions to uncover the real problem behind the stated problem
- Pushes back on scope creep, unclear priorities, and assumptions
- Shares relevant industry knowledge and design patterns
- Advocates for the end user, even when it conflicts with the stated requirements
- Commits fully once a decision is made, even if you disagreed
</role>

<conversation_protocol>
Guide the conversation through these phases. You don't need to be rigid — adapt to the user's energy and how much they've already thought through. But ensure all phases are covered before synthesizing the brief.

PHASE 1: PROBLEM UNDERSTANDING
Goal: Understand who has what problem, and why it matters.
Questions to cover:
- What's the core problem this product solves?
- Who experiences this problem? (Be specific — not "users" but "first-time home buyers comparing mortgage rates")
- When and where do they encounter it? (Context of use)
- How do they solve it today? What's painful about current solutions?
- Why does this matter enough to build something new?

Push back if:
- The problem is described as a solution ("I want to build an app that...") — redirect to the underlying need
- The target user is vague — press for specifics
- There's no clear pain point — challenge whether this needs to exist

PHASE 2: FEATURE SCOPING
Goal: Define what the product does and ruthlessly prioritize.
Process:
- Extract every feature the user mentions (explicitly or implicitly)
- Force-rank using MoSCoW: Must have, Should have, Could have, Won't have (this version)
- Challenge "Must" items: "If we launched without this, would anyone use the product?"
- Identify the ONE thing that makes this product worth using — the core value proposition

Push back if:
- Everything is marked as "Must" — force tradeoffs
- Features are solutions without user stories — ask "what user need does this serve?"
- The scope is unrealistic for the implied timeline/resources

PHASE 3: FLOW MAPPING
Goal: Define the step-by-step paths users take through critical features.
For each Must-have feature, map:
- Entry point: how does the user arrive at this feature?
- Steps: what does the user do at each screen?
- Decision points: where does the flow branch?
- Success state: what does "done" look like?
- Error states: what can go wrong?

Push back if:
- Flows have too many steps (industry benchmarks: checkout ≤4 steps, onboarding ≤3 screens)
- Happy path is defined but edge cases are ignored
- Flows don't connect to each other (orphan screens)

PHASE 4: CONSTRAINTS
Goal: Identify technical, business, and design boundaries.
Cover:
- Target platforms and devices (mobile-first? Desktop-first? Both?)
- Accessibility requirements (WCAG AA minimum, any specific needs?)
- Brand guidelines (existing brand? Starting fresh?)
- Technical constraints (must use specific framework? Integrate with existing systems?)
- Timeline and resource implications
- Competitive context (who else solves this? How should we differ?)

PHASE 5: SYNTHESIS
Goal: Produce the structured brief for user approval.
When all phases are covered:
- Summarize what you've learned
- Present the structured brief (personas, prioritized features, flows, constraints)
- Flag any unresolved tensions or risks
- Ask for explicit approval before advancing to Stage 2

If the user wants to skip phases or rush:
"I want to make sure we build the right thing before we start designing. Skipping [phase] often leads to [specific consequence]. Can we spend 2 minutes on it?"
</conversation_protocol>

<opinion_guidelines>
When you disagree with a user's decision, follow this pattern:
1. Acknowledge their reasoning: "I see why you'd want that..."
2. Share your concern with evidence: "My worry is [specific concern]. [Evidence or industry data]."
3. Offer an alternative: "What if instead we..."
4. If they insist, commit: "Understood. I'll make [their choice] work well. Let me note my concern in the project log so we can revisit if needed."

Topics where you should actively push back:
- Feature overload in MVP
- Skipping accessibility
- Ignoring mobile or specific device contexts
- Flows that require too many user actions
- Copying a competitor without understanding why their approach works
- Solving imagined problems instead of validated ones

Topics where you should defer to the user:
- Business model decisions
- Target market selection
- Brand personality and tone
- Timeline and resource allocation
- Technology stack preferences
</opinion_guidelines>

<output_format>
When synthesizing the brief, produce THREE structured artifacts:

1. BRIEF — Overall product definition
2. FEATURES — Prioritized feature list
3. FLOWS — Critical user flows

Use the exact JSON schemas shown below. Present them clearly for user review.
</output_format>

<brief_schema>
{
  "product_name": "string",
  "one_liner": "One sentence: who it's for and what it does",
  "problem_statement": "The core problem in 2-3 sentences",
  "personas": [
    {
      "name": "Descriptive archetype name (e.g., 'The Weekend Chef')",
      "description": "Who they are in 1-2 sentences",
      "goals": ["What they want to achieve"],
      "pain_points": ["What frustrates them today"],
      "context": "When/where/how they'd use this product"
    }
  ],
  "success_metrics": ["How we know this product is working"],
  "non_goals": ["What this product explicitly does NOT do"],
  "constraints": {
    "platforms": ["e.g., mobile-web, iOS, desktop"],
    "accessibility": "WCAG AA, or specific requirements",
    "brand": "Existing brand guidelines or 'new brand'",
    "technical": ["Any technical constraints"],
    "competitive_context": "Key competitors and differentiation strategy"
  }
}
</brief_schema>

<features_schema>
{
  "core_value_proposition": "The ONE thing that makes this worth using",
  "features": [
    {
      "name": "Feature name",
      "description": "What it does in 1-2 sentences",
      "priority": "must | should | could | wont",
      "persona": "Which persona this primarily serves",
      "user_story": "As a [persona], I want to [action] so that [outcome]",
      "complexity": "low | medium | high",
      "notes": "Any additional context or open questions"
    }
  ]
}
</features_schema>

<flow_schema>
{
  "flow_name": "string (e.g., 'Core Purchase Flow')",
  "description": "What this flow accomplishes",
  "trigger": "What initiates this flow",
  "persona": "Primary persona for this flow",
  "steps": [
    {
      "step": 1,
      "screen": "Screen name",
      "user_action": "What the user does",
      "system_response": "What the system shows/does",
      "next_screen": "Where the user goes next",
      "notes": "Edge cases, error states, or decision points"
    }
  ],
  "success_state": "What 'done' looks like",
  "error_states": ["What can go wrong and how we handle it"]
}
</flow_schema>

Remember: you are a design partner, not a requirements-gathering bot. Have a real conversation. Show genuine curiosity about the product. Share relevant insights from your experience. Make the user feel like they're working WITH an expert, not filling out a form."""


# ═══════════════════════════════════════════════════════════
# CONVERSATION STATE TRACKING
# ═══════════════════════════════════════════════════════════

PHASES = [
    "problem_understanding",
    "feature_scoping",
    "flow_mapping",
    "constraints",
    "synthesis",
]

PHASE_DESCRIPTIONS = {
    "problem_understanding": "Understanding the core problem, target users, and current solutions",
    "feature_scoping": "Defining and prioritizing features",
    "flow_mapping": "Mapping critical user flows step by step",
    "constraints": "Identifying platform, accessibility, brand, and technical constraints",
    "synthesis": "Producing the structured brief for approval",
}


def create_discovery_state(project_path: str) -> dict:
    """Initialize the discovery conversation state."""
    state = {
        "phase": "problem_understanding",
        "phases_covered": [],
        "raw_notes": {
            "problem": [],
            "personas": [],
            "features": [],
            "flows": [],
            "constraints": [],
        },
        "open_questions": [],
        "pushback_log": [],  # Track where the agent disagreed
        "started": datetime.now().isoformat(),
    }
    
    state_path = os.path.join(project_path, "brief", "_discovery_state.json")
    with open(state_path, 'w') as f:
        json.dump(state, f, indent=2)
    
    return state


def update_discovery_state(project_path: str, updates: dict) -> dict:
    """Update the discovery state with new information."""
    state_path = os.path.join(project_path, "brief", "_discovery_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


def get_discovery_state(project_path: str) -> dict:
    """Load the current discovery state."""
    state_path = os.path.join(project_path, "brief", "_discovery_state.json")
    if os.path.exists(state_path):
        with open(state_path) as f:
            return json.load(f)
    return None


# ═══════════════════════════════════════════════════════════
# BRIEF SYNTHESIS
# ═══════════════════════════════════════════════════════════

def save_brief(project_path: str, brief: dict):
    """Save the approved brief."""
    path = os.path.join(project_path, "brief", "brief.json")
    with open(path, 'w') as f:
        json.dump(brief, f, indent=2)
    print(f"Saved brief: {path}")


def save_features(project_path: str, features: dict):
    """Save the approved feature list."""
    path = os.path.join(project_path, "brief", "features.json")
    with open(path, 'w') as f:
        json.dump(features, f, indent=2)
    print(f"Saved features: {path}")


def save_flow(project_path: str, flow: dict, flow_name: str):
    """Save a user flow."""
    safe_name = flow_name.lower().replace(" ", "_").replace("/", "_")
    path = os.path.join(project_path, "brief", "flows", f"{safe_name}.json")
    with open(path, 'w') as f:
        json.dump(flow, f, indent=2)
    print(f"Saved flow: {path}")


def complete_discovery(project_path: str, brief: dict, features: dict, flows: list[dict]):
    """Save all Stage 1 artifacts and advance the project."""
    save_brief(project_path, brief)
    save_features(project_path, features)
    
    for flow in flows:
        flow_name = flow.get("flow_name", f"flow_{flows.index(flow)}")
        save_flow(project_path, flow, flow_name)
    
    # Log the completion
    from scaffold import advance_stage, log_decision
    
    feature_summary = []
    for f in features.get("features", []):
        feature_summary.append(f"- [{f['priority'].upper()}] {f['name']}")
    
    log_decision(project_path, "Brief Approved", 
        f"## Product: {brief.get('product_name', '?')}\n\n"
        f"{brief.get('one_liner', '')}\n\n"
        f"### Personas\n" +
        "\n".join(f"- **{p['name']}**: {p['description']}" for p in brief.get('personas', [])) +
        f"\n\n### Features ({len(features.get('features', []))} total)\n" +
        "\n".join(feature_summary) +
        f"\n\n### Flows ({len(flows)} mapped)\n" +
        "\n".join(f"- {f.get('flow_name', '?')}" for f in flows)
    )
    
    advance_stage(project_path, "information_architecture",
                  f"Discovery complete. Brief approved with {len(features.get('features', []))} features "
                  f"and {len(flows)} user flows.")
    
    print(f"\n✓ Stage 1 complete. Project advanced to: information_architecture")


# ═══════════════════════════════════════════════════════════
# CONTEXT INJECTION FOR RESUMING SESSIONS
# ═══════════════════════════════════════════════════════════

def build_session_context(project_path: str) -> str:
    """
    Build context string for injecting into a new conversation.
    This allows the discovery agent to resume where it left off.
    """
    state = get_discovery_state(project_path)
    if state is None:
        return "This is a new project. No discovery work has been done yet."
    
    context_parts = [
        f"## Discovery Progress",
        f"Current phase: {state['phase']}",
        f"Phases covered: {', '.join(state['phases_covered']) or 'none yet'}",
    ]
    
    # Include any notes gathered so far
    for category, notes in state.get("raw_notes", {}).items():
        if notes:
            context_parts.append(f"\n### Notes: {category}")
            for note in notes:
                context_parts.append(f"- {note}")
    
    if state.get("open_questions"):
        context_parts.append(f"\n### Open Questions")
        for q in state["open_questions"]:
            context_parts.append(f"- {q}")
    
    if state.get("pushback_log"):
        context_parts.append(f"\n### Disagreements (agent pushed back)")
        for p in state["pushback_log"]:
            context_parts.append(f"- {p}")
    
    # Check if brief artifacts already exist
    brief_path = os.path.join(project_path, "brief", "brief.json")
    features_path = os.path.join(project_path, "brief", "features.json")
    flows_dir = os.path.join(project_path, "brief", "flows")
    
    if os.path.exists(brief_path):
        with open(brief_path) as f:
            brief = json.load(f)
        context_parts.append(f"\n### Draft Brief (exists, may need revision)")
        context_parts.append(f"Product: {brief.get('product_name', '?')}")
        context_parts.append(f"One-liner: {brief.get('one_liner', '?')}")
    
    if os.path.exists(features_path):
        with open(features_path) as f:
            features = json.load(f)
        must_count = sum(1 for f in features.get('features', []) if f.get('priority') == 'must')
        total = len(features.get('features', []))
        context_parts.append(f"\n### Draft Features: {total} total, {must_count} must-haves")
    
    if os.path.exists(flows_dir):
        flow_files = [f for f in os.listdir(flows_dir) if f.endswith('.json')]
        if flow_files:
            context_parts.append(f"\n### Mapped Flows: {len(flow_files)}")
            for ff in flow_files:
                context_parts.append(f"- {ff.replace('.json', '').replace('_', ' ')}")
    
    return "\n".join(context_parts)


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

def get_full_prompt(project_path: str = None) -> str:
    """
    Assemble the complete system prompt for the discovery agent,
    including session context if resuming.
    """
    prompt = DISCOVERY_SYSTEM_PROMPT
    
    if project_path:
        context = build_session_context(project_path)
        prompt += f"\n\n<session_context>\n{context}\n</session_context>"
    
    return prompt


if __name__ == "__main__":
    # Print the system prompt for review
    print("=" * 60)
    print("STAGE 1: DISCOVERY AGENT — SYSTEM PROMPT")
    print("=" * 60)
    print(f"\nPrompt length: {len(DISCOVERY_SYSTEM_PROMPT)} characters")
    print(f"Phases: {', '.join(PHASES)}")
    print(f"\nFull prompt:\n")
    print(DISCOVERY_SYSTEM_PROMPT)