"""Self-Extension Engine — builds new skills from natural language descriptions."""

from __future__ import annotations

import ast
import json
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from rich.console import Console
from rich.panel import Panel
from rich.syntax import Syntax

from alive.core.replay import BirthStoryRecorder, save_birth_story
from alive.core.skill_registry import SkillInfo, SkillRegistry, get_registry
from alive.llm.base import LLMProvider
from alive.sandbox.dependency import DependencyManager, get_dependency_manager
from alive.sandbox.runner import run_skill_code

MAX_ATTEMPTS = 5


# ── Data structures ─────────────────────────────────────────────────────────


@dataclass
class BuildResult:
    """Parsed output from the LLM generation phase."""

    skill_name: str
    description: str
    code: str
    test_code: str
    dependencies: list[str]
    input_schema: dict[str, Any]
    output_schema: dict[str, Any]
    skill_doc: str
    sample_input: dict[str, Any]


@dataclass
class BuildAttempt:
    """Record of one build attempt."""

    attempt_number: int
    code: str
    test_code: str
    success: bool
    error: str | None = None
    stdout: str | None = None
    stderr: str | None = None


# ── LLM prompt templates ────────────────────────────────────────────────────

_ANALYZE_SYSTEM = """\
You are Alive's skill analyzer. You analyze capability requests to plan \
Python skill implementations.

Given a capability description and a list of existing skills, respond with \
a brief analysis covering:
- What the skill needs to do (inputs, outputs, edge cases)
- What Python libraries would be best suited
- Whether any existing skills can be composed or referenced
- A proposed function signature for run()

Be concise. No code yet — just the plan."""

_GENERATE_SYSTEM = """\
You are Alive's skill generator. You write Python skills that follow a strict template.

SKILL TEMPLATE RULES:
- A single file with a `run()` function as the entry point
- The function returns a dict
- Type hints on all parameters and return type
- Handle errors gracefully — return {{"success": False, "error": "..."}} \
rather than raising exceptions
- Only use standard library OR well-known PyPI packages
- Never use input() or any interactive I/O
- Keep the code focused — one skill, one purpose

You MUST respond in EXACTLY this format with these labeled code blocks:

```name
skill_name_here
```

```description
One-line description here.
```

```dependencies
package1
package2
```

```input_schema
{{"param1": "str", "param2": "int (optional, default 10)"}}
```

```output_schema
{{"result": "str", "success": "bool"}}
```

```python skill
# The complete skill.py contents
```

```python test
# Tests — must use: from skill import run
```

```python sample_input
{{"param1": "example value"}}
```

```markdown doc
# skill_name
Full SKILL.md documentation here.
```

IMPORTANT:
- The skill name must be valid snake_case
- Tests must import with `from skill import run` (not a package path)
- sample_input must be a valid Python dict that exercises the happy path
- Dependencies must NOT include standard library modules"""

_FIX_SYSTEM = _GENERATE_SYSTEM + """

You are fixing a skill that failed testing. You will receive the original \
code, the test code, and the error output. Fix the code to make it work.

IMPORTANT: Respond with the COMPLETE set of blocks again (not just a diff)."""


# ── SkillBuilder class ──────────────────────────────────────────────────────


class SkillBuilder:
    """Builds a new skill from a natural language description."""

    def __init__(
        self,
        description: str,
        *,
        console: Console | None = None,
        verbose: bool = False,
        provider: LLMProvider | None = None,
        registry: SkillRegistry | None = None,
    ) -> None:
        self._description = description
        self._console = console or Console()
        self._verbose = verbose
        self._provider = provider
        self._registry = registry
        self._attempts: list[BuildAttempt] = []
        self._recorder: BirthStoryRecorder | None = None

    async def build(self) -> SkillInfo:
        """Run the full build pipeline: analyze -> generate -> test -> retry -> save."""
        provider = self._get_provider()
        self._console.print(
            Panel(
                f"[bold]Building skill:[/] {self._description}\n"
                f"[dim]Provider: {provider.name}[/]",
                title="[bold cyan]Self-Extension Engine[/]",
                border_style="cyan",
            )
        )

        # Phase 1: Analyze
        with self._console.status("[bold cyan]  Analyzing capability...", spinner="dots"):
            analysis = await self._analyze()

        # Phase 2: Generate
        with self._console.status("[bold cyan]  Generating code...", spinner="dots"):
            result = await self._generate(analysis)

        if self._verbose:
            self._show_code(result)

        # Start recording the birth story
        self._recorder = BirthStoryRecorder(
            skill_name=result.skill_name,
            original_request=self._description,
            plan_step=analysis[:200],
        )

        # Phase 3+4: Test with retry loop
        for attempt_num in range(1, MAX_ATTEMPTS + 1):
            self._console.print(f"\n[bold cyan]  Attempt {attempt_num}/{MAX_ATTEMPTS}[/]")

            with self._console.status("[bold cyan]  Installing dependencies...", spinner="dots"):
                self._install_dependencies(result.dependencies)

            with self._console.status("[bold cyan]  Testing in sandbox...", spinner="dots"):
                attempt = self._test(result)
            self._attempts.append(attempt)

            # Record attempt in birth story
            if self._recorder:
                self._recorder.record_attempt(
                    attempt_num,
                    result.code,
                    success=attempt.success,
                    error=attempt.error,
                )

            if attempt.success:
                self._console.print("[bold green]  Tests passed.[/]")
                break

            # Show error
            error_text = attempt.stderr or attempt.error or "Unknown error"
            self._console.print(
                Panel(
                    error_text[:500],
                    title=f"[bold red]Test Failed (attempt {attempt_num})[/]",
                    border_style="red",
                )
            )

            if attempt_num == MAX_ATTEMPTS:
                self._console.print(
                    Panel(
                        f"Failed to build skill after {MAX_ATTEMPTS} attempts.\n"
                        f"Description: {self._description}",
                        title="[bold red]Build Failed[/]",
                        border_style="red",
                    )
                )
                # Record the gap for self-model tracking
                try:
                    from alive.core.self_model import record_gap

                    record_gap(self._description, error_text[:200])
                except Exception as exc:
                    import logging
                    logging.getLogger("alive").debug("Gap tracking failed: %s", exc)
                raise RuntimeError(
                    f"Failed to build skill '{result.skill_name}' "
                    f"after {MAX_ATTEMPTS} attempts. "
                    f"Last error: {error_text[:200]}"
                )

            # Fix
            with self._console.status("[bold yellow]  Fixing code...", spinner="dots"):
                result = await self._fix(result, attempt)

            if self._verbose:
                self._show_code(result)

        # Phase 5: Save
        with self._console.status("[bold cyan]  Saving skill...", spinner="dots"):
            skill_info = self._save(result)

        # Save the birth story
        if self._recorder:
            try:
                n_attempts = len(self._attempts)
                approach = (
                    f"Built successfully on first try."
                    if n_attempts == 1
                    else f"Took {n_attempts} attempts. Final code resolved earlier failures."
                )
                story = self._recorder.finalize(final_approach=approach)
                save_birth_story(story)
            except Exception as exc:
                import logging
                logging.getLogger("alive").debug("Failed to save birth story: %s", exc)

        self._console.print(
            Panel(
                f"[bold green]{result.skill_name}[/]\n\n"
                f"{result.description}\n\n"
                f"[dim]Dependencies:[/] {', '.join(result.dependencies) or 'none'}\n"
                f"[dim]Attempts:[/] {len(self._attempts)}\n"
                f"[dim]Location:[/] ~/.alive/skills/{result.skill_name}/",
                title="[bold green]Skill Built Successfully[/]",
                border_style="green",
            )
        )
        return skill_info

    # ── Phase methods ────────────────────────────────────────────────────

    async def _analyze(self) -> str:
        """Phase 1: Analyze the capability request."""
        registry = self._registry or get_registry()
        existing = registry.search(self._description)[:5]

        if existing:
            skills_text = "\n".join(
                f"- {s.name}: {s.description} "
                f"(inputs: {s.input_schema}, outputs: {s.output_schema})"
                for s in existing
            )
        else:
            skills_text = "(none relevant)"

        user_prompt = (
            f'Capability requested: "{self._description}"\n\n'
            f"Existing skills that might help:\n{skills_text}\n\n"
            f"Analyze this capability and plan the implementation."
        )
        return await self._llm(_ANALYZE_SYSTEM, user_prompt)

    async def _generate(self, analysis: str) -> BuildResult:
        """Phase 2: Ask LLM to generate skill code."""
        user_prompt = (
            f'Build a skill for: "{self._description}"\n\n'
            f"Analysis:\n{analysis}"
        )
        response = await self._llm(_GENERATE_SYSTEM, user_prompt)
        return self._parse_llm_response(response)

    def _test(self, result: BuildResult) -> BuildAttempt:
        """Phase 3: Run generated code in a subprocess sandbox."""
        dm = self._get_dep_manager()
        sr = run_skill_code(
            result.code,
            result.sample_input,
            dep_manager=dm,
        )
        return BuildAttempt(
            attempt_number=len(self._attempts) + 1,
            code=result.code,
            test_code=result.test_code,
            success=sr.success,
            stdout=json.dumps(sr.output, default=str) if sr.output else None,
            stderr=sr.error,
            error=sr.error,
        )

    async def _fix(self, result: BuildResult, attempt: BuildAttempt) -> BuildResult:
        """Phase 4: Send error back to LLM for fixing."""
        attempt_num = len(self._attempts)
        different = (
            "\n\nThe previous approach is NOT WORKING. "
            "Try a FUNDAMENTALLY DIFFERENT approach — "
            "different algorithm, different library, different strategy."
            if attempt_num >= 3
            else ""
        )

        error_output = (attempt.stderr or "") + (attempt.stdout or "")
        user_prompt = (
            f"The skill failed testing. Fix it.\n\n"
            f"ORIGINAL CODE:\n```python\n{result.code}\n```\n\n"
            f"TEST CODE:\n```python\n{result.test_code}\n```\n\n"
            f"ERROR OUTPUT:\n```\n{error_output[:2000]}\n```\n\n"
            f"Attempt: {attempt_num}/{MAX_ATTEMPTS}"
            f"{different}\n\n"
            f"Respond with the COMPLETE set of blocks again."
        )
        response = await self._llm(_FIX_SYSTEM, user_prompt)
        return self._parse_llm_response(response)

    def _install_dependencies(self, dependencies: list[str]) -> None:
        """Install dependencies into the shared venv."""
        if not dependencies:
            return
        dm = self._get_dep_manager()
        dm.install(dependencies)

    def _save(self, result: BuildResult) -> SkillInfo:
        """Phase 5: Register the completed skill."""
        registry = self._registry or get_registry()

        # Handle name collision
        name = result.skill_name
        if name in {s.name for s in registry.list_all()}:
            for i in range(2, 100):
                candidate = f"{name}_{i}"
                if candidate not in {s.name for s in registry.list_all()}:
                    name = candidate
                    break

        return registry.register(
            name=name,
            code=result.code,
            description=result.description,
            dependencies=result.dependencies,
            input_schema=result.input_schema,
            output_schema=result.output_schema,
            sample_input=result.sample_input,
            test_code=result.test_code,
            skill_doc=result.skill_doc,
        )

    # ── Helpers ──────────────────────────────────────────────────────────

    def _get_provider(self) -> LLMProvider:
        if self._provider is None:
            from alive.llm import get_provider

            self._provider = get_provider()
        return self._provider

    async def _llm(self, system_prompt: str, user_prompt: str) -> str:
        provider = self._get_provider()
        return await provider.complete(system_prompt, user_prompt)

    def _get_dep_manager(self) -> DependencyManager:
        """Return the dependency manager (lazy singleton)."""
        return get_dependency_manager()

    def _get_venv_python(self) -> Path:
        """Return path to the venv Python binary."""
        dm = self._get_dep_manager()
        dm.ensure_venv()
        return dm.python

    def _parse_llm_response(self, response: str) -> BuildResult:
        """Parse the structured LLM response into a BuildResult."""
        code = _extract_block("skill", response)
        if not code:
            raise ValueError("LLM response missing skill code block")

        raw_name = _extract_block("name", response) or self._derive_name()
        skill_name = re.sub(r"[^a-z0-9_]", "", raw_name.lower().replace("-", "_").replace(" ", "_"))
        skill_name = skill_name.strip("_") or "unnamed_skill"

        description = _extract_block("description", response) or self._description

        test_code = _extract_block("test", response) or (
            f'"""Auto-generated tests for {skill_name}."""\n\n'
            f"from skill import run\n\n\n"
            f"def test_{skill_name}_runs():\n"
            f"    assert callable(run)\n"
        )

        skill_doc = _extract_block("doc", response) or f"# {skill_name}\n\n{description}\n"

        deps_raw = _extract_block("dependencies", response)
        dependencies = [d.strip() for d in deps_raw.splitlines() if d.strip()] if deps_raw else []

        input_schema = _safe_parse_json(_extract_block("input_schema", response))
        output_schema = _safe_parse_json(_extract_block("output_schema", response))

        sample_raw = _extract_block("sample_input", response)
        sample_input = _safe_eval_dict(sample_raw) if sample_raw else {}

        return BuildResult(
            skill_name=skill_name,
            description=description,
            code=code,
            test_code=test_code,
            dependencies=dependencies,
            input_schema=input_schema,
            output_schema=output_schema,
            skill_doc=skill_doc,
            sample_input=sample_input,
        )

    def _derive_name(self) -> str:
        """Derive a skill name from the description."""
        words = re.findall(r"[a-z]+", self._description.lower())
        # Filter stop words
        stop = {"a", "an", "the", "to", "from", "and", "or", "of", "in", "for", "with", "that", "this", "is"}
        significant = [w for w in words if w not in stop][:4]
        return "_".join(significant) or "unnamed_skill"

    def _show_code(self, result: BuildResult) -> None:
        """Display generated code with syntax highlighting (verbose mode)."""
        self._console.print(
            Panel(
                Syntax(result.code, "python", theme="monokai", line_numbers=True),
                title=f"[bold]{result.skill_name}/skill.py[/]",
                border_style="dim",
            )
        )


# ── Parsing helpers (module-level) ──────────────────────────────────────────


def _extract_block(label: str, response: str) -> str:
    """Extract content from a labeled markdown code block."""
    patterns = [
        # ```python skill  or  ```python test  etc.
        rf"```\w+\s+{re.escape(label)}\s*\n(.*?)```",
        # ```markdown doc
        rf"```\w+\s+{re.escape(label)}\s*\n(.*?)```",
        # ```name  or  ```description  etc.
        rf"```{re.escape(label)}\s*\n(.*?)```",
    ]
    for pattern in patterns:
        match = re.search(pattern, response, re.DOTALL)
        if match:
            return match.group(1).strip()
    return ""


def _safe_parse_json(text: str) -> dict[str, Any]:
    """Parse JSON, falling back to ast.literal_eval for Python dicts."""
    if not text:
        return {}
    try:
        result = json.loads(text)
        return result if isinstance(result, dict) else {}
    except json.JSONDecodeError:
        try:
            result = ast.literal_eval(text)
            return result if isinstance(result, dict) else {}
        except (ValueError, SyntaxError, TypeError):
            return {}


def _safe_eval_dict(text: str) -> dict[str, Any]:
    """Safely evaluate a Python dict literal."""
    if not text:
        return {}
    try:
        result = ast.literal_eval(text)
        return result if isinstance(result, dict) else {}
    except (ValueError, SyntaxError, TypeError):
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            return {}


# ── Public API ──────────────────────────────────────────────────────────────


async def build_skill(
    description: str,
    *,
    console: Console | None = None,
    verbose: bool = False,
    provider: LLMProvider | None = None,
    registry: SkillRegistry | None = None,
) -> SkillInfo:
    """Build a new skill from a natural language capability description.

    This is the main entry point for the Self-Extension Engine. It:
    1. Analyzes the capability request
    2. Generates Python code via LLM
    3. Tests the code in a sandboxed subprocess
    4. Retries up to 5 times on failure
    5. Saves the skill to the registry on success

    Args:
        description: What the skill should do, in plain English.
        console: Rich console for output. Defaults to a new Console.
        verbose: If True, show generated code with syntax highlighting.
        provider: LLM provider override. Defaults to ALIVE_PROVIDER.
        registry: SkillRegistry override. Defaults to global registry.

    Returns:
        SkillInfo for the newly registered skill.

    Raises:
        RuntimeError: If the skill could not be built after 5 attempts.
    """
    builder = SkillBuilder(
        description,
        console=console,
        verbose=verbose,
        provider=provider,
        registry=registry,
    )
    return await builder.build()
