"""Exec tool for executing shell commands."""

from __future__ import annotations

import asyncio
import shlex
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from oclawma.safety import RiskLevel, SafetyGuard
from oclawma.tools.base import (
    BaseTool,
    ToolExecutionError,
    ToolParameter,
    ToolResult,
    ToolSchema,
    ToolTimeoutError,
)


@dataclass
class ExecOptions:
    """Options for shell command execution."""

    cwd: Path | None = None
    timeout: float | None = 60.0
    env: dict[str, str] | None = None
    shell: bool = True
    capture_output: bool = True
    capture_stderr: bool = True
    merge_stderr: bool = True


class ExecTool(BaseTool):
    """Tool for executing shell commands.

    Provides safe shell command execution with configurable options
    for working directory, timeout, and environment variables.
    Includes safety guard integration for risky operation detection.
    """

    DEFAULT_TIMEOUT = 60.0
    MAX_TIMEOUT = 300.0  # 5 minutes max

    def __init__(
        self,
        default_timeout: float = DEFAULT_TIMEOUT,
        safety_guard: SafetyGuard | None = None,
    ) -> None:
        """Initialize the exec tool.

        Args:
            default_timeout: Default timeout in seconds
            safety_guard: Optional safety guard instance for risk analysis
        """
        self._default_timeout = min(default_timeout, self.MAX_TIMEOUT)
        self._safety_guard = safety_guard
        super().__init__()

    def _define_schema(self) -> ToolSchema:
        """Define the exec tool schema."""
        return ToolSchema(
            name="exec",
            description="Execute a shell command and return the output.",
            parameters=[
                ToolParameter(
                    name="command",
                    description="The shell command to execute",
                    type="string",
                    required=True,
                ),
                ToolParameter(
                    name="cwd",
                    description="Working directory for command execution (default: current directory)",
                    type="string",
                    required=False,
                ),
                ToolParameter(
                    name="timeout",
                    description=f"Timeout in seconds (default: {int(self._default_timeout)}, max: 300)",
                    type="integer",
                    required=False,
                ),
                ToolParameter(
                    name="env",
                    description="Environment variables as key=value pairs (comma-separated)",
                    type="string",
                    required=False,
                ),
                ToolParameter(
                    name="shell",
                    description="Execute through shell (default: true)",
                    type="boolean",
                    required=False,
                    default=True,
                ),
            ],
            returns="string",
            examples=[
                'exec(command="ls -la")',
                'exec(command="git status", cwd="/home/user/project")',
                'exec(command="sleep 5", timeout=10)',
                'exec(command="echo $MY_VAR", env="MY_VAR=hello")',
            ],
        )

    def _check_safety(
        self,
        command: str,
    ) -> tuple[list[Any], RiskLevel]:
        """Check command for safety risks.

        Args:
            command: The command to check

        Returns:
            Tuple of (risks list, max risk level)
        """
        if self._safety_guard is None:
            return [], RiskLevel.LOW

        risks = self._safety_guard.analyze_command(command)
        max_risk = self._safety_guard.get_max_risk_level(risks)
        return risks, max_risk

    def _simulate_command(
        self,
        command: str,
        cwd: Path | None,
    ) -> ToolResult:
        """Simulate a command execution for dry-run mode.

        Args:
            command: The command to simulate
            cwd: Working directory

        Returns:
            ToolResult indicating simulation success
        """
        return ToolResult(
            tool_name=self.name,
            success=True,
            output=f"[DRY RUN] Command would execute: {command}",
            error_message=None,
            metadata={
                "command": command,
                "cwd": str(cwd) if cwd else None,
                "dry_run": True,
                "simulated": True,
            },
        )

    async def execute(
        self,
        command: str,
        cwd: str | None = None,
        timeout: int | None = None,
        env: str | None = None,
        shell: bool = True,
    ) -> ToolResult:
        """Execute a shell command.

        Args:
            command: The command to execute
            cwd: Working directory
            timeout: Timeout in seconds
            env: Environment variables string (comma-separated key=value)
            shell: Whether to use shell execution

        Returns:
            ToolResult with command output
        """
        # Validate command
        if not command or not command.strip():
            return ToolResult(
                tool_name=self.name,
                success=False,
                output="",
                error_message="Command cannot be empty",
            )

        # Parse timeout
        exec_timeout = self._default_timeout
        if timeout is not None:
            exec_timeout = min(float(timeout), self.MAX_TIMEOUT)

        # Parse working directory
        working_dir = None
        if cwd:
            working_dir = Path(cwd).expanduser().resolve()
            if not working_dir.exists():
                return ToolResult(
                    tool_name=self.name,
                    success=False,
                    output="",
                    error_message=f"Working directory does not exist: {cwd}",
                    metadata={"cwd": str(cwd)},
                )
            if not working_dir.is_dir():
                return ToolResult(
                    tool_name=self.name,
                    success=False,
                    output="",
                    error_message=f"Path is not a directory: {cwd}",
                    metadata={"cwd": str(cwd)},
                )

        # Parse environment variables
        env_vars = None
        if env:
            env_vars = self._parse_env_string(env)

        # Safety check
        risks, max_risk = self._check_safety(command)
        is_dry_run = self._safety_guard.dry_run if self._safety_guard else False

        # Log the operation attempt (before execution)
        if self._safety_guard:
            self._safety_guard.log_audit(
                tool_name=self.name,
                operation="command_execution_attempt",
                parameters={
                    "command": command,
                    "cwd": str(working_dir) if working_dir else None,
                    "timeout": exec_timeout,
                },
                risk_level=max_risk,
                was_blocked=False,
                was_dry_run=is_dry_run,
            )

        # Handle dry-run mode
        if is_dry_run:
            result = self._simulate_command(command, working_dir)
            # Log the dry-run result
            if self._safety_guard:
                self._safety_guard.log_audit(
                    tool_name=self.name,
                    operation="command_execution",
                    parameters={
                        "command": command,
                        "cwd": str(working_dir) if working_dir else None,
                    },
                    risk_level=max_risk,
                    was_blocked=False,
                    was_dry_run=True,
                    result=result.to_dict(),
                )
            return result

        # Execute the command
        try:
            result = await self._run_command(
                command=command,
                cwd=working_dir,
                timeout=exec_timeout,
                env=env_vars,
                shell=shell,
            )

            # Log successful execution
            if self._safety_guard:
                self._safety_guard.log_audit(
                    tool_name=self.name,
                    operation="command_execution",
                    parameters={
                        "command": command,
                        "cwd": str(working_dir) if working_dir else None,
                    },
                    risk_level=max_risk,
                    was_blocked=False,
                    was_dry_run=False,
                    result=result.to_dict(),
                )

            return result

        except ToolTimeoutError:
            error_result = ToolResult(
                tool_name=self.name,
                success=False,
                output="",
                error_message=f"Command timed out after {exec_timeout} seconds",
                metadata={
                    "command": command,
                    "timeout": exec_timeout,
                    "cwd": str(working_dir) if working_dir else None,
                },
            )
            # Log timeout
            if self._safety_guard:
                self._safety_guard.log_audit(
                    tool_name=self.name,
                    operation="command_execution_timeout",
                    parameters={
                        "command": command,
                        "cwd": str(working_dir) if working_dir else None,
                        "timeout": exec_timeout,
                    },
                    risk_level=max_risk,
                    was_blocked=False,
                    was_dry_run=False,
                    result=error_result.to_dict(),
                )
            return error_result

        except Exception as e:
            error_result = ToolResult(
                tool_name=self.name,
                success=False,
                output="",
                error_message=str(e),
                metadata={
                    "command": command,
                    "cwd": str(working_dir) if working_dir else None,
                },
            )
            # Log error
            if self._safety_guard:
                self._safety_guard.log_audit(
                    tool_name=self.name,
                    operation="command_execution_error",
                    parameters={
                        "command": command,
                        "cwd": str(working_dir) if working_dir else None,
                        "error": str(e),
                    },
                    risk_level=max_risk,
                    was_blocked=False,
                    was_dry_run=False,
                    result=error_result.to_dict(),
                )
            return error_result

    def _parse_env_string(self, env_str: str) -> dict[str, str]:
        """Parse environment variable string into dictionary.

        Format: KEY1=value1,KEY2=value2 or KEY1=value1 KEY2=value2

        Args:
            env_str: Environment variables string

        Returns:
            Dictionary of environment variables
        """
        env_vars = {}

        # Try comma-separated first
        pairs = [p.strip() for p in env_str.split(",")] if "," in env_str else shlex.split(env_str)

        for pair in pairs:
            if "=" in pair:
                key, value = pair.split("=", 1)
                env_vars[key.strip()] = value.strip()

        return env_vars

    async def _run_command(
        self,
        command: str,
        cwd: Path | None,
        timeout: float,
        env: dict[str, str] | None,
        shell: bool,
    ) -> ToolResult:
        """Run the actual command.

        Args:
            command: Command string
            cwd: Working directory
            timeout: Timeout in seconds
            env: Environment variables
            shell: Use shell execution

        Returns:
            ToolResult with output
        """
        # Prepare environment
        cmd_env = None
        if env:
            import os

            cmd_env = {**os.environ, **env}

        try:
            # Create subprocess
            process = await asyncio.create_subprocess_shell(
                command,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                cwd=cwd,
                env=cmd_env,
            )

            # Wait for completion with timeout
            try:
                stdout, stderr = await asyncio.wait_for(
                    process.communicate(),
                    timeout=timeout,
                )
            except asyncio.TimeoutError:
                process.kill()
                await process.wait()
                raise ToolTimeoutError(
                    f"Command timed out after {timeout}s",
                    tool_name=self.name,
                ) from None

            # Decode output
            stdout_str = stdout.decode("utf-8", errors="replace") if stdout else ""
            stderr_str = stderr.decode("utf-8", errors="replace") if stderr else ""

            # Combine stdout and stderr if there's stderr output
            output = stdout_str
            if stderr_str:
                if output:
                    output += "\n"
                output += f"[stderr]\n{stderr_str}"

            # Success if return code is 0
            success = process.returncode == 0

            return ToolResult(
                tool_name=self.name,
                success=success,
                output=output,
                error_message=None if success else f"Exit code: {process.returncode}",
                metadata={
                    "command": command,
                    "exit_code": process.returncode,
                    "cwd": str(cwd) if cwd else None,
                },
            )

        except ToolTimeoutError:
            raise
        except Exception as e:
            raise ToolExecutionError(
                f"Failed to execute command: {e}",
                tool_name=self.name,
                cause=e,
            ) from e

    def validate_params(self, params: dict[str, Any]) -> dict[str, Any]:
        """Validate and sanitize exec parameters.

        Args:
            params: Parameters to validate

        Returns:
            Validated parameters
        """
        validated = super().validate_params(params)

        # Sanitize command (basic)
        if "command" in validated:
            command = validated["command"].strip()
            # Remove any null bytes
            command = command.replace("\x00", "")
            validated["command"] = command

        return validated
