"""Command execution orchestration."""

from __future__ import annotations

import asyncio
import base64
import contextlib
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

import structlog

from curadise_agent.auth.crypto import SignatureInfo, SignatureVerifier
from curadise_agent.executor.sandbox import SandboxConfig, SandboxedProcess
from curadise_agent.models.api.command import CommandRequest, CommandResponse, CommandStatus

if TYPE_CHECKING:
    from curadise_agent.config.schema import ExecutorConfig
    from curadise_agent.executor.validator import CommandValidator

log = structlog.get_logger(__name__)


@dataclass
class ExecutorStats:
    """Statistics about command execution."""

    total_executed: int = 0
    successful: int = 0
    failed: int = 0
    rejected: int = 0
    timed_out: int = 0
    average_duration_ms: float = 0.0


@dataclass
class CommandRunner:
    """
    Orchestrates command execution with validation and sandboxing.

    Handles the complete lifecycle of command execution including:
    - Signature verification
    - Command validation
    - Sandboxed execution
    - Output collection and truncation
    - Result reporting
    """

    validator: CommandValidator
    verifier: SignatureVerifier | None = None
    sandbox_config: SandboxConfig = field(default_factory=SandboxConfig)
    max_output_size: int = 1_048_576  # 1MB
    enabled: bool = True
    _stats: ExecutorStats = field(default_factory=ExecutorStats, init=False)
    _running_commands: dict[str, asyncio.Task[CommandResponse]] = field(
        default_factory=dict, init=False
    )

    @property
    def stats(self) -> ExecutorStats:
        """Get execution statistics."""
        return self._stats

    @property
    def running_count(self) -> int:
        """Get number of currently running commands."""
        return len(self._running_commands)

    async def execute(self, command: CommandRequest) -> CommandResponse:
        """
        Execute a command with full validation and sandboxing.

        Args:
            command: Command to execute

        Returns:
            CommandResponse with results
        """
        if not self.enabled:
            return CommandResponse.rejected(command.id, "Executor is disabled")

        log.info(
            "command_received",
            command_id=command.id,
            command_type=command.command_type.value,
        )

        # Validate signature if required
        if self.validator.require_signature:
            if not command.signature:
                self._stats.rejected += 1
                return CommandResponse.rejected(command.id, "Command signature required")

            if not self._verify_signature(command):
                self._stats.rejected += 1
                return CommandResponse.rejected(command.id, "Invalid command signature")

        # Validate command against allowlist/denylist
        validation = self.validator.validate(command)
        if not validation.valid:
            self._stats.rejected += 1
            log.warning(
                "command_rejected",
                command_id=command.id,
                reason=validation.reason,
            )
            return CommandResponse.rejected(command.id, validation.reason or "Command rejected")

        # Execute command
        try:
            response = await self._execute_sandboxed(command)
            self._update_stats(response)
            return response
        except Exception as e:
            self._stats.failed += 1
            log.error("command_execution_error", command_id=command.id, error=str(e))
            return CommandResponse.failed(command.id, str(e))

    async def execute_async(self, command: CommandRequest) -> str:
        """
        Start command execution asynchronously.

        Args:
            command: Command to execute

        Returns:
            Command ID for tracking
        """
        task = asyncio.create_task(self.execute(command))
        self._running_commands[command.id] = task
        return command.id

    async def get_result(self, command_id: str, timeout: float = 0.0) -> CommandResponse | None:
        """
        Get the result of an async command.

        Args:
            command_id: Command ID
            timeout: Time to wait (0 = don't wait)

        Returns:
            CommandResponse if available, None otherwise
        """
        task = self._running_commands.get(command_id)
        if task is None:
            return None

        if task.done():
            del self._running_commands[command_id]
            return task.result()

        if timeout > 0:
            try:
                await asyncio.wait_for(asyncio.shield(task), timeout=timeout)
                del self._running_commands[command_id]
                return task.result()
            except TimeoutError:
                return None

        return None

    async def cancel(self, command_id: str) -> bool:
        """
        Cancel a running command.

        Args:
            command_id: Command ID

        Returns:
            True if cancelled
        """
        task = self._running_commands.get(command_id)
        if task is None:
            return False

        task.cancel()
        with contextlib.suppress(asyncio.CancelledError):
            await task

        del self._running_commands[command_id]
        log.info("command_cancelled", command_id=command_id)
        return True

    def _verify_signature(self, command: CommandRequest) -> bool:
        """Verify command signature."""
        if self.verifier is None:
            log.warning("no_signature_verifier_configured")
            return False

        if not command.signature:
            return False

        try:
            signature_bytes = base64.b64decode(command.signature)
            sig_info = SignatureInfo(
                algorithm=command.signature_algorithm,
                key_id=command.key_id,
                signature=signature_bytes,
            )

            data = command.get_signature_data()
            return self.verifier.verify(data, sig_info)

        except Exception as e:
            log.error("signature_verification_failed", error=str(e))
            return False

    async def _execute_sandboxed(self, command: CommandRequest) -> CommandResponse:
        """Execute command in sandbox."""
        started_at = datetime.now(UTC)

        sandbox = SandboxedProcess(
            command=command.command,
            args=command.args,
            config=self.sandbox_config,
            working_dir=command.working_dir,
            env=command.env,
        )

        try:
            await sandbox.start()

            exit_code, stdout_bytes, stderr_bytes = await sandbox.wait(timeout=command.timeout)

            completed_at = datetime.now(UTC)
            duration_ms = (completed_at - started_at).total_seconds() * 1000

            # Decode and truncate output
            stdout, stdout_truncated = self._process_output(stdout_bytes)
            stderr, stderr_truncated = self._process_output(stderr_bytes)

            log.info(
                "command_completed",
                command_id=command.id,
                exit_code=exit_code,
                duration_ms=duration_ms,
            )

            return CommandResponse(
                command_id=command.id,
                status=CommandStatus.COMPLETED,
                exit_code=exit_code,
                stdout=stdout,
                stderr=stderr,
                started_at=started_at,
                completed_at=completed_at,
                duration_ms=duration_ms,
                truncated=stdout_truncated or stderr_truncated,
            )

        except TimeoutError:
            log.warning(
                "command_timeout",
                command_id=command.id,
                timeout=command.timeout,
            )
            return CommandResponse.timeout(command.id, started_at)

        except Exception as e:
            log.error(
                "command_execution_failed",
                command_id=command.id,
                error=str(e),
            )
            return CommandResponse.failed(command.id, str(e), started_at)

    def _process_output(self, output: bytes) -> tuple[str, bool]:
        """
        Process and optionally truncate output.

        Returns:
            Tuple of (output_string, was_truncated)
        """
        truncated = len(output) > self.max_output_size
        if truncated:
            output = output[: self.max_output_size]

        return output.decode("utf-8", errors="replace"), truncated

    def _update_stats(self, response: CommandResponse) -> None:
        """Update execution statistics."""
        self._stats.total_executed += 1

        if response.status == CommandStatus.COMPLETED and response.exit_code == 0:
            self._stats.successful += 1
        elif response.status == CommandStatus.TIMEOUT:
            self._stats.timed_out += 1
        elif response.status == CommandStatus.REJECTED:
            self._stats.rejected += 1
        else:
            self._stats.failed += 1

        # Update average duration
        if response.duration_ms > 0:
            n = self._stats.successful + self._stats.failed + self._stats.timed_out
            if n > 0:
                self._stats.average_duration_ms = (
                    self._stats.average_duration_ms * (n - 1) + response.duration_ms
                ) / n

    def get_status(self) -> dict[str, Any]:
        """Get executor status."""
        return {
            "enabled": self.enabled,
            "running_commands": self.running_count,
            "stats": {
                "total_executed": self._stats.total_executed,
                "successful": self._stats.successful,
                "failed": self._stats.failed,
                "rejected": self._stats.rejected,
                "timed_out": self._stats.timed_out,
                "average_duration_ms": self._stats.average_duration_ms,
            },
        }


def create_command_runner(
    config: ExecutorConfig,
    verifier: SignatureVerifier | None = None,
) -> CommandRunner:
    """
    Factory function to create a command runner.

    Args:
        config: Executor configuration
        verifier: Optional signature verifier

    Returns:
        Configured CommandRunner
    """
    from curadise_agent.executor.validator import create_validator

    validator = create_validator(config)

    sandbox_config = SandboxConfig(
        max_cpu_time=int(config.max_execution_time),
    )

    return CommandRunner(
        validator=validator,
        verifier=verifier,
        sandbox_config=sandbox_config,
        max_output_size=config.max_output_size,
        enabled=config.enabled,
    )
