"""
Context Builder

Builds rich context prompts for Claude sessions based on mode and study.
"""

import json
import sqlite3
from pathlib import Path
from typing import Any, Dict, List, Literal, Optional

# Atomizer root directory
ATOMIZER_ROOT = Path(__file__).parent.parent.parent.parent.parent


class ContextBuilder:
    """Builds context prompts for Claude sessions"""

    def __init__(self):
        self.atomizer_root = ATOMIZER_ROOT
        self.studies_dir = ATOMIZER_ROOT / "studies"

    def build(
        self,
        mode: Literal["user", "power"],
        study_id: Optional[str] = None,
        conversation_history: Optional[List[Dict[str, Any]]] = None,
        canvas_state: Optional[Dict[str, Any]] = None,
        spec_path: Optional[str] = None,
    ) -> str:
        """
        Build full system prompt with context.

        Args:
            mode: "user" for safe operations, "power" for full access
            study_id: Optional study name to provide context for
            conversation_history: Optional recent messages for continuity
            canvas_state: Optional canvas state (nodes, edges) from the UI
            spec_path: Optional path to the atomizer_spec.json file

        Returns:
            Complete system prompt string
        """
        parts = [self._base_context(mode)]

        # Canvas context takes priority - if user is working on a canvas, include it
        if canvas_state:
            node_count = len(canvas_state.get("nodes", []))
            print(f"[ContextBuilder] Including canvas context with {node_count} nodes")
            parts.append(self._canvas_context(canvas_state, spec_path))
        else:
            print("[ContextBuilder] No canvas state provided")

        if study_id:
            parts.append(self._study_context(study_id))
        else:
            parts.append(self._global_context())

        if conversation_history:
            parts.append(self._conversation_context(conversation_history))

        parts.append(self._mode_instructions(mode, spec_path))

        return "\n\n---\n\n".join(parts)

    def build_study_context(self, study_id: str) -> str:
        """Build just the study context (for updates)"""
        return self._study_context(study_id)

    def _base_context(self, mode: str) -> str:
        """Base identity and capabilities"""
        return f"""# Atomizer Assistant

You are the Atomizer Assistant - an expert system for structural optimization using FEA.

**Current Mode**: {mode.upper()}

Your role:
- Help engineers with FEA optimization workflows
- Create, configure, and run optimization studies
- Analyze results and provide insights
- Explain FEA concepts and methodology

Important guidelines:
- Be concise and professional
- Use technical language appropriate for engineers
- You are "Atomizer Assistant", not a generic AI
- Use the available MCP tools to perform actions
- When asked about studies, use the appropriate tools to get real data
"""

    def _study_context(self, study_id: str) -> str:
        """Context for a specific study"""
        study_dir = self.studies_dir / study_id
        if not study_dir.exists():
            return f"# Current Study: {study_id}\n\n**Status**: Study directory not found."

        context = f"# Current Study: {study_id}\n\n"

        # Check for AtomizerSpec v2.0 first (preferred)
        spec_path = study_dir / "1_setup" / "atomizer_spec.json"
        if not spec_path.exists():
            spec_path = study_dir / "atomizer_spec.json"

        if spec_path.exists():
            context += self._spec_context(spec_path)
        else:
            # Fall back to legacy optimization_config.json
            context += self._legacy_config_context(study_dir)

        # Check for results
        db_path = study_dir / "3_results" / "study.db"
        if db_path.exists():
            try:
                conn = sqlite3.connect(db_path)
                count = conn.execute(
                    "SELECT COUNT(*) FROM trials WHERE state = 'COMPLETE'"
                ).fetchone()[0]

                best = conn.execute("""
                    SELECT MIN(tv.value) FROM trial_values tv
                    JOIN trials t ON tv.trial_id = t.trial_id
                    WHERE t.state = 'COMPLETE'
                """).fetchone()[0]

                context += f"\n## Results Status\n\n"
                context += f"- **Trials completed**: {count}\n"
                if best is not None:
                    context += f"- **Best objective**: {best:.6f}\n"

                conn.close()
            except Exception:
                pass

        return context

    def _spec_context(self, spec_path: Path) -> str:
        """Build context from AtomizerSpec v2.0 file"""
        context = "**Format**: AtomizerSpec v2.0\n\n"

        try:
            with open(spec_path) as f:
                spec = json.load(f)

            context += "## Configuration\n\n"

            # Design variables
            dvs = spec.get("design_variables", [])
            if dvs:
                context += "**Design Variables:**\n"
                for dv in dvs[:10]:
                    bounds = dv.get("bounds", {})
                    bound_str = f"[{bounds.get('min', '?')}, {bounds.get('max', '?')}]"
                    enabled = "✓" if dv.get("enabled", True) else "✗"
                    context += f"- {dv.get('name', 'unnamed')}: {bound_str} {enabled}\n"
                if len(dvs) > 10:
                    context += f"- ... and {len(dvs) - 10} more\n"

            # Extractors
            extractors = spec.get("extractors", [])
            if extractors:
                context += "\n**Extractors:**\n"
                for ext in extractors:
                    ext_type = ext.get("type", "unknown")
                    outputs = ext.get("outputs", [])
                    output_names = [o.get("name", "?") for o in outputs[:3]]
                    builtin = "builtin" if ext.get("builtin", True) else "custom"
                    context += f"- {ext.get('name', 'unnamed')} ({ext_type}, {builtin}): outputs {output_names}\n"

            # Objectives
            objs = spec.get("objectives", [])
            if objs:
                context += "\n**Objectives:**\n"
                for obj in objs:
                    direction = obj.get("direction", "minimize")
                    weight = obj.get("weight", 1.0)
                    context += f"- {obj.get('name', 'unnamed')} ({direction}, weight={weight})\n"

            # Constraints
            constraints = spec.get("constraints", [])
            if constraints:
                context += "\n**Constraints:**\n"
                for c in constraints:
                    op = c.get("operator", "<=")
                    thresh = c.get("threshold", "?")
                    context += f"- {c.get('name', 'unnamed')}: {op} {thresh}\n"

            # Optimization settings
            opt = spec.get("optimization", {})
            algo = opt.get("algorithm", {})
            budget = opt.get("budget", {})
            method = algo.get("type", "TPE")
            max_trials = budget.get("max_trials", "not set")
            context += f"\n**Optimization**: {method}, max_trials: {max_trials}\n"

            # Surrogate
            surrogate = opt.get("surrogate", {})
            if surrogate.get("enabled"):
                sur_type = surrogate.get("type", "gaussian_process")
                context += f"**Surrogate**: {sur_type} enabled\n"

        except (json.JSONDecodeError, IOError) as e:
            context += f"\n*Spec file exists but could not be parsed: {e}*\n"

        return context

    def _legacy_config_context(self, study_dir: Path) -> str:
        """Build context from legacy optimization_config.json"""
        context = "**Format**: Legacy optimization_config.json\n\n"

        config_path = study_dir / "1_setup" / "optimization_config.json"
        if not config_path.exists():
            config_path = study_dir / "optimization_config.json"

        if config_path.exists():
            try:
                with open(config_path) as f:
                    config = json.load(f)

                context += "## Configuration\n\n"

                # Design variables
                dvs = config.get("design_variables", [])
                if dvs:
                    context += "**Design Variables:**\n"
                    for dv in dvs[:10]:
                        bounds = f"[{dv.get('lower', '?')}, {dv.get('upper', '?')}]"
                        context += f"- {dv.get('name', 'unnamed')}: {bounds}\n"
                    if len(dvs) > 10:
                        context += f"- ... and {len(dvs) - 10} more\n"

                # Objectives
                objs = config.get("objectives", [])
                if objs:
                    context += "\n**Objectives:**\n"
                    for obj in objs:
                        direction = obj.get("direction", "minimize")
                        context += f"- {obj.get('name', 'unnamed')} ({direction})\n"

                # Constraints
                constraints = config.get("constraints", [])
                if constraints:
                    context += "\n**Constraints:**\n"
                    for c in constraints:
                        context += f"- {c.get('name', 'unnamed')}: {c.get('type', 'unknown')}\n"

                # Method
                method = config.get("method", "TPE")
                max_trials = config.get("max_trials", "not set")
                context += f"\n**Method**: {method}, max_trials: {max_trials}\n"

            except (json.JSONDecodeError, IOError) as e:
                context += f"\n*Config file exists but could not be parsed: {e}*\n"
        else:
            context += "*No configuration file found.*\n"

        return context

    def _global_context(self) -> str:
        """Context when no study is selected"""
        context = "# Available Studies\n\n"

        if self.studies_dir.exists():
            studies = [
                d.name
                for d in self.studies_dir.iterdir()
                if d.is_dir() and not d.name.startswith("_")
            ]

            if studies:
                context += "The following studies are available:\n\n"
                for name in sorted(studies)[:20]:
                    context += f"- {name}\n"
                if len(studies) > 20:
                    context += f"\n... and {len(studies) - 20} more\n"
            else:
                context += "No studies found. Use `create_study` tool to create one.\n"
        else:
            context += "Studies directory not found.\n"

        context += "\n## Quick Actions\n\n"
        context += "- **Create study**: Describe what you want to optimize\n"
        context += "- **List studies**: Use `list_studies` tool for details\n"
        context += "- **Open study**: Ask about a specific study by name\n"

        return context

    def _conversation_context(self, history: List[Dict[str, Any]]) -> str:
        """Recent conversation for continuity"""
        if not history:
            return ""

        context = "# Recent Conversation\n\n"
        for msg in history[-10:]:
            role = "User" if msg.get("role") == "user" else "Assistant"
            content = msg.get("content", "")[:500]
            if len(msg.get("content", "")) > 500:
                content += "..."
            context += f"**{role}**: {content}\n\n"

        return context

    def _canvas_context(self, canvas_state: Dict[str, Any], spec_path: Optional[str] = None) -> str:
        """
        Build context from canvas state (nodes and edges).

        This is CRITICAL for Claude to understand the current workflow
        being built in the Canvas UI.
        """
        context = "# Current Canvas State\n\n"
        context += "**You are assisting the user with a Canvas Builder workflow.**\n"
        context += "The canvas represents an optimization pipeline being configured visually.\n\n"

        nodes = canvas_state.get("nodes", [])
        edges = canvas_state.get("edges", [])
        study_name = canvas_state.get("studyName", "Untitled")
        study_path = canvas_state.get("studyPath", None)

        context += f"**Study Name**: {study_name}\n"
        if study_path:
            context += f"**Study Path**: {study_path}\n"
        if spec_path:
            context += f"**Spec File**: `{spec_path}`\n"
        context += "\n"

        # Group nodes by type
        node_types = {}
        for node in nodes:
            node_type = node.get("type", "unknown")
            if node_type not in node_types:
                node_types[node_type] = []
            node_types[node_type].append(node)

        # Model node
        if "model" in node_types:
            model = node_types["model"][0]
            data = model.get("data", {})
            context += "## Model\n"
            context += f"- **Label**: {data.get('label', 'Model')}\n"
            context += f"- **File Path**: {data.get('filePath', 'Not set')}\n"
            context += f"- **File Type**: {data.get('fileType', 'Not set')}\n\n"

        # Solver node
        if "solver" in node_types:
            solver = node_types["solver"][0]
            data = solver.get("data", {})
            context += "## Solver\n"
            context += f"- **Type**: {data.get('solverType', 'Not set')}\n\n"

        # Design variables
        if "designVar" in node_types:
            context += "## Design Variables\n\n"
            context += "| Name | Expression | Min | Max | Baseline | Unit | Enabled |\n"
            context += "|------|------------|-----|-----|----------|------|---------|\n"
            for dv in node_types["designVar"]:
                data = dv.get("data", {})
                name = data.get("label", "?")
                expr = data.get("expressionName", data.get("label", "?"))
                min_val = data.get("minValue", "?")
                max_val = data.get("maxValue", "?")
                baseline = data.get("baseline", "-")
                unit = data.get("unit", "-")
                enabled = "✓" if data.get("enabled", True) else "✗"
                context += f"| {name} | {expr} | {min_val} | {max_val} | {baseline} | {unit} | {enabled} |\n"
            context += "\n"

        # Extractors
        if "extractor" in node_types:
            context += "## Extractors\n\n"
            for ext in node_types["extractor"]:
                data = ext.get("data", {})
                context += f"### {data.get('extractorName', data.get('label', 'Extractor'))}\n"
                context += f"- **ID**: {data.get('extractorId', 'Not set')}\n"
                context += f"- **Type**: {data.get('extractorType', 'Not set')}\n"
                if data.get("extractMethod"):
                    context += f"- **Method**: {data.get('extractMethod')}\n"
                if data.get("innerRadius"):
                    context += f"- **Inner Radius**: {data.get('innerRadius')}\n"
                if data.get("nModes"):
                    context += f"- **Zernike Modes**: {data.get('nModes')}\n"
                if data.get("subcases"):
                    context += f"- **Subcases**: {data.get('subcases')}\n"
                if data.get("config"):
                    config = data.get("config", {})
                    if config.get("subcaseLabels"):
                        context += f"- **Subcase Labels**: {config.get('subcaseLabels')}\n"
                    if config.get("referenceSubcase"):
                        context += f"- **Reference Subcase**: {config.get('referenceSubcase')}\n"
                context += "\n"

        # Objectives
        if "objective" in node_types:
            context += "## Objectives\n\n"
            context += "| Name | Direction | Weight | Penalty |\n"
            context += "|------|-----------|--------|---------|\n"
            for obj in node_types["objective"]:
                data = obj.get("data", {})
                name = data.get("name", data.get("label", "?"))
                direction = data.get("direction", "minimize")
                weight = data.get("weight", 1)
                penalty = data.get("penaltyWeight", "-")
                context += f"| {name} | {direction} | {weight} | {penalty} |\n"
            context += "\n"

        # Constraints
        if "constraint" in node_types:
            context += "## Constraints\n\n"
            context += "| Name | Operator | Value |\n"
            context += "|------|----------|-------|\n"
            for con in node_types["constraint"]:
                data = con.get("data", {})
                name = data.get("name", data.get("label", "?"))
                operator = data.get("operator", "?")
                value = data.get("value", "?")
                context += f"| {name} | {operator} | {value} |\n"
            context += "\n"

        # Algorithm
        if "algorithm" in node_types:
            algo = node_types["algorithm"][0]
            data = algo.get("data", {})
            context += "## Algorithm\n"
            context += f"- **Method**: {data.get('method', 'Not set')}\n"
            context += f"- **Max Trials**: {data.get('maxTrials', 'Not set')}\n"
            if data.get("sigma0"):
                context += f"- **CMA-ES Sigma0**: {data.get('sigma0')}\n"
            if data.get("restartStrategy"):
                context += f"- **Restart Strategy**: {data.get('restartStrategy')}\n"
            context += "\n"

        # Surrogate
        if "surrogate" in node_types:
            sur = node_types["surrogate"][0]
            data = sur.get("data", {})
            context += "## Surrogate\n"
            context += f"- **Enabled**: {data.get('enabled', False)}\n"
            context += f"- **Type**: {data.get('modelType', 'Not set')}\n"
            context += f"- **Min Trials**: {data.get('minTrials', 'Not set')}\n\n"

        # Edge connections summary
        context += "## Connections\n\n"
        context += f"Total edges: {len(edges)}\n"
        context += "Flow: Design Variables → Model → Solver → Extractors → Objectives/Constraints → Algorithm\n\n"

        # Instructions will be in _mode_instructions based on spec_path
        return context

    def _mode_instructions(self, mode: str, spec_path: Optional[str] = None) -> str:
        """Mode-specific instructions"""
        if mode == "power":
            instructions = """# Power Mode Instructions

You have **FULL ACCESS** to modify Atomizer studies. **DO NOT ASK FOR PERMISSION** - just do it.

## CRITICAL: How to Modify the Spec

"""
            if spec_path:
                instructions += f"""**The spec file is at**: `{spec_path}`

When asked to add/modify/remove design variables, extractors, objectives, or constraints:
1. **Read the spec file first** using the Read tool
2. **Edit the spec file** using the Edit tool to make precise changes
3. **Confirm what you changed** in your response

### AtomizerSpec v2.0 Structure

The spec has these main arrays you can modify:
- `design_variables` - Parameters to optimize
- `extractors` - Physics extraction functions  
- `objectives` - What to minimize/maximize
- `constraints` - Limits that must be satisfied

### Example: Add a Design Variable

To add a design variable called "thickness" with bounds [1, 10]:

1. Read the spec: `Read({spec_path})`
2. Find the `"design_variables": [...]` array
3. Add a new entry like:
```json
{{
  "id": "dv_thickness",
  "name": "thickness",
  "expression_name": "thickness",
  "type": "continuous",
  "bounds": {{"min": 1, "max": 10}},
  "baseline": 5,
  "units": "mm",
  "enabled": true
}}
```
4. Use Edit tool to insert this into the array

### Example: Add an Objective

To add a "minimize mass" objective:
```json
{{
  "id": "obj_mass",
  "name": "mass",
  "direction": "minimize",
  "weight": 1.0,
  "source": {{
    "extractor_id": "ext_mass",
    "output_name": "mass"
  }}
}}
```

### Example: Add an Extractor

To add a mass extractor:
```json
{{
  "id": "ext_mass",
  "name": "mass",
  "type": "mass",
  "builtin": true,
  "outputs": [{{"name": "mass", "units": "kg"}}]
}}
```

"""
            else:
                instructions += """No spec file is currently set. Ask the user which study they want to work with.

"""

            instructions += """## IMPORTANT Rules:
- You have --dangerously-skip-permissions enabled
- **ACT IMMEDIATELY** when asked to add/modify/remove things
- Use the **Edit** tool to modify the spec file directly
- Generate unique IDs like `dv_<name>`, `ext_<name>`, `obj_<name>`, `con_<name>`
- Explain what you changed AFTER doing it, not before
- Do NOT say "I need permission" - you already have it
"""
            return instructions
        else:
            return """# User Mode Instructions

You can help with optimization workflows:
- Create and configure studies
- Run optimizations
- Analyze results
- Generate reports
- Explain FEA concepts

**For modifying studies**, the user needs to switch to Power Mode.

In user mode you can:
- Read and explain study configurations
- Analyze optimization results
- Provide recommendations
- Answer questions about FEA and optimization
"""