"""Planner — decomposes natural language requests into executable steps."""

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from typing import Any, Literal

from rich.console import Console
from rich.panel import Panel
from rich.tree import Tree

from alive.core.skill_registry import SkillRegistry, get_registry
from alive.llm.base import LLMProvider

StepType = Literal["EXISTING_SKILL", "NEW_SKILL", "LLM_ONLY"]
VALID_STEP_TYPES = {"EXISTING_SKILL", "NEW_SKILL", "LLM_ONLY"}


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


@dataclass
class PlanStep:
    """A single step in an execution plan."""

    type: StepType
    description: str
    skill: str | None = None  # EXISTING_SKILL: which skill to run
    input: dict[str, Any] | None = None  # EXISTING_SKILL: kwargs for run()
    capability: str | None = None  # NEW_SKILL: what the builder should create
    prompt: str | None = None  # LLM_ONLY: text generation prompt


@dataclass
class ExecutionPlan:
    """A decomposed plan for fulfilling a user request."""

    steps: list[PlanStep]
    reasoning: str
    original_request: str


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

_PLAN_SYSTEM = """\
You are the planner for Alive, a self-evolving AI agent. Given a \
user request and a list of available skills, decompose the request \
into steps. For each step, either:
- EXISTING_SKILL: name the skill to use and what input to give it
- NEW_SKILL: describe the capability needed (the builder will create it)
- LLM_ONLY: this step just needs text generation, no skill needed

Respond in JSON:
{
  "steps": [
    {"type": "EXISTING_SKILL", "skill": "web_fetch", \
"input": {"url": "..."}, "description": "Fetch the webpage"},
    {"type": "NEW_SKILL", \
"capability": "Extract prices from HTML using CSS selectors", \
"description": "Parse the fetched HTML to find product prices"},
    {"type": "LLM_ONLY", \
"prompt": "Summarize these prices and recommend the cheapest", \
"description": "Analyze and summarize results"}
  ],
  "reasoning": "Brief explanation of why this plan works"
}

IMPORTANT: Respond with ONLY valid JSON. No markdown, no explanation outside the JSON."""

_RETRY_SUFFIX = (
    "\n\nYour previous response was not valid JSON. "
    "Respond with ONLY a JSON object, no markdown formatting, "
    "no code blocks, no extra text."
)


# ── JSON parsing ─────────────────────────────────────────────────────────────


def _extract_json(response: str) -> dict[str, Any]:
    """Extract a JSON object from an LLM response.

    Tries direct parsing first, then looks for JSON inside markdown code blocks.
    Raises ValueError if no valid JSON found.
    """
    # Try direct parse
    stripped = response.strip()
    try:
        data = json.loads(stripped)
        if isinstance(data, dict):
            return data
    except json.JSONDecodeError:
        pass

    # Try extracting from markdown code block
    match = re.search(r"```(?:json)?\s*\n(.*?)```", stripped, re.DOTALL)
    if match:
        try:
            data = json.loads(match.group(1).strip())
            if isinstance(data, dict):
                return data
        except json.JSONDecodeError:
            pass

    raise ValueError(f"Could not parse JSON from response: {stripped[:200]}")


def _parse_plan_response(response: str, original_request: str) -> ExecutionPlan:
    """Parse an LLM JSON response into an ExecutionPlan.

    Raises ValueError if the response is malformed or missing required fields.
    """
    data = _extract_json(response)

    if "steps" not in data or not isinstance(data["steps"], list):
        raise ValueError("Response missing 'steps' array")

    if not data["steps"]:
        raise ValueError("Response has empty 'steps' array")

    steps: list[PlanStep] = []
    for i, step_data in enumerate(data["steps"]):
        if not isinstance(step_data, dict):
            raise ValueError(f"Step {i} is not a dict")

        step_type = step_data.get("type", "")
        if step_type not in VALID_STEP_TYPES:
            raise ValueError(
                f"Step {i} has invalid type '{step_type}'. "
                f"Must be one of: {', '.join(sorted(VALID_STEP_TYPES))}"
            )

        description = step_data.get("description", "")
        if not description:
            raise ValueError(f"Step {i} missing 'description'")

        if step_type == "EXISTING_SKILL":
            skill = step_data.get("skill")
            if not skill:
                raise ValueError(
                    f"Step {i} (EXISTING_SKILL) missing 'skill' name"
                )
            steps.append(
                PlanStep(
                    type="EXISTING_SKILL",
                    description=description,
                    skill=skill,
                    input=step_data.get("input") or {},
                )
            )
        elif step_type == "NEW_SKILL":
            capability = step_data.get("capability")
            if not capability:
                raise ValueError(
                    f"Step {i} (NEW_SKILL) missing 'capability' description"
                )
            steps.append(
                PlanStep(
                    type="NEW_SKILL",
                    description=description,
                    capability=capability,
                )
            )
        else:  # LLM_ONLY
            steps.append(
                PlanStep(
                    type="LLM_ONLY",
                    description=description,
                    prompt=step_data.get("prompt") or description,
                )
            )

    reasoning = data.get("reasoning", "")
    return ExecutionPlan(
        steps=steps,
        reasoning=reasoning,
        original_request=original_request,
    )


# ── Display ──────────────────────────────────────────────────────────────────

_STEP_STYLES = {
    "EXISTING_SKILL": "green",
    "NEW_SKILL": "yellow",
    "LLM_ONLY": "blue",
}


def display_plan(plan: ExecutionPlan, console: Console | None = None) -> None:
    """Render an execution plan as a Rich tree."""
    console = console or Console()

    tree = Tree(
        f"[bold]Execution Plan[/] for: [dim]{plan.original_request}[/]",
    )

    for i, step in enumerate(plan.steps, 1):
        style = _STEP_STYLES.get(step.type, "white")
        num = f"[{style}]{chr(0x2460 + i - 1)}[/{style}]"  # ①②③...

        if step.type == "EXISTING_SKILL":
            label = (
                f"{num} [{style}]EXISTING_SKILL[/{style}] "
                f"[bold]{step.skill}[/bold]: {step.description}"
            )
        elif step.type == "NEW_SKILL":
            label = (
                f"{num} [{style}]NEW_SKILL[/{style}]: "
                f"{step.description} "
                f'[dim](needs: "{step.capability}")[/dim]'
            )
        else:
            label = (
                f"{num} [{style}]LLM_ONLY[/{style}]: {step.description}"
            )

        tree.add(label)

    console.print(tree)

    if plan.reasoning:
        console.print(f"\n[dim]Reasoning: {plan.reasoning}[/]")


# ── Fallback ─────────────────────────────────────────────────────────────────


def _fallback_plan(request: str) -> ExecutionPlan:
    """Create a single LLM_ONLY fallback plan when parsing fails."""
    return ExecutionPlan(
        steps=[
            PlanStep(
                type="LLM_ONLY",
                description="Respond to the user's request directly",
                prompt=request,
            )
        ],
        reasoning="Could not decompose the request into steps; "
        "falling back to direct LLM response.",
        original_request=request,
    )


# ── Core planner function ───────────────────────────────────────────────────


async def plan_request(
    request: str,
    *,
    provider: LLMProvider | None = None,
    registry: SkillRegistry | None = None,
    console: Console | None = None,
) -> ExecutionPlan:
    """Decompose a natural language request into an executable plan.

    Sends the request and available skills to the LLM, parses the JSON
    response into an ExecutionPlan. Retries once on parse failure, then
    falls back to a single LLM_ONLY step.

    Args:
        request: The user's natural language request.
        provider: LLM provider override. Defaults to ALIVE_PROVIDER.
        registry: SkillRegistry override. Defaults to global registry.
        console: Rich console for status output.

    Returns:
        ExecutionPlan with one or more steps.
    """
    if provider is None:
        from alive.llm import get_provider

        provider = get_provider()

    registry = registry or get_registry()
    con = console or Console()

    # Build skills context
    skills = registry.list_all()
    if skills:
        skills_text = "\n".join(
            f"- {s.name}: {s.description} (inputs: {s.input_schema})"
            for s in skills
        )
    else:
        skills_text = "(no skills available yet)"

    user_prompt = (
        f'User request: "{request}"\n\n'
        f"Available skills:\n{skills_text}\n\n"
        f"Decompose this request into steps."
    )

    # First attempt
    with con.status("[bold cyan]  Planning...", spinner="dots"):
        response = await provider.complete(_PLAN_SYSTEM, user_prompt)

    try:
        return _parse_plan_response(response, request)
    except ValueError:
        pass

    # Retry with explicit JSON instruction
    with con.status("[bold yellow]  Retrying plan...", spinner="dots"):
        response = await provider.complete(
            _PLAN_SYSTEM, user_prompt + _RETRY_SUFFIX
        )

    try:
        return _parse_plan_response(response, request)
    except ValueError:
        return _fallback_plan(request)
