"""Executor — executes plans and steps using registered tools.

The Executor is the runtime that takes a Plan and runs each Step
through the registered tool dispatch. It handles:
- Tool lookup and dispatch
- Parameter validation before execution
- Timeout enforcement
- Automatic retries for transient failures
- Fallback tool switching
- Budget tracking
- Step tracing for debugging
"""

import asyncio
import time
from typing import Any

from xyglue.plan.models import Plan, Step
from xyglue.execute.models import (
    ExecutionResult,
    ExecutionContext,
    StepTrace,
)
from xyglue.tools.base import Tool


class Executor:
    """Executes plans step by step with tool dispatch.

    Tools are registered by name. When a step references a tool that
    isn't registered, the executor simulates success (rule-based mode)
    so the pipeline works without real tools attached.
    """

    def __init__(self, default_timeout_ms: int = 30_000) -> None:
        self._tools: dict[str, Tool] = {}
        self._default_timeout_ms = default_timeout_ms
        self._execution_log: list[StepTrace] = []

    # ── Tool management ───────────────────────────────────────────────

    def register_tool(self, tool: Tool) -> None:
        self._tools[tool.name] = tool

    def unregister_tool(self, name: str) -> bool:
        return self._tools.pop(name, None) is not None

    def get_tool(self, name: str) -> Tool | None:
        return self._tools.get(name)

    def list_tools(self) -> list[str]:
        return sorted(self._tools.keys())

    # ── Plan execution ────────────────────────────────────────────────

    async def execute_plan(
        self, plan: Plan, ctx: ExecutionContext
    ) -> list[ExecutionResult]:
        """Execute all steps in a plan sequentially."""
        results: list[ExecutionResult] = []
        for step in plan.steps:
            step_ctx = ExecutionContext(
                run_id=ctx.run_id,
                step_number=step.number,
                previous_results=results,
                accumulated_state=ctx.accumulated_state,
                budget_used_usd=ctx.budget_used_usd,
                budget_limit_usd=ctx.budget_limit_usd,
                start_time_ms=ctx.start_time_ms,
                timeout_ms=ctx.timeout_ms,
            )

            if step_ctx.is_over_budget:
                results.append(ExecutionResult(
                    step_number=step.number,
                    success=False,
                    error="Budget exceeded",
                    tool_used=step.tool,
                ))
                break

            result = await self.execute_step(step, step_ctx)
            results.append(result)

            if result.success and result.output is not None:
                ctx.accumulated_state[f"step_{step.number}_output"] = result.output

            if result.failed:
                break
        return results

    # ── Step execution ────────────────────────────────────────────────

    async def execute_step(
        self, step: Step, ctx: ExecutionContext
    ) -> ExecutionResult:
        """Execute a single step: validate, run tool, handle errors."""
        start = time.time()
        trace = StepTrace(
            step_number=step.number,
            tool=step.tool,
            action=step.action,
            params_summary=_summarize_params(step.params),
        )

        tool = self._tools.get(step.tool)

        if tool is None:
            result = ExecutionResult(
                step_number=step.number,
                success=True,
                output=f"Simulated: {step.description}",
                tool_used=step.tool,
                duration_ms=int((time.time() - start) * 1000),
                trace=trace,
            )
            trace.duration_ms = result.duration_ms
            self._execution_log.append(trace)
            return result

        # Validate parameters
        try:
            issues = await tool.validate(step.params)
            if issues:
                trace.success = False
                trace.error = f"Validation failed: {'; '.join(issues)}"
                trace.duration_ms = int((time.time() - start) * 1000)
                self._execution_log.append(trace)
                return ExecutionResult(
                    step_number=step.number,
                    success=False,
                    error=trace.error,
                    tool_used=step.tool,
                    duration_ms=trace.duration_ms,
                    trace=trace,
                )
        except Exception as exc:
            trace.success = False
            trace.error = f"Validation exception: {exc}"
            trace.duration_ms = int((time.time() - start) * 1000)
            self._execution_log.append(trace)
            return ExecutionResult(
                step_number=step.number,
                success=False,
                error=str(exc),
                tool_used=step.tool,
                duration_ms=int((time.time() - start) * 1000),
                trace=trace,
            )

        # Execute with timeout
        timeout_ms = step.timeout_ms or self._default_timeout_ms
        result = await self._execute_with_timeout(tool, step, timeout_ms)

        # If failed and step has auto-retry (idempotent + retryable error)
        if result.failed and result.retryable and step.idempotent and step.can_retry:
            step.record_retry()
            trace.was_retried = True
            retry_result = await self._execute_with_timeout(tool, step, timeout_ms)
            if retry_result.success:
                retry_result.was_retried = True
                result = retry_result

        # If still failed and step has a fallback tool
        if result.failed and step.has_fallback:
            fallback_result = await self.try_fallbacks(step, ctx)
            if fallback_result and fallback_result.success:
                fallback_result.used_fallback = True
                trace.used_fallback = True
                result = fallback_result

        # Finalize trace
        result.duration_ms = int((time.time() - start) * 1000)
        result.trace = trace
        trace.duration_ms = result.duration_ms
        trace.success = result.success
        trace.error = result.error
        self._execution_log.append(trace)
        return result

    async def _execute_with_timeout(
        self, tool: Tool, step: Step, timeout_ms: int
    ) -> ExecutionResult:
        """Run a tool with timeout enforcement."""
        start = time.time()
        timeout_s = timeout_ms / 1000.0
        try:
            tool_result = await asyncio.wait_for(
                tool.execute(step.params), timeout=timeout_s
            )
            return ExecutionResult(
                step_number=step.number,
                success=tool_result.success,
                output=tool_result.output,
                error=tool_result.error,
                tool_used=step.tool,
                retryable=getattr(tool_result, "retryable", False),
                duration_ms=int((time.time() - start) * 1000),
                metadata=tool_result.metadata,
            )
        except asyncio.TimeoutError:
            return ExecutionResult(
                step_number=step.number,
                success=False,
                error=f"Tool '{step.tool}' timed out after {timeout_ms}ms",
                tool_used=step.tool,
                retryable=True,
                duration_ms=int((time.time() - start) * 1000),
            )
        except Exception as exc:
            return ExecutionResult(
                step_number=step.number,
                success=False,
                error=str(exc),
                tool_used=step.tool,
                duration_ms=int((time.time() - start) * 1000),
            )

    # ── Fallback execution ────────────────────────────────────────────

    async def try_fallbacks(
        self, step: Step, ctx: ExecutionContext
    ) -> ExecutionResult | None:
        """Try executing with the step's fallback tool."""
        if not step.fallback_tool:
            return None
        fallback_step = Step(
            number=step.number,
            action=step.action,
            tool=step.fallback_tool,
            params=step.fallback_params or step.params,
            description=f"Fallback: {step.description}",
            timeout_ms=step.timeout_ms,
        )
        return await self.execute_step(fallback_step, ctx)

    # ── Execution log ─────────────────────────────────────────────────

    def get_execution_log(self) -> list[StepTrace]:
        return list(self._execution_log)

    def clear_log(self) -> None:
        self._execution_log.clear()

    def get_stats(self) -> dict[str, Any]:
        """Summary statistics for all executed steps."""
        if not self._execution_log:
            return {"total_steps": 0}
        total = len(self._execution_log)
        successes = sum(1 for t in self._execution_log if t.success)
        retried = sum(1 for t in self._execution_log if t.was_retried)
        fallbacks = sum(1 for t in self._execution_log if t.used_fallback)
        avg_duration = (
            sum(t.duration_ms for t in self._execution_log) / total
        )
        return {
            "total_steps": total,
            "successes": successes,
            "failures": total - successes,
            "retried": retried,
            "fallbacks_used": fallbacks,
            "avg_duration_ms": round(avg_duration, 1),
        }


def _summarize_params(params: dict[str, Any], max_len: int = 100) -> str:
    """Create a short summary of step params for tracing."""
    if not params:
        return "{}"
    parts: list[str] = []
    for key, value in params.items():
        val_str = str(value)
        if len(val_str) > 30:
            val_str = val_str[:27] + "..."
        parts.append(f"{key}={val_str}")
    result = ", ".join(parts)
    if len(result) > max_len:
        result = result[: max_len - 3] + "..."
    return result
