import asyncio
import asyncio.subprocess
import io
import os
import platform
import shlex
import shutil
import tarfile
from pathlib import Path, PurePosixPath

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


_STREAM_CHUNK_SIZE = 65536  # 64 KB


class AppleContainerEnvironment(BaseEnvironment):
    """Environment using Apple Container (lightweight Linux VMs on Apple silicon)."""

    # Class-level lock per image name to prevent parallel builds of the same image.
    _image_build_locks: dict[str, asyncio.Lock] = {}

    @classmethod
    def preflight(cls) -> None:
        if platform.machine() != "arm64":
            raise SystemExit(
                "Apple Container requires a Mac with Apple silicon (arm64)."
            )
        if not shutil.which("container"):
            raise SystemExit(
                "Apple Container requires the 'container' CLI to be installed. "
                "Download it from https://github.com/apple/container/releases"
            )

    def __init__(
        self,
        environment_dir: Path,
        environment_name: str,
        session_id: str,
        trial_paths: TrialPaths,
        task_env_config: EnvironmentConfig,
        keep_containers: bool = False,
        *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,
            **kwargs,
        )

        self._keep_containers = keep_containers
        self._image_name = f"hb__{environment_name.lower()}"
        self._container_name = session_id.lower().replace(".", "-")
        self._use_prebuilt = False

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

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

    @property
    def supports_gpus(self) -> bool:
        return False

    @property
    def can_disable_internet(self) -> bool:
        return False

    def _validate_definition(self):
        dockerfile_path = self.environment_dir / "Dockerfile"
        if not dockerfile_path.exists():
            raise FileNotFoundError(
                f"{dockerfile_path} not found. Apple Container environment requires "
                "a Dockerfile."
            )

    async def _run_container_command(
        self,
        args: list[str],
        check: bool = True,
        timeout_sec: int | None = None,
        stdin_data: bytes | None = None,
    ) -> ExecResult:
        """Run a `container` CLI command and return the result."""
        full_command = ["container", *args]

        process = await asyncio.create_subprocess_exec(
            *full_command,
            stdin=asyncio.subprocess.PIPE if stdin_data else asyncio.subprocess.DEVNULL,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

        try:
            if timeout_sec:
                stdout_bytes, stderr_bytes = await asyncio.wait_for(
                    process.communicate(input=stdin_data), timeout=timeout_sec
                )
            else:
                stdout_bytes, stderr_bytes = await process.communicate(input=stdin_data)
        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"
            ) from None

        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"Container 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 bool(self.task_env_config.docker_image)

        if not self._use_prebuilt:
            lock = self._image_build_locks.setdefault(
                self.environment_name, asyncio.Lock()
            )
            async with lock:
                await self._run_container_command(
                    [
                        "build",
                        "-t",
                        self._image_name,
                        "-f",
                        str((self.environment_dir / "Dockerfile").resolve().absolute()),
                        str(self.environment_dir.resolve().absolute()),
                    ],
                    timeout_sec=int(self.task_env_config.build_timeout_sec),
                )

        image: str = (
            self.task_env_config.docker_image
            if self._use_prebuilt and self.task_env_config.docker_image
            else self._image_name
        )

        # Best-effort cleanup of stale containers from previous runs.
        try:
            await self._run_container_command(
                ["stop", self._container_name], check=True
            )
        except RuntimeError:
            pass
        try:
            await self._run_container_command(["rm", self._container_name], check=True)
        except RuntimeError:
            pass

        # Build the run command.
        run_cmd: list[str] = ["run", "-d", "--name", self._container_name]

        # Resource limits.
        run_cmd.extend(["-c", str(self.task_env_config.cpus)])
        run_cmd.extend(["-m", f"{self.task_env_config.memory_mb}M"])

        # Bind-mount log directories.
        mounts = {
            str(self.trial_paths.verifier_dir.resolve().absolute()): str(
                EnvironmentPaths.verifier_dir
            ),
            str(self.trial_paths.agent_dir.resolve().absolute()): str(
                EnvironmentPaths.agent_dir
            ),
            str(self.trial_paths.artifacts_dir.resolve().absolute()): str(
                EnvironmentPaths.artifacts_dir
            ),
        }
        for host_path, container_path in mounts.items():
            run_cmd.extend(["-v", f"{host_path}:{container_path}"])

        run_cmd.append(image)
        run_cmd.extend(["sh", "-c", "sleep infinity"])

        await self._run_container_command(run_cmd)

    async def stop(self, delete: bool):
        # Best-effort: fix ownership of bind-mounted directories.
        await self._chown_to_host_user(str(EnvironmentPaths.logs_dir), recursive=True)

        if self._keep_containers and delete:
            self.logger.warning(
                "Both `keep_containers` and `--delete` option are set. "
                "keep_containers takes precedence."
            )

        # Always stop the container.
        try:
            await self._run_container_command(["stop", self._container_name])
        except RuntimeError as e:
            self.logger.warning(f"Container stop failed: {e}")

        if not self._keep_containers:
            try:
                await self._run_container_command(["rm", self._container_name])
            except RuntimeError as e:
                self.logger.warning(f"Container rm failed: {e}")

            if delete and not self._use_prebuilt:
                try:
                    await self._run_container_command(["image", "rm", self._image_name])
                except RuntimeError as e:
                    self.logger.warning(f"Image removal failed: {e}")

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

        exec_cmd: list[str] = ["exec"]

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

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

        exec_cmd.append(self._container_name)
        exec_cmd.extend(["bash", "-c", command])

        return await self._run_container_command(
            exec_cmd, check=False, timeout_sec=timeout_sec
        )

    async def _chown_to_host_user(self, path: str, recursive: bool = False) -> None:
        """Best-effort chown of a container path to the host user's UID:GID."""
        if not hasattr(os, "getuid"):
            return
        flag = "-R " if recursive else ""
        await self.exec(f"chown {flag}{os.getuid()}:{os.getgid()} {shlex.quote(path)}")

    async def _upload_tar(self, buf: io.BytesIO, target_dir: str) -> None:
        """Pipe a tar archive into the container, extracting at target_dir."""
        await self.exec(f"mkdir -p {shlex.quote(target_dir)}", timeout_sec=10)
        await self._run_container_command(
            ["exec", "-i", self._container_name, "tar", "xf", "-", "-C", target_dir],
            stdin_data=buf.getvalue(),
            check=True,
        )

    async def _download_tar(
        self, tar_args: list[str], target_path: Path, description: str
    ) -> None:
        """Run tar in the container and stream-extract the result locally."""
        target_path.mkdir(parents=True, exist_ok=True)

        process = await asyncio.create_subprocess_exec(
            "container",
            "exec",
            self._container_name,
            "tar",
            "cf",
            "-",
            *tar_args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

        read_fd, write_fd = os.pipe()
        loop = asyncio.get_running_loop()

        async def _pump() -> None:
            try:
                assert process.stdout is not None
                while True:
                    chunk = await process.stdout.read(_STREAM_CHUNK_SIZE)
                    if not chunk:
                        break
                    try:
                        await loop.run_in_executor(None, os.write, write_fd, chunk)
                    except OSError:
                        break  # Pipe broken — extractor likely failed
            finally:
                os.close(write_fd)

        def _extract() -> None:
            with os.fdopen(read_fd, "rb") as rf:
                with tarfile.open(fileobj=rf, mode="r|") as tar:
                    tar.extractall(path=str(target_path), filter="data")

        async def _extract_in_thread() -> None:
            await loop.run_in_executor(None, _extract)

        gather_error: Exception | None = None
        try:
            async with asyncio.TaskGroup() as tg:
                tg.create_task(_pump())
                tg.create_task(_extract_in_thread())
        except* Exception as eg:
            gather_error = eg.exceptions[0]
            process.kill()

        stderr_bytes = await process.stderr.read() if process.stderr else b""
        await process.wait()

        # Process failure takes priority over extraction errors (e.g. tarfile
        # raising ReadError on the empty/truncated stream).
        if process.returncode != 0:
            raise RuntimeError(
                f"Failed to download {description} from container "
                f"{self._container_name}. "
                f"Stderr: {stderr_bytes.decode(errors='replace')}"
            )

        if gather_error is not None:
            raise gather_error

    async def upload_file(self, source_path: Path | str, target_path: str):
        source_path = Path(source_path)
        target = PurePosixPath(target_path)

        buf = io.BytesIO()
        with tarfile.open(fileobj=buf, mode="w") as tar:
            tar.add(str(source_path), arcname=target.name)
        buf.seek(0)

        await self._upload_tar(buf, str(target.parent))

    async def upload_dir(self, source_dir: Path | str, target_dir: str):
        source_dir = Path(source_dir)

        buf = io.BytesIO()
        with tarfile.open(fileobj=buf, mode="w") as tar:
            for child in source_dir.iterdir():
                tar.add(str(child), arcname=child.name)
        buf.seek(0)

        await self._upload_tar(buf, target_dir)

    async def download_file(self, source_path: str, target_path: Path | str):
        target_path = Path(target_path)
        source = PurePosixPath(source_path)

        await self._chown_to_host_user(source_path)
        await self._download_tar(
            ["-C", str(source.parent), source.name],
            target_path.parent,
            f"file {source_path}",
        )

        # tar preserves the source filename; rename to match the caller's target.
        extracted = target_path.parent / source.name
        if extracted != target_path:
            extracted.rename(target_path)

    async def download_dir(self, source_dir: str, target_dir: Path | str):
        target_dir = Path(target_dir)

        await self._chown_to_host_user(source_dir, recursive=True)
        await self._download_tar(
            ["-C", source_dir, "."],
            target_dir,
            f"directory {source_dir}",
        )

    async def attach(self) -> None:
        os.execvp(
            "container",
            ["container", "exec", "-it", self._container_name, "bash"],
        )
