import asyncio
import asyncio.subprocess
import os
import shlex
from pathlib import Path

from pydantic import BaseModel

from harbor.environments.base import BaseEnvironment, ExecResult
from harbor.models.environment_type import EnvironmentType
from harbor.models.task.config import EnvironmentConfig
from harbor.models.trial.paths import EnvironmentPaths, TrialPaths


class DockerEnvironmentEnvVars(BaseModel):
    main_image_name: str
    context_dir: str
    test_dir: str
    host_verifier_logs_path: str
    host_agent_logs_path: str
    env_verifier_logs_path: str
    env_agent_logs_path: str
    prebuilt_image_name: str | None = None
    cpus: int = 1
    memory: str = "1G"

    def to_env_dict(self, include_os_env: bool = True) -> dict[str, str]:
        env_dict = {} if not include_os_env else os.environ.copy()

        for field_name, value in self.model_dump(exclude_none=True).items():
            if value is None:
                continue

            env_dict[f"{field_name.upper()}"] = str(value)

        return env_dict


class DockerEnvironment(BaseEnvironment):
    _DOCKER_COMPOSE_BUILD_PATH = Path(__file__).parent / "docker-compose-build.yaml"
    _DOCKER_COMPOSE_PREBUILT_PATH = (
        Path(__file__).parent / "docker-compose-prebuilt.yaml"
    )

    def __init__(
        self,
        environment_dir: Path,
        environment_name: str,
        session_id: str,
        trial_paths: TrialPaths,
        task_env_config: EnvironmentConfig,
        *args,
        **kwargs,
    ):
        super().__init__(
            environment_dir=environment_dir,
            environment_name=environment_name,
            session_id=session_id,
            trial_paths=trial_paths,
            task_env_config=task_env_config,
            *args,
            **kwargs,
        )

        self._env_vars = DockerEnvironmentEnvVars(
            main_image_name=f"hb__{environment_name}",
            context_dir=str(self.environment_dir.resolve().absolute()),
            test_dir=str(EnvironmentPaths.tests_dir),
            host_verifier_logs_path=str(trial_paths.verifier_dir.resolve().absolute()),
            host_agent_logs_path=str(trial_paths.agent_dir.resolve().absolute()),
            env_verifier_logs_path=str(EnvironmentPaths.verifier_dir),
            env_agent_logs_path=str(EnvironmentPaths.agent_dir),
            prebuilt_image_name=task_env_config.docker_image,
            cpus=task_env_config.cpus,
            memory=f"{task_env_config.memory_mb}M",
        )
        self._use_prebuilt = False

    @staticmethod
    def type() -> EnvironmentType:
        return EnvironmentType.DOCKER

    @property
    def is_mounted(self) -> bool:
        return True

    @property
    def _dockerfile_path(self) -> Path:
        return self.environment_dir / "Dockerfile"

    @property
    def _environment_docker_compose_path(self) -> Path:
        return self.environment_dir / "docker-compose.yaml"

    @property
    def _docker_compose_path(self) -> Path:
        """Returns the appropriate docker-compose file based on whether we're using a prebuilt image."""
        if self._environment_docker_compose_path.exists():
            return self._environment_docker_compose_path

        if self._use_prebuilt:
            return self._DOCKER_COMPOSE_PREBUILT_PATH

        return self._DOCKER_COMPOSE_BUILD_PATH

    def _validate_definition(self):
        if (
            not self._dockerfile_path.exists()
            and not self._environment_docker_compose_path.exists()
        ):
            raise FileNotFoundError(
                f"{self._dockerfile_path} and {self._environment_docker_compose_path} "
                "not found. Please ensure at least one of these files exist."
            )

    async def _run_docker_compose_command(
        self, command: list[str], check: bool = True, timeout_sec: int | None = None
    ) -> ExecResult:
        """Run a docker compose command and return the result."""
        full_command = [
            "docker",
            "compose",
            "-p",
            self.session_id.lower().replace(".", "-"),
            "-f",
            str(self._docker_compose_path.resolve().absolute()),
            *command,
        ]

        process = await asyncio.create_subprocess_exec(
            *full_command,
            env=self._env_vars.to_env_dict(),
            stdin=asyncio.subprocess.DEVNULL,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
        )

        try:
            if timeout_sec:
                stdout_bytes, stderr_bytes = await asyncio.wait_for(
                    process.communicate(), timeout=timeout_sec
                )
            else:
                stdout_bytes, stderr_bytes = await process.communicate()
        except asyncio.TimeoutError:
            process.terminate()
            try:
                stdout_bytes, stderr_bytes = await asyncio.wait_for(
                    process.communicate(), timeout=5
                )
            except asyncio.TimeoutError:
                process.kill()
                stdout_bytes, stderr_bytes = await process.communicate()
            raise RuntimeError(f"Command timed out after {timeout_sec} seconds")

        stdout = stdout_bytes.decode(errors="replace") if stdout_bytes else None
        stderr = stderr_bytes.decode(errors="replace") if stderr_bytes else None

        result = ExecResult(
            stdout=stdout,
            stderr=stderr,
            return_code=process.returncode or 0,
        )

        if check and result.return_code != 0:
            raise RuntimeError(
                f"Docker compose command failed for environment {self.environment_name}. "
                f"Command: {' '.join(full_command)}. "
                f"Return code: {result.return_code}. "
                f"Stdout: {result.stdout}. "
                f"Stderr: {result.stderr}. "
            )

        return result

    async def start(self, force_build: bool):
        self._use_prebuilt = (
            not force_build and self.task_env_config.docker_image is not None
        )

        if force_build:
            await self._run_docker_compose_command(["build"])

        await self._run_docker_compose_command(["up", "-d"])

    async def stop(self, delete: bool):
        if delete:
            try:
                await self._run_docker_compose_command(
                    ["down", "--rmi", "all", "--volumes", "--remove-orphans"]
                )
            except RuntimeError as e:
                print(f"Warning: Docker compose down failed: {e}")

            await self._cleanup_build_cache()
        else:
            try:
                await self._run_docker_compose_command(["down"])
            except RuntimeError as e:
                print(f"Warning: Docker compose down failed: {e}")

    async def upload_file(self, source_path: Path | str, target_path: str):
        await self._run_docker_compose_command(
            [
                "cp",
                str(source_path),
                f"main:{target_path}",
            ],
            check=True,
        )

    async def upload_dir(self, source_dir: Path | str, target_dir: str):
        await self._run_docker_compose_command(
            [
                "cp",
                str(source_dir),
                f"main:{target_dir}",
            ],
            check=True,
        )

    async def download_file(self, source_path: str, target_path: Path | str):
        await self._run_docker_compose_command(
            [
                "cp",
                f"main:{source_path}",
                str(target_path),
            ],
            check=True,
        )

    async def download_dir(self, source_dir: str, target_dir: Path | str):
        await self._run_docker_compose_command(
            [
                "cp",
                f"main:{source_dir}",
                str(target_dir),
            ],
            check=True,
        )

    async def exec(
        self,
        command: str,
        cwd: str | None = None,
        env: dict[str, str] | None = None,
        timeout_sec: int | None = None,
    ) -> ExecResult:
        exec_command = ["exec", "-it"]

        if cwd:
            exec_command.extend(["-w", cwd])

        if env:
            for key, value in env.items():
                exec_command.extend(["-e", f"{key}={shlex.quote(value)}"])

        exec_command.append("main")
        exec_command.extend(["bash", "-ic", command])

        return await self._run_docker_compose_command(
            exec_command, check=False, timeout_sec=timeout_sec
        )

    async def _cleanup_build_cache(self) -> None:
        """Clean up Docker build cache to free up disk space."""
        try:
            # Use --max-used-space to keep cache under the specified size limit
            process = await asyncio.create_subprocess_exec(
                "docker",
                "buildx",
                "prune",
                "--force",
                "--max-used-space",
                "30GB",
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
            await process.wait()
        except Exception:
            # Fallback to docker builder prune if buildx is not available
            try:
                process = await asyncio.create_subprocess_exec(
                    "docker",
                    "builder",
                    "prune",
                    "--force",
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                )
                await process.wait()
            except Exception:
                # Silent fail - build cache cleanup is not critical
                pass
