"""Composer — executes an ExecutionPlan step by step."""

from __future__ import annotations

import json
import time
from dataclasses import dataclass
from typing import Any, Callable, Literal

from rich.console import Console
from rich.panel import Panel
from rich.prompt import Prompt

from alive.core.planner import ExecutionPlan, PlanStep, display_plan
from alive.core.skill_registry import SkillRegistry, get_registry
from alive.llm.base import LLMProvider

ErrorAction = Literal["retry", "skip", "abort"]


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


@dataclass
class StepResult:
    """Result of executing a single plan step."""

    step_index: int
    step: PlanStep
    success: bool
    output: str
    error: str | None = None
    duration_ms: int = 0
    skipped: bool = False


@dataclass
class CompositionResult:
    """Final result of executing an entire plan."""

    plan: ExecutionPlan
    step_results: list[StepResult]
    final_output: str
    success: bool


class CompositionError(RuntimeError):
    """Raised when composition is aborted due to a step failure."""

    def __init__(self, step_index: int, step: PlanStep, error: str) -> None:
        self.step_index = step_index
        self.step = step
        super().__init__(
            f"Step {step_index + 1} ({step.description}) failed: {error}"
        )


# ── Error handling ───────────────────────────────────────────────────────────


def _interactive_error_handler(
    step_index: int,
    step: PlanStep,
    error: str,
    console: Console,
) -> ErrorAction:
    """Prompt the user for retry/skip/abort via the terminal."""
    console.print(
        Panel(
            f"{error[:500]}",
            title=f"[bold red]Step {step_index + 1} failed: {step.description}[/]",
            border_style="red",
        )
    )
    try:
        choice = Prompt.ask(
            "[bold]What should I do?[/]",
            choices=["retry", "skip", "abort"],
            default="retry",
            console=console,
        )
        return choice  # type: ignore[return-value]
    except EOFError:
        console.print("[dim]Non-interactive mode — aborting.[/]")
        return "abort"


# ── LLM prompts ─────────────────────────────────────────────────────────────

_LLM_ONLY_SYSTEM = """\
You are Alive, a helpful AI assistant executing one step of a multi-step plan. \
Respond concisely and helpfully."""

_SYNTHESIZE_SYSTEM = """\
You are Alive, a helpful AI assistant. Synthesize the results of a multi-step \
plan into a clear, helpful response for the user. Be concise but thorough."""


# ── Composer class ───────────────────────────────────────────────────────────


class Composer:
    """Executes an ExecutionPlan step by step, chaining outputs."""

    def __init__(
        self,
        *,
        provider: LLMProvider | None = None,
        registry: SkillRegistry | None = None,
        console: Console | None = None,
        on_error: Callable[[int, PlanStep, str, Console], ErrorAction] | None = None,
    ) -> None:
        self._provider = provider
        self._registry = registry
        self._console = console or Console()
        self._on_error = on_error or _interactive_error_handler
        self._step_results: list[StepResult] = []

    async def execute(self, plan: ExecutionPlan) -> CompositionResult:
        """Execute an entire plan step by step.

        Returns CompositionResult with all step results and a final synthesized output.
        Raises CompositionError if a step fails and the user chooses to abort.
        """
        self._step_results = []
        provider = self._get_provider()
        registry = self._registry or get_registry()
        n = len(plan.steps)

        # Show the plan
        display_plan(plan, console=self._console)
        self._console.print()

        for i, step in enumerate(plan.steps):
            step_label = f"Step {i + 1}/{n}: {step.description}"

            while True:  # retry loop
                with self._console.status(
                    f"[bold cyan]  {step_label}...", spinner="dots"
                ):
                    start = time.monotonic()
                    result = await self._execute_step(
                        i, step, provider, registry
                    )
                    result.duration_ms = int(
                        (time.monotonic() - start) * 1000
                    )

                if result.success or result.skipped:
                    icon = (
                        "[yellow]⊘[/yellow]"
                        if result.skipped
                        else "[green]✓[/green]"
                    )
                    self._console.print(f"  {icon} {step_label}")
                    self._step_results.append(result)
                    break

                # Step failed — ask user
                self._console.print(
                    f"  [red]✗[/red] {step_label}"
                )
                action = self._on_error(
                    i, step, result.error or "Unknown error", self._console
                )

                if action == "retry":
                    self._console.print(
                        f"  [bold yellow]↻[/bold yellow] Retrying step {i + 1}..."
                    )
                    continue
                elif action == "skip":
                    result.skipped = True
                    self._step_results.append(result)
                    self._console.print(
                        f"  [yellow]⊘[/yellow] Skipped step {i + 1}"
                    )
                    break
                else:  # abort
                    raise CompositionError(
                        i, step, result.error or "Unknown error"
                    )

        # Synthesize final output
        with self._console.status(
            "[bold cyan]  Synthesizing results...", spinner="dots"
        ):
            final = await self._synthesize(plan, provider)

        success = all(
            r.success or r.skipped for r in self._step_results
        )

        return CompositionResult(
            plan=plan,
            step_results=self._step_results,
            final_output=final,
            success=success,
        )

    # ── Step handlers ────────────────────────────────────────────────────

    async def _execute_step(
        self,
        index: int,
        step: PlanStep,
        provider: LLMProvider,
        registry: SkillRegistry,
    ) -> StepResult:
        """Dispatch a step to the appropriate handler."""
        if step.type == "EXISTING_SKILL":
            return self._run_existing_skill(index, step, registry)
        elif step.type == "NEW_SKILL":
            return await self._run_new_skill(index, step, provider, registry)
        else:  # LLM_ONLY
            return await self._run_llm_only(index, step, provider)

    def _run_existing_skill(
        self,
        index: int,
        step: PlanStep,
        registry: SkillRegistry,
    ) -> StepResult:
        """Execute an existing skill via the sandbox runner."""
        from alive.sandbox.runner import run_skill

        sr = run_skill(step.skill or "", step.input)

        # Track usage stats
        try:
            registry.update_stats(step.skill or "", success=sr.success)
        except KeyError:
            pass  # skill might not be in registry

        # Auto-suggest evolution for slow skills
        if sr.success and sr.duration_ms > 5000 and self._provider:
            try:
                import asyncio

                from alive.core.evolve import check_and_suggest

                asyncio.get_event_loop().create_task(
                    check_and_suggest(
                        step.skill or "",
                        sr.duration_ms,
                        provider=self._provider,
                        console=self._console,
                    )
                )
            except Exception:
                pass  # best-effort, don't break the pipeline

        if sr.success:
            return StepResult(
                step_index=index,
                step=step,
                success=True,
                output=json.dumps(sr.output, default=str) if sr.output else "{}",
                duration_ms=sr.duration_ms,
            )
        return StepResult(
            step_index=index,
            step=step,
            success=False,
            output="",
            error=sr.error or "Skill execution failed",
            duration_ms=sr.duration_ms,
        )

    async def _run_new_skill(
        self,
        index: int,
        step: PlanStep,
        provider: LLMProvider,
        registry: SkillRegistry,
    ) -> StepResult:
        """Build a new skill via skill_builder, then run it."""
        from alive.core.skill_builder import build_skill
        from alive.sandbox.runner import run_skill

        try:
            info = await build_skill(
                step.capability or step.description,
                console=self._console,
                provider=provider,
                registry=registry,
            )
        except RuntimeError as e:
            # Gap already recorded by skill_builder; just return failure
            return StepResult(
                step_index=index,
                step=step,
                success=False,
                output="",
                error=f"Failed to build skill: {e}",
            )

        # Run the newly built skill with its sample input
        sr = run_skill(info.name, info.sample_input or {})
        try:
            registry.update_stats(info.name, success=sr.success)
        except KeyError:
            pass

        if sr.success:
            return StepResult(
                step_index=index,
                step=step,
                success=True,
                output=json.dumps(sr.output, default=str) if sr.output else "{}",
            )
        return StepResult(
            step_index=index,
            step=step,
            success=False,
            output="",
            error=sr.error or "Newly built skill failed to execute",
        )

    async def _run_llm_only(
        self,
        index: int,
        step: PlanStep,
        provider: LLMProvider,
    ) -> StepResult:
        """Execute an LLM-only step with context from previous results."""
        context = self._build_context()
        prompt = step.prompt or step.description

        if context:
            full_prompt = (
                f"Previous step results:\n{context}\n\n"
                f"Now: {prompt}"
            )
        else:
            full_prompt = prompt

        try:
            response = await provider.complete(_LLM_ONLY_SYSTEM, full_prompt)
            return StepResult(
                step_index=index,
                step=step,
                success=True,
                output=response,
            )
        except Exception as e:
            return StepResult(
                step_index=index,
                step=step,
                success=False,
                output="",
                error=str(e),
            )

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

    def _build_context(self) -> str:
        """Format previous step results as context for LLM prompts."""
        if not self._step_results:
            return ""
        lines = []
        for r in self._step_results:
            if r.skipped:
                lines.append(f"Step {r.step_index + 1} ({r.step.description}): [skipped]")
            elif r.success:
                # Truncate long outputs
                out = r.output[:2000] if len(r.output) > 2000 else r.output
                lines.append(f"Step {r.step_index + 1} ({r.step.description}): {out}")
            else:
                lines.append(f"Step {r.step_index + 1} ({r.step.description}): [failed: {r.error}]")
        return "\n".join(lines)

    async def _synthesize(
        self,
        plan: ExecutionPlan,
        provider: LLMProvider,
    ) -> str:
        """Synthesize all step results into a final response."""
        successful = [r for r in self._step_results if r.success and not r.skipped]

        # Single step — no synthesis needed
        if len(successful) == 1:
            return successful[0].output

        # No successful steps
        if not successful:
            return "No steps completed successfully."

        context = self._build_context()
        user_prompt = (
            f'Original request: "{plan.original_request}"\n\n'
            f"Step results:\n{context}\n\n"
            f"Synthesize these results into a clear, helpful response."
        )

        try:
            return await provider.complete(_SYNTHESIZE_SYSTEM, user_prompt)
        except Exception as exc:
            import logging
            logging.getLogger("alive").debug("Synthesis LLM call failed: %s", exc)
            # Fallback: concatenate outputs
            return "\n\n".join(
                f"**{r.step.description}:**\n{r.output}" for r in successful
            )

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

            self._provider = get_provider()
        return self._provider


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


async def compose(
    plan: ExecutionPlan,
    *,
    provider: LLMProvider | None = None,
    registry: SkillRegistry | None = None,
    console: Console | None = None,
    on_error: Callable[[int, PlanStep, str, Console], ErrorAction] | None = None,
) -> CompositionResult:
    """Execute an ExecutionPlan and return the composed result.

    This is the main entry point for the composition engine.

    Args:
        plan: The execution plan to run.
        provider: LLM provider override. Defaults to ALIVE_PROVIDER.
        registry: SkillRegistry override. Defaults to global registry.
        console: Rich console for output.
        on_error: Error handler callback. Defaults to interactive CLI prompt.

    Returns:
        CompositionResult with step results and synthesized final output.

    Raises:
        CompositionError: If a step fails and the user chooses to abort.
    """
    composer = Composer(
        provider=provider,
        registry=registry,
        console=console,
        on_error=on_error,
    )
    return await composer.execute(plan)
