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

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

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: Optional[Path] = None
    timeout: Optional[float] = 60.0
    env: Optional[dict[str, str]] = 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: Optional[SafetyGuard] = 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: Optional[Path],
    ) -> 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: Optional[str] = None,
        timeout: Optional[int] = None,
        env: Optional[str] = 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 = {}

        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: Optional[Path],
        timeout: float,
        env: Optional[dict[str, str]],
        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
