"""Script Tool Provider - Execute bundled skill scripts as tools.

This module provides the infrastructure for executing scripts bundled with skills
as callable tools. When a skill loads with a scripts/ directory or explicit tool
declarations, ScriptToolProvider registers handlers that wrap script execution.

Architecture:
    Skill (with scripts/)
       └──► ScriptToolProvider.register_skill_tools()
              └──► ToolRegistry.register_builtin()
                     └──► Handler wraps subprocess execution

Security:
- Scripts can run in local mode (with resource limits) or Docker mode (isolated)
- Local mode applies POSIX resource limits (memory, CPU, file size) when RuntimeContext is provided
- Docker mode provides full container isolation with security hardening
- Environment variables like OPENAI_API_KEY are inherited
- Working directory is set to the skill's source directory
- Only allowed script extensions are executed

Usage:
    from infrastructure.tools.script_provider import ScriptToolProvider

    provider = ScriptToolProvider()
    provider.register_skill_tools(skill, tool_registry)

    # Now tools like "sora.create" can be invoked via the registry
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
import os
import resource
import shlex
import shutil
import signal
import time
from typing import TYPE_CHECKING, Any

from pydantic import BaseModel, Field, computed_field

from core.pathways.models.tool import ToolOutput
from core.skills.loader import SCRIPT_INTERPRETERS
from core.skills.models import RuntimeContext

if TYPE_CHECKING:
    from pathlib import Path

    from core.pathways.context import Context
    from core.skills.models import ScriptToolDeclaration, Skill
    from core.tools import ToolRegistry

logger = logging.getLogger(__name__)


# =============================================================================
# Errors
# =============================================================================


class ScriptExecutionError(Exception):
    """Raised when a script fails to execute."""

    def __init__(
        self,
        script: str,
        exit_code: int,
        stderr: str,
        stdout: str = "",
    ):
        self.script = script
        self.exit_code = exit_code
        self.stderr = stderr
        self.stdout = stdout
        super().__init__(f"Script '{script}' failed with exit code {exit_code}: {stderr[:500]}")


class ScriptNotFoundError(Exception):
    """Raised when a script file is not found."""

    def __init__(self, script_path: str):
        self.script_path = script_path
        super().__init__(f"Script not found: {script_path}")


class InterpreterNotFoundError(Exception):
    """Raised when the interpreter for a script is not found."""

    def __init__(self, interpreter: str, script: str):
        self.interpreter = interpreter
        self.script = script
        super().__init__(f"Interpreter '{interpreter}' not found for script '{script}'")


# =============================================================================
# Script Execution Result
# =============================================================================


class ScriptResult(BaseModel):
    """Result of a script execution."""

    exit_code: int
    stdout: str
    stderr: str
    duration_ms: float
    script_path: str
    args: list[str] = Field(default_factory=list)

    @computed_field
    @property
    def success(self) -> bool:
        """Check if the script executed successfully."""
        return self.exit_code == 0

    def to_tool_output(self) -> ToolOutput:
        """Convert to canonical ToolOutput for tool response."""
        data: dict[str, Any] = {
            "exit_code": self.exit_code,
            "stdout": self.stdout,
            "stderr": self.stderr,
            "duration_ms": self.duration_ms,
            "script_path": self.script_path,
            "args": list(self.args),
        }
        return ToolOutput(success=self.success, data=data)


# =============================================================================
# Script Tool Provider
# =============================================================================


class ScriptToolProvider:
    """Provides tools backed by skill scripts.

    Registers script-backed tools from skills into a ToolRegistry.
    Handles subprocess execution, argument building, and result formatting.

    Example:
        skill = await load_skill_from_directory(Path("external-skills/sora"))
        provider = ScriptToolProvider()
        registered = provider.register_skill_tools(skill, registry)
        print(f"Registered {registered} tools from {skill.id}")
    """

    def __init__(
        self,
        *,
        timeout_seconds: float = 300.0,  # 5 minute default
        inherit_env: bool = True,
        allowed_extensions: set[str] | None = None,
    ):
        """Initialize the script tool provider.

        Args:
            timeout_seconds: Maximum execution time for scripts
            inherit_env: Whether to inherit environment variables
            allowed_extensions: Allowed script extensions (default: all in SCRIPT_INTERPRETERS)
        """
        self._timeout = timeout_seconds
        self._inherit_env = inherit_env
        self._allowed_extensions = allowed_extensions or set(SCRIPT_INTERPRETERS.keys())
        self._registered_tools: dict[str, tuple[Skill, ScriptToolDeclaration]] = {}

    def register_skill_tools(
        self,
        skill: Skill,
        registry: ToolRegistry,
    ) -> int:
        """Register all script tools from a skill into the registry.

        Args:
            skill: Skill with script_tools to register
            registry: Tool registry to register handlers into

        Returns
        -------
            Number of tools registered
        """
        if not skill.script_tools:
            return 0

        registered = 0
        for tool_decl in skill.script_tools:
            tool_id = f"{skill.id}.{tool_decl.id}"

            try:
                self._register_single_tool(skill, tool_decl, tool_id, registry)
                self._registered_tools[tool_id] = (skill, tool_decl)
                registered += 1
                logger.info(f"Registered script tool: {tool_id}")
            except Exception as e:
                logger.exception(f"Failed to register tool {tool_id}: {e}")

        return registered

    def _register_single_tool(
        self,
        skill: Skill,
        tool_decl: ScriptToolDeclaration,
        tool_id: str,
        registry: ToolRegistry,
    ) -> None:
        """Register a single script tool."""
        # Validate script exists
        if not skill.source_dir:
            msg = f"Skill {skill.id} has no source_dir set"
            raise ValueError(msg)

        script_path = skill.source_dir / tool_decl.script
        if not script_path.exists():
            raise ScriptNotFoundError(str(script_path))

        # Build input schema from params
        input_schema = self.build_input_schema(tool_decl)

        # Create the handler that wraps script execution
        async def handler(
            arguments: dict[str, Any],
            context: Context,
        ) -> ToolOutput:
            return await self.execute_script_tool(skill, tool_decl, arguments, context)

        # Register with the registry using ToolDefinition
        from core.tools import ToolDefinition

        registry.register(
            ToolDefinition(
                name=tool_id,
                description=tool_decl.description or f"Run {tool_decl.script}",
                parameters=input_schema,
                handler=handler,
                category="script",
            )
        )

    @staticmethod
    def _param_to_cli_args(param_name: str, param_type: str, value: Any) -> list[str]:
        """Convert a typed parameter to CLI flag(s).

        Dispatch table replaces if/elif on param_type.
        """
        _CONVERTERS: dict[str, Any] = {
            "boolean": lambda n, v: [f"--{n}"] if v else [],
            "array": lambda n, v: [a for item in v for a in (f"--{n}", str(item))],
        }
        converter = _CONVERTERS.get(param_type, lambda n, v: [f"--{n}", str(v)])
        return converter(param_name, value)

    def build_input_schema(
        self,
        tool_decl: ScriptToolDeclaration,
    ) -> dict[str, Any]:
        """Build JSON Schema for tool inputs from param declarations."""
        properties: dict[str, Any] = {}
        required: list[str] = []

        # Add declared params
        for param in tool_decl.params:
            prop: dict[str, Any] = {
                "type": self._map_param_type(param.param_type),
            }
            if param.description:
                prop["description"] = param.description
            if param.default is not None:
                prop["default"] = param.default

            properties[param.name] = prop
            if param.required:
                required.append(param.name)

        # Always allow arbitrary args to be passed through
        properties["_extra_args"] = {
            "type": "array",
            "items": {"type": "string"},
            "description": "Additional command-line arguments to pass to the script",
        }

        schema: dict[str, Any] = {
            "type": "object",
            "properties": properties,
        }
        if required:
            schema["required"] = required

        return schema

    def _map_param_type(self, param_type: str) -> str:
        """Map our param types to JSON Schema types."""
        mapping = {
            "string": "string",
            "number": "number",
            "integer": "integer",
            "boolean": "boolean",
            "array": "array",
            "object": "object",
        }
        return mapping.get(param_type, "string")

    async def execute_script_tool(
        self,
        skill: Skill,
        tool_decl: ScriptToolDeclaration,
        arguments: dict[str, Any],
        context: Context,
    ) -> ToolOutput:
        """Execute a script and return the result.

        Builds command line from arguments and runs via subprocess or Docker.
        """
        if not skill.source_dir:
            return ToolOutput(
                success=False,
                error=f"Skill {skill.id} has no source_dir",
                metadata={"skill_id": skill.id},
            )

        script_path = skill.source_dir / tool_decl.script
        if not script_path.exists():
            return ToolOutput(
                success=False,
                error=f"Script not found: {script_path}",
                metadata={"script": str(script_path)},
            )

        # Check if we should use Docker runtime
        if tool_decl.runtime and tool_decl.runtime.mode == "docker":
            return await self._execute_in_docker(skill, tool_decl, script_path, arguments, context)

        # Local execution path with resource limits (if RuntimeContext provided)
        # Determine interpreter
        interpreter = tool_decl.interpreter or self._detect_interpreter(script_path)
        if not interpreter:
            return ToolOutput(
                success=False,
                error=f"No interpreter for script: {script_path.suffix}",
                metadata={"script": str(script_path)},
            )

        # Verify interpreter exists
        interpreter_cmd = interpreter.split()[0]
        if not shutil.which(interpreter_cmd):
            return ToolOutput(
                success=False,
                error=f"Interpreter not found: {interpreter_cmd}",
                metadata={"interpreter": interpreter_cmd, "script": str(script_path)},
            )

        # Build command
        cmd = self._build_command(
            interpreter=interpreter,
            script_path=script_path,
            tool_decl=tool_decl,
            arguments=arguments,
        )

        # Build environment
        env = self._build_environment(tool_decl)

        # Determine working directory
        cwd = skill.source_dir
        if tool_decl.working_dir:
            cwd = skill.source_dir / tool_decl.working_dir

        # Determine timeout from runtime config if available
        timeout = self._timeout
        if tool_decl.runtime:
            timeout = float(tool_decl.runtime.timeout_s)

        # Execute
        try:
            result = await self._run_subprocess(
                cmd,
                cwd=cwd,
                env=env,
                runtime=tool_decl.runtime,
            )
            return result.to_tool_output()
        except TimeoutError:
            return ToolOutput(
                success=False,
                error=f"Script timed out after {timeout}s",
                metadata={"script": str(script_path), "timeout_s": timeout},
            )
        except Exception as e:
            logger.error(f"Script execution failed: {e}", exc_info=True)
            return ToolOutput(
                success=False,
                error=str(e),
                metadata={"script": str(script_path)},
            )

    async def _execute_in_docker(
        self,
        skill: Skill,
        tool_decl: ScriptToolDeclaration,
        script_path: Path,
        arguments: dict[str, Any],
        context: Context,
    ) -> ToolOutput:
        """Execute a script inside a Docker container.

        Mounts the skill directory into the container and runs the script.
        """
        runtime = tool_decl.runtime
        if not runtime or not runtime.image:
            return ToolOutput(
                success=False,
                error="Docker runtime requires an image",
                metadata={"runtime": "docker"},
            )

        # Check Docker availability
        if not shutil.which("docker"):
            return ToolOutput(
                success=False,
                error="Docker not found. Install Docker to use containerized tools.",
                metadata={"runtime": "docker"},
            )

        start = time.monotonic()

        # Build the command to run inside the container
        # Script path inside container (mounted at /skill)
        container_script = f"/skill/{tool_decl.script}"

        # Determine interpreter
        interpreter = tool_decl.interpreter or self._detect_interpreter(script_path)
        if not interpreter:
            interpreter = "python3"  # Default for Docker

        # Build script arguments
        script_args = list(tool_decl.default_args)
        extra_args = arguments.pop("_extra_args", [])

        for param in tool_decl.params:
            value = arguments.get(param.name)
            if value is None:
                continue
            script_args.extend(self._param_to_cli_args(param.name, param.param_type, value))

        if extra_args:
            script_args.extend(extra_args)

        # Build Docker command
        docker_cmd = [
            "docker",
            "run",
            "--rm",
            "-v",
            f"{skill.source_dir}:/skill:ro",
            "-w",
            "/skill",
            "--memory",
            f"{runtime.memory_mb}m",
        ]

        # Add GPU support if requested
        if runtime.gpu:
            docker_cmd.extend(["--gpus", "all"])

        # Add network isolation for security (can be made configurable)
        # docker_cmd.append("--network=none")

        # Add environment variables
        for key, value in tool_decl.env.items():
            docker_cmd.extend(["-e", f"{key}={value}"])

        # Add image
        docker_cmd.append(runtime.image)

        # Build the command to run inside container
        # If packages are specified, install them first
        if runtime.packages:
            pip_install = f"pip install -q {' '.join(runtime.packages)}"
            # Properly quote arguments for shell
            quoted_args = " ".join(shlex.quote(arg) for arg in script_args)
            run_script = f"{' '.join(interpreter.split())} {container_script} {quoted_args}"
            docker_cmd.extend(["sh", "-c", f"{pip_install} && {run_script}"])
        else:
            # Add the actual command directly
            docker_cmd.extend(interpreter.split())
            docker_cmd.append(container_script)
            docker_cmd.extend(script_args)

        logger.info(f"Executing in Docker: {' '.join(docker_cmd[:10])}...")

        try:
            process = await asyncio.create_subprocess_exec(
                *docker_cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )

            timeout = runtime.timeout_s if runtime else self._timeout
            stdout, stderr = await asyncio.wait_for(
                process.communicate(),
                timeout=timeout,
            )

            duration_ms = (time.monotonic() - start) * 1000

            ok = process.returncode == 0
            data: dict[str, Any] = {
                "exit_code": process.returncode or 0,
                "stdout": stdout.decode("utf-8", errors="replace"),
                "stderr": stderr.decode("utf-8", errors="replace") if stderr else "",
                "duration_ms": duration_ms,
                "runtime": "docker",
                "image": runtime.image,
                "script": str(script_path),
            }

            if not ok:
                logger.warning(
                    f"Docker script failed with exit code {data['exit_code']}: "
                    f"{str(data.get('stderr', ''))[:200]}"
                )

            return ToolOutput(success=ok, data=data)

        except TimeoutError:
            return ToolOutput(
                success=False,
                error=f"Docker execution timed out after {runtime.timeout_s}s",
                metadata={"script": str(script_path), "runtime": "docker"},
            )
        except Exception as e:
            logger.error(f"Docker execution failed: {e}", exc_info=True)
            return ToolOutput(
                success=False,
                error=str(e),
                metadata={"script": str(script_path), "runtime": "docker"},
            )

    def _build_command(
        self,
        interpreter: str,
        script_path: Path,
        tool_decl: ScriptToolDeclaration,
        arguments: dict[str, Any],
    ) -> list[str]:
        """Build the command line for script execution."""
        # Start with interpreter (may have args like "npx ts-node")
        cmd = interpreter.split()

        # Add script path
        cmd.append(str(script_path))

        # Add default args (e.g., ["create"] for subcommands)
        cmd.extend(tool_decl.default_args)

        # Add arguments from tool invocation
        extra_args = arguments.pop("_extra_args", [])

        for param in tool_decl.params:
            value = arguments.get(param.name)
            if value is None:
                continue
            cmd.extend(self._param_to_cli_args(param.name, param.param_type, value))

        # Add any extra args
        if extra_args:
            cmd.extend(extra_args)

        return cmd

    def _build_environment(
        self,
        tool_decl: ScriptToolDeclaration,
    ) -> dict[str, str]:
        """Build environment for script execution."""
        env: dict[str, str] = {}

        if self._inherit_env:
            env.update(os.environ)

        # Add tool-specific env vars
        env.update(tool_decl.env)

        # Ensure uv cache is writable (common issue in sandboxes)
        if "UV_CACHE_DIR" not in env:
            env["UV_CACHE_DIR"] = "/tmp/uv-cache"

        return env

    def _detect_interpreter(self, script_path: Path) -> str | None:
        """Detect interpreter from script extension."""
        ext = script_path.suffix.lower()
        if ext not in self._allowed_extensions:
            return None
        return SCRIPT_INTERPRETERS.get(ext)

    async def _run_subprocess(
        self,
        cmd: list[str],
        *,
        cwd: Path,
        env: dict[str, str],
        runtime: RuntimeContext | None = None,
    ) -> ScriptResult:
        """Run a subprocess and capture output.

        Applies resource limits from RuntimeContext if provided (for local execution).
        """
        start = time.monotonic()

        # Determine timeout from runtime config or use default
        timeout = self._timeout
        if runtime:
            timeout = float(runtime.timeout_s)

        logger.debug(f"Executing: {' '.join(cmd)} in {cwd} (timeout={timeout}s)")

        # Build preexec_fn for resource limits (POSIX only)
        def _apply_resource_limits() -> None:
            """Apply resource limits from RuntimeContext if available."""
            if not runtime or os.name != "posix":
                return

            try:
                # Set process group for better cleanup
                with contextlib.suppress(Exception):
                    os.setsid()

                # CPU time limit (seconds)
                cpu_seconds = runtime.timeout_s
                resource.setrlimit(
                    resource.RLIMIT_CPU, (cpu_seconds, cpu_seconds + 1)
                )

                # Memory limit (bytes)
                memory_bytes = int(runtime.memory_mb) * 1024 * 1024
                resource.setrlimit(
                    resource.RLIMIT_AS, (memory_bytes, memory_bytes)
                )

                # File size limit (5MB default)
                resource.setrlimit(
                    resource.RLIMIT_FSIZE, (5 * 1024 * 1024, 5 * 1024 * 1024)
                )

                # Open files limit
                resource.setrlimit(resource.RLIMIT_NOFILE, (64, 64))

            except Exception:
                # Non-fatal - we still have wall-clock timeout and process isolation
                logger.debug("Failed to set resource limits", exc_info=True)

        # Use preexec_fn only on POSIX systems and when runtime config is provided
        preexec_fn = _apply_resource_limits if (runtime and os.name == "posix") else None

        process = await asyncio.create_subprocess_exec(
            *cmd,
            cwd=cwd,
            env=env,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            preexec_fn=preexec_fn,
        )

        try:
            stdout, stderr = await asyncio.wait_for(
                process.communicate(),
                timeout=timeout,
            )
        except TimeoutError:
            # Kill process group if possible (POSIX)
            try:
                if os.name == "posix":
                    os.killpg(process.pid, signal.SIGKILL)
                else:
                    process.kill()
            except Exception:
                logger.debug("Failed to kill process %d after timeout", process.pid, exc_info=True)
            await process.wait()
            raise

        duration_ms = (time.monotonic() - start) * 1000

        result = ScriptResult(
            exit_code=process.returncode or 0,
            stdout=stdout.decode("utf-8", errors="replace"),
            stderr=stderr.decode("utf-8", errors="replace"),
            duration_ms=duration_ms,
            script_path=str(cmd[1]) if len(cmd) > 1 else "",
            args=cmd[2:] if len(cmd) > 2 else [],
        )

        if not result.success:
            logger.warning(
                f"Script failed with exit code {result.exit_code}: {result.stderr[:200]}"
            )

        return result

    def list_registered_tools(self) -> list[str]:
        """List all registered script tool IDs."""
        return list(self._registered_tools.keys())

    def get_tool_info(self, tool_id: str) -> dict[str, Any] | None:
        """Get info about a registered tool."""
        entry = self._registered_tools.get(tool_id)
        if not entry:
            return None

        skill, tool_decl = entry
        return {
            "tool_id": tool_id,
            "skill_id": skill.id,
            "script": tool_decl.script,
            "description": tool_decl.description,
            "default_args": tool_decl.default_args,
            "params": [p.name for p in tool_decl.params],
        }


# =============================================================================
# Convenience functions
# =============================================================================


def create_script_tool_provider(
    *,
    timeout_seconds: float = 300.0,
) -> ScriptToolProvider:
    """Create a script tool provider with default settings."""
    return ScriptToolProvider(timeout_seconds=timeout_seconds)


__all__ = [
    "InterpreterNotFoundError",
    # Errors
    "ScriptExecutionError",
    "ScriptNotFoundError",
    # Result
    "ScriptResult",
    # Provider
    "ScriptToolProvider",
    "create_script_tool_provider",
]
