"""Tool runner implementing the four-stage tool lifecycle.

Stages: **discover** -> **activate** -> **execute** -> **deactivate**.

The runner is the thin orchestration layer that bridges the
[ToolRegistry][cleveragents.tool.registry.ToolRegistry] (where specs live)
and the actual handler callables.

## Lifecycle Stages

1. **Discover** -- ``discover(registry?)`` queries the registry for all
   available tools. An alternate registry may be passed for dynamic
   discovery.
2. **Activate** -- ``activate(tool_name)`` validates that the tool exists
   and marks it as ready for execution. Raises ``ToolError`` if not found.
3. **Execute** -- ``execute(tool_name, inputs)`` runs the tool handler.
   Guarantees JSON-serialisable I/O, error normalisation, and duration
   tracking.  When a ``CheckpointService`` is wired and the tool has
   ``writes=True``, automatic checkpoints are created before and after
   execution (``before_tool_execute`` / ``after_tool_execute`` triggers).
4. **Deactivate** -- ``deactivate(tool_name)`` cleans up. Returns ``True``
   if the tool was active, ``False`` otherwise.
"""

from __future__ import annotations

import json
import threading
import time
from typing import TYPE_CHECKING, Any

import structlog

from cleveragents.domain.models.core.execution_environment_preference import (
    EnvironmentPreferenceMode,
)
from cleveragents.domain.models.core.plan import ExecutionEnvironment
from cleveragents.tool.container_executor import ContainerToolExecutor
from cleveragents.tool.registry import ToolRegistry
from cleveragents.tool.runtime import ToolError, ToolResult, ToolSpec

logger = structlog.get_logger(__name__)

if TYPE_CHECKING:
    from cleveragents.application.services.checkpoint_service import CheckpointService
    from cleveragents.application.services.execution_environment_resolver import (
        ExecutionEnvironmentResolver,
    )

# Default set of active automatic checkpoint triggers.
# Exported as a public constant so callers (e.g. SubplanExecutionService,
# test step definitions) can reference the canonical set without duplicating it.
DEFAULT_AUTO_TRIGGERS: frozenset[str] = frozenset(
    {
        "before_tool_execute",
        "after_tool_execute",
        "on_subplan_spawn",
        "on_error",
    }
)

# Private alias kept for backward compatibility within this module.
_DEFAULT_AUTO_TRIGGERS = DEFAULT_AUTO_TRIGGERS


class ToolRunner:
    """Executes tools through the four-stage lifecycle.

    Parameters
    ----------
    registry:
        The ``ToolRegistry`` that this runner queries for tool specs.
    execution_environment_resolver:
        Optional resolver for execution environment routing.
    checkpoint_service:
        Optional ``CheckpointService`` for automatic per-tool-write
        checkpoint creation.  When provided and a tool has
        ``capabilities.writes=True``, checkpoints are created before
        and after execution (``before_tool_execute`` /
        ``after_tool_execute`` triggers).  When ``None``, checkpoint
        hooks are silently skipped.
    auto_checkpoint_triggers:
        Set of trigger names that are active.  Defaults to all four
        triggers.  Pass a subset to disable specific triggers.
    """

    def __init__(
        self,
        registry: ToolRegistry,
        execution_environment_resolver: ExecutionEnvironmentResolver | None = None,
        container_executor: ContainerToolExecutor | None = None,
        checkpoint_service: CheckpointService | None = None,
        auto_checkpoint_triggers: frozenset[str] | None = None,
    ) -> None:
        from cleveragents.application.services.execution_environment_resolver import (
            ExecutionEnvironmentResolver as _Resolver,
        )

        self._registry = registry
        self._active: dict[str, ToolSpec] = {}
        self._active_lock = threading.RLock()
        self._env_resolver: ExecutionEnvironmentResolver = (
            execution_environment_resolver or _Resolver()
        )
        self._container_executor = container_executor
        self._checkpoint_service = checkpoint_service
        self._auto_checkpoint_triggers: frozenset[str] = (
            auto_checkpoint_triggers
            if auto_checkpoint_triggers is not None
            else _DEFAULT_AUTO_TRIGGERS
        )

    # -- Checkpoint helpers ---------------------------------------------------

    def is_trigger_active(self, trigger: str) -> bool:
        """Return True if the given automatic checkpoint trigger is active.

        Public API so that callers such as ``PlanExecutor`` can query the
        active trigger set without accessing private attributes.

        Args:
            trigger: Trigger name (e.g. ``"on_error"``).

        Returns:
            ``True`` if the trigger is in the active set.
        """
        return trigger in self._auto_checkpoint_triggers

    # Private alias for internal callers within this module.
    _is_trigger_active = is_trigger_active

    def _try_create_tool_checkpoint(
        self,
        plan_id: str,
        tool_name: str,
        phase: str,
        sandbox_ref: str = "",
    ) -> None:
        """Create a checkpoint for a write-tool execution event.

        Silently skips when no ``CheckpointService`` is configured or
        when checkpoint creation fails (non-fatal).

        Args:
            plan_id: The plan identifier.
            tool_name: The tool being executed.
            phase: Trigger phase label (e.g. ``"before_tool_execute"``).
            sandbox_ref: Optional sandbox git ref for the checkpoint.
        """
        if self._checkpoint_service is None:
            return
        if not plan_id:
            return
        try:
            self._checkpoint_service.create_checkpoint(
                plan_id=plan_id,
                sandbox_ref=sandbox_ref or f"tool-{phase}-{tool_name}",
                reason=f"Automatic checkpoint: {phase}",
                source_tool=tool_name,
                phase="execute",
                checkpoint_type="pre_write" if "before" in phase else "post_step",
            )
        except Exception:
            logger.debug(
                "tool_checkpoint_creation_failed",
                plan_id=plan_id,
                tool_name=tool_name,
                phase=phase,
                exc_info=True,
            )

    # -- Stage 1: Discover ----------------------------------------------------

    def discover(self, registry: ToolRegistry | None = None) -> list[ToolSpec]:
        """Find available tools from the registry.

        If *registry* is provided it overrides the instance-level
        registry for this call (useful for dynamic discovery).
        """
        source = registry if registry is not None else self._registry
        return source.list_tools()

    # -- Stage 2: Activate ----------------------------------------------------

    def activate(self, tool_name: str) -> ToolSpec:
        """Prepare a tool for execution.

        Validates that the tool exists in the registry and marks it as
        active.  Raises ``ToolError`` if the tool is not found.
        """
        spec = self._registry.get(tool_name)
        if spec is None:
            raise ToolError(
                tool_name=tool_name,
                error_type="ActivationError",
                details=f"Tool '{tool_name}' not found in registry",
            )
        with self._active_lock:
            self._active[tool_name] = spec
        return spec

    # -- Stage 3: Execute -----------------------------------------------------

    def resolve_execution_environment(
        self,
        tool_env: str | None = None,
        plan_env: str | None = None,
        project_env: str | None = None,
    ) -> ExecutionEnvironment:
        """Resolve the execution environment for the current invocation.

        Delegates to :class:`ExecutionEnvironmentResolver`.
        """
        return self._env_resolver.resolve(
            tool_env=tool_env,
            plan_env=plan_env,
            project_env=project_env,
        )

    def execute(
        self,
        tool_name: str,
        inputs: dict[str, Any],
        *,
        tool_env: str | None = None,
        plan_env: str | None = None,
        plan_priority: str | None = None,
        project_env: str | None = None,
        project_priority: str | None = None,
        linked_resource_types: list[str] | None = None,
        project_name: str = "",
        timeout_seconds: int | None = None,
        plan_id: str = "",
        sandbox_ref: str = "",
    ) -> ToolResult:
        """Run the tool with strict JSON-serialisable IO.

        When *tool_env*, *plan_env*, or *project_env* indicate
        ``container`` execution, the runner validates that a container
        resource is linked before delegating to the container handler.

        Any exception raised by the handler is caught and normalised
        into a ``ToolResult`` with ``success=False``.

        When a ``CheckpointService`` is wired and the tool has
        ``capabilities.writes=True``, automatic checkpoints are created
        before and after execution (``before_tool_execute`` /
        ``after_tool_execute`` triggers), provided those triggers are
        active in ``auto_checkpoint_triggers``.

        Parameters
        ----------
        plan_priority:
            Priority semantics for *plan_env*: ``"override"`` (always
            wins over auto-detected devcontainers) or ``"fallback"``
            (defers to auto-detected devcontainers when present).
            Defaults to ``"fallback"`` when not supplied.
        project_priority:
            Priority semantics for *project_env*: ``"override"`` or
            ``"fallback"``.  Defaults to ``"fallback"``.
        timeout_seconds:
            Optional timeout override forwarded to the container
            executor when the execution environment is ``CONTAINER``.
        plan_id:
            Optional plan identifier for automatic checkpoint creation.
            When empty, checkpoint hooks are skipped.
        sandbox_ref:
            Optional sandbox git ref passed to checkpoint creation.
        """
        with self._active_lock:
            spec = self._active.get(tool_name)
        if spec is None:
            spec = self._registry.get(tool_name)
        if spec is None:
            raise ToolError(
                tool_name=tool_name,
                error_type="ExecutionError",
                details=f"Tool '{tool_name}' not found",
            )

        from cleveragents.application.services.execution_environment_resolver import (
            ContainerUnavailableError,
        )

        # Consult tool-level execution environment preference.
        # Caller override takes precedence over tool preference —
        # PREFERRED only upgrades when no caller override is set.
        pref = spec.execution_environment
        effective_tool_env = tool_env

        if pref.mode == EnvironmentPreferenceMode.SPECIFIC:
            # Route to container; the named target_resource identifies
            # which container to use.
            effective_tool_env = "container"
            # TODO: pass pref.target_resource to the resolver /
            # container executor so the specific named resource is
            # selected.  The current resolver and ContainerToolExecutor
            # APIs do not accept a target_resource parameter.
            logger.debug(
                "tool_env_preference.specific_override",
                tool=tool_name,
                target_resource=pref.target_resource,
            )
        elif pref.mode == EnvironmentPreferenceMode.REQUIRED:
            # Force container execution
            effective_tool_env = "container"
            logger.info(
                "tool_env_preference.required_container",
                tool=tool_name,
            )
        elif (
            pref.mode == EnvironmentPreferenceMode.PREFERRED
            and effective_tool_env is None
        ):
            # Try container if available, otherwise fall back
            effective_tool_env = "container"
            logger.debug(
                "tool_env_preference.preferred_upgrade",
                tool=tool_name,
            )

        # Resolve execution environment using the 6-level precedence chain.
        # devcontainer_available is True when any linked resource is a
        # devcontainer-instance (Level 3 in the precedence chain).
        _linked = linked_resource_types or []
        devcontainer_available = self._env_resolver.has_devcontainer(_linked)
        try:
            env = self._env_resolver.resolve_with_precedence(
                tool_env=effective_tool_env,
                plan_env=plan_env,
                plan_priority=plan_priority,
                project_env=project_env,
                project_priority=project_priority,
                devcontainer_available=devcontainer_available,
            )
            # Validate container availability when container execution is
            # resolved (mirrors the old resolve_and_validate contract).
            if env.value in ("container", "container_ref"):
                self._env_resolver.validate_container_available(_linked, project_name)
        except ContainerUnavailableError as exc:
            if pref.mode == EnvironmentPreferenceMode.REQUIRED:
                logger.info(
                    "tool_env_preference.required_failed",
                    tool=tool_name,
                    error=str(exc),
                )
                return ToolResult(
                    success=False,
                    output={},
                    error=(
                        f"Tool '{tool_name}' requires container execution "
                        f"but no container is available: {exc}"
                    ),
                    duration_ms=0.0,
                )
            if pref.mode == EnvironmentPreferenceMode.PREFERRED:
                # Fall back to host
                logger.info(
                    "tool_env_preference.preferred_fallback_to_host",
                    tool=tool_name,
                    reason=str(exc),
                )
                env = ExecutionEnvironment.HOST
            elif pref.mode == EnvironmentPreferenceMode.SPECIFIC:
                logger.info(
                    "tool_env_preference.specific_target_unavailable",
                    tool=tool_name,
                    target_resource=pref.target_resource,
                    error=str(exc),
                )
                return ToolResult(
                    success=False,
                    output={},
                    error=(
                        f"Target container resource "
                        f"'{pref.target_resource}' unavailable: {exc}"
                    ),
                    duration_ms=0.0,
                )
            else:
                return ToolResult(
                    success=False,
                    output={},
                    error=str(exc),
                    duration_ms=0.0,
                )
        except Exception as exc:
            return ToolResult(
                success=False,
                output={},
                error=f"Execution environment error: {exc}",
                duration_ms=0.0,
            )

        # Container routing — delegate to container handler
        if env == ExecutionEnvironment.CONTAINER:
            if self._container_executor is None:
                return ToolResult(
                    success=False,
                    output={},
                    error=(
                        "Container execution is not available. "
                        "A ContainerToolExecutor must be configured on the "
                        "ToolRunner before container-routed tools can execute."
                    ),
                    duration_ms=0.0,
                )

            # Validate input_schema and capabilities (T2) before
            # delegating to the container executor so container-routed
            # tools receive the same spec validation as host-routed ones.
            try:
                json.dumps(inputs, allow_nan=False)
            except (TypeError, ValueError) as exc:
                return ToolResult(
                    success=False,
                    output={},
                    error=f"Inputs are not JSON-serialisable: {exc}",
                    duration_ms=0.0,
                )

            if spec.input_schema:
                # Iterate the ``required`` list directly so that fields
                # listed in ``required`` but absent from ``properties``
                # (e.g. ``additionalProperties`` patterns) are still
                # detected (P2-20).
                missing = [
                    k for k in spec.input_schema.get("required", []) if k not in inputs
                ]
                if missing:
                    return ToolResult(
                        success=False,
                        output={},
                        error=(
                            f"Missing required input fields for "
                            f"'{tool_name}': {missing}"
                        ),
                        duration_ms=0.0,
                    )

            try:
                return self._container_executor.execute_tool(
                    tool_name,
                    inputs,
                    timeout_seconds=timeout_seconds,
                )
            except Exception as exc:
                # Intentionally broad: the runner's contract is to
                # normalise *any* handler failure into a ToolResult
                # with ``success=False`` so callers never see raw
                # exceptions from container execution.
                return ToolResult(
                    success=False,
                    output={},
                    error=f"Container execution error: {type(exc).__name__}: {exc}",
                    duration_ms=0.0,
                )

        # Validate that inputs are JSON-serialisable.  Host-routed tools
        # use the default ``allow_nan=True`` for backward compatibility;
        # only the container path enforces RFC 7159 (P1-1).
        try:
            json.dumps(inputs)
        except (TypeError, ValueError) as exc:
            return ToolResult(
                success=False,
                output={},
                error=f"Inputs are not JSON-serialisable: {exc}",
                duration_ms=0.0,
            )

        # Determine if this tool writes to the sandbox.
        _tool_writes = spec.capabilities.writes

        # before_tool_execute checkpoint: fire before any write tool runs.
        if _tool_writes and self._is_trigger_active("before_tool_execute"):
            self._try_create_tool_checkpoint(
                plan_id=plan_id,
                tool_name=tool_name,
                phase="before_tool_execute",
                sandbox_ref=sandbox_ref,
            )

        start = time.monotonic()
        try:
            raw_output = spec.handler(inputs)
        except Exception as exc:
            # Intentionally broad: the runner normalises handler
            # exceptions into a failed ToolResult so callers get a
            # uniform result type regardless of handler implementation.
            elapsed = (time.monotonic() - start) * 1000.0
            return ToolResult(
                success=False,
                output={},
                error=f"{type(exc).__name__}: {exc}",
                duration_ms=elapsed,
            )

        elapsed = (time.monotonic() - start) * 1000.0

        # Normalise output to dict
        if not isinstance(raw_output, dict):
            raw_output = {"result": raw_output}

        # Validate that output is JSON-serialisable.  Host-routed tools
        # use the default ``allow_nan=True`` for backward compatibility.
        try:
            json.dumps(raw_output)
        except (TypeError, ValueError) as exc:
            return ToolResult(
                success=False,
                output={},
                error=f"Output is not JSON-serialisable: {exc}",
                duration_ms=elapsed,
            )

        # after_tool_execute checkpoint: fire after a write tool succeeds.
        if _tool_writes and self._is_trigger_active("after_tool_execute"):
            self._try_create_tool_checkpoint(
                plan_id=plan_id,
                tool_name=tool_name,
                phase="after_tool_execute",
                sandbox_ref=sandbox_ref,
            )

        return ToolResult(
            success=True,
            output=raw_output,
            duration_ms=elapsed,
        )

    # -- Stage 4: Deactivate --------------------------------------------------

    def deactivate(self, tool_name: str) -> bool:
        """Clean up after tool execution.

        Returns ``True`` if the tool was active and has been
        deactivated, ``False`` otherwise.
        """
        with self._active_lock:
            if tool_name in self._active:
                del self._active[tool_name]
                return True
            return False
