"""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
"""

import asyncio
import contextlib
import logging
import os
import resource
import shlex
import shutil
import signal
import time
from pathlib import Path
from typing import Any

from pydantic import BaseModel, Field, computed_field

from core.pathways.context import Context
from core.pathways.models.tool import ToolOutput
from core.skills.loader import SCRIPT_INTERPRETERS
from core.skills.models import RuntimeContext, ScriptToolDeclaration, Skill
from core.tools import ToolDefinition, ToolRegistry

logger = logging.getLogger(__name__)

_CLI_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))],
}


# =============================================================================
# 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 = 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
        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.
        """
        converter = _CLI_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.exception(f"Script execution failed: {e}")
            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"])

        # Network isolation is not enabled by default; can be made configurable.

        # 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.exception(f"Docker execution failed: {e}")
            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",
]
