from __future__ import annotations

import asyncio
import subprocess
import sys
from pathlib import Path

from ai_orchestrator.config.settings import Settings
from ai_orchestrator.schemas.task_state import utc_now
from ai_orchestrator.schemas.tool_result import ToolResult


class DockerSandbox:
    def __init__(self, settings: Settings) -> None:
        self.settings = settings

    async def run_command(
        self,
        *,
        command: str,
        repo_path: Path,
        timeout_seconds: int | None = None,
        use_docker: bool | None = None,
        sandbox_image: str | None = None,
    ) -> ToolResult:
        timeout = timeout_seconds or self.settings.default_command_timeout_seconds
        use_docker = self.settings.runtime_use_docker if use_docker is None else use_docker
        if use_docker:
            invocation = self._docker_invocation(command, repo_path, sandbox_image)
            return await self._run_process(invocation, cwd=None, timeout_seconds=timeout, command=command)
        invocation = self._local_shell_invocation(command)
        return await self._run_process(invocation, cwd=repo_path, timeout_seconds=timeout, command=command)

    def _local_shell_invocation(self, command: str) -> list[str]:
        if sys.platform == "win32":
            return ["powershell", "-NoProfile", "-Command", command]
        return ["bash", "-lc", command]

    def _docker_invocation(self, command: str, repo_path: Path, sandbox_image: str | None) -> list[str]:
        image = sandbox_image or self.settings.docker_image
        invocation = [
            "docker",
            "run",
            "--rm",
            "--cpus",
            str(self.settings.docker_cpus),
            "--memory",
            self.settings.docker_memory,
            "-v",
            f"{repo_path.resolve()}:/workspace/repo",
            "-w",
            "/workspace/repo",
        ]
        if self.settings.docker_network_disabled:
            invocation.extend(["--network", "none"])
        invocation.extend([image, "sh", "-lc", command])
        return invocation

    async def _run_process(
        self,
        invocation: list[str],
        *,
        cwd: Path | None,
        timeout_seconds: int,
        command: str,
    ) -> ToolResult:
        started_at = utc_now()
        try:
            completed = await asyncio.to_thread(
                subprocess.run,
                invocation,
                cwd=str(cwd) if cwd else None,
                capture_output=True,
                text=True,
                encoding="utf-8",
                errors="replace",
                timeout=timeout_seconds,
                check=False,
            )
        except subprocess.TimeoutExpired:
            return ToolResult(
                success=False,
                stdout="",
                stderr=f"Command timed out after {timeout_seconds} seconds.",
                exit_code=None,
                command=command,
                timed_out=True,
                started_at=started_at,
                completed_at=utc_now(),
            )

        return ToolResult(
            success=completed.returncode == 0,
            stdout=completed.stdout,
            stderr=completed.stderr,
            exit_code=completed.returncode,
            command=command,
            started_at=started_at,
            completed_at=utc_now(),
        )
