from __future__ import annotations

import os
from datetime import UTC, datetime
from typing import Any

from atlas_runner._docker_cli import DockerCLI
from atlas_runner.models import ContainerState, ContainerStatus, RunResult, RunSpec

_DOCKER_ZERO_TS = "0001-01-01T00:00:00Z"


def _parse_docker_ts(raw: str) -> datetime | None:
    """Parse a Docker timestamp, returning *None* for the zero-value."""
    if not raw or raw == _DOCKER_ZERO_TS:
        return None
    # Docker timestamps can have nanosecond precision (9 digits after `.`).
    # Python's fromisoformat only handles up to 6 (microseconds), so truncate.
    if "." in raw:
        base, frac_and_tz = raw.split(".", 1)
        # Separate fractional seconds from timezone suffix
        frac = ""
        tz = ""
        for i, ch in enumerate(frac_and_tz):
            if ch.isdigit():
                frac += ch
            else:
                tz = frac_and_tz[i:]
                break
        frac = frac[:6].ljust(6, "0")
        raw = f"{base}.{frac}{tz}"
    return datetime.fromisoformat(raw).astimezone(UTC)


def _parse_container_status(data: dict[str, Any]) -> ContainerStatus:
    """Map ``docker inspect`` JSON to a ``ContainerStatus``."""
    state = data["State"]
    name = data["Name"].lstrip("/")
    return ContainerStatus(
        container_id=data["Id"],
        name=name,
        state=ContainerState(state["Status"]),
        exit_code=state.get("ExitCode"),
        started_at=_parse_docker_ts(state.get("StartedAt", "")),
        finished_at=_parse_docker_ts(state.get("FinishedAt", "")),
    )


class DockerBackend:
    """Docker backend — translates ``RunSpec`` into ``docker`` CLI calls."""

    def __init__(self, docker_host: str | None = None) -> None:
        self._cli = DockerCLI(docker_host=docker_host)

    # ------------------------------------------------------------------
    # Protocol methods
    # ------------------------------------------------------------------

    def login(self, registry: str, username: str, password: str) -> None:
        self._cli.login(registry, username, password)

    def pull(self, image: str) -> None:
        self._cli.pull(image)

    def run(self, spec: RunSpec) -> RunResult:
        # Idempotent: return existing container if name already exists
        if self._cli.container_exists(spec.name):
            data = self._cli.inspect(spec.name)
            status = _parse_container_status(data)
            return RunResult(
                container_id=status.container_id,
                name=status.name,
                started_at=status.started_at or datetime.now(UTC),
            )

        args = self._build_run_args(spec)
        container_id = self._cli.run(args)
        data = self._cli.inspect(spec.name)
        status = _parse_container_status(data)
        return RunResult(
            container_id=container_id,
            name=status.name,
            started_at=status.started_at or datetime.now(UTC),
        )

    def stop(self, name: str) -> None:
        self._cli.stop(name)

    def rm(self, name: str) -> None:
        self._cli.rm(name)

    def inspect(self, name: str) -> ContainerStatus:
        data = self._cli.inspect(name)
        return _parse_container_status(data)

    def ps(self) -> list[ContainerStatus]:
        entries = self._cli.ps()
        results: list[ContainerStatus] = []
        for entry in entries:
            cid = entry.get("ID", entry.get("Id", ""))
            if cid:
                data = self._cli.inspect(cid)
                results.append(_parse_container_status(data))
        return results

    def logs(self, name: str) -> str:
        return self._cli.logs(name)

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    @staticmethod
    def _build_run_args(spec: RunSpec) -> list[str]:
        """Translate a ``RunSpec`` into ``docker run`` CLI arguments."""
        args: list[str] = ["-d", "--name", spec.name]

        # Labels
        for key, val in spec.labels.items():
            args.extend(["--label", f"{key}={val}"])

        # Build a local env copy so we never mutate the caller's RunSpec
        env = dict(spec.env)

        # GPU
        if spec.gpus is not None:
            devices = ",".join(str(g) for g in spec.gpus)
            args.extend(["--gpus", f"device={devices}"])
            # Only inject NVIDIA_DRIVER_CAPABILITIES (harmless, tells libs
            # which features to enable).  Do NOT set NVIDIA_VISIBLE_DEVICES
            # when using --gpus — the two mechanisms conflict and Docker
            # rejects the request with "cannot set both Count and DeviceIDs".
            env.setdefault("NVIDIA_DRIVER_CAPABILITIES", "compute,utility")

        # CPU
        if spec.cpu_cores is not None:
            args.extend(["--cpus", str(spec.cpu_cores)])
        elif spec.cpuset is not None:
            args.extend(["--cpuset-cpus", spec.cpuset])

        # Memory
        if spec.memory_gb is not None:
            args.extend(["--memory", f"{spec.memory_gb}g"])

        # Scratch
        if spec.scratch is not None:
            os.makedirs(spec.scratch, exist_ok=True)
            args.extend(["-v", f"{spec.scratch}:/scratch"])

        # Mounts
        for mount in spec.mounts:
            mount_str = f"{mount.src}:{mount.dst}"
            if mount.ro:
                mount_str += ":ro"
            args.extend(["-v", mount_str])

        # Security
        if spec.cap_drop_all:
            args.extend(["--cap-drop", "ALL"])
        if spec.no_new_privileges:
            args.extend(["--security-opt", "no-new-privileges"])

        # Network
        if spec.network is not None:
            args.extend(["--network", spec.network])

        # Workdir
        if spec.workdir is not None:
            args.extend(["--workdir", spec.workdir])

        # Environment variables
        for key, val in env.items():
            args.extend(["-e", f"{key}={val}"])

        # Ports
        for host_port, container_port in spec.ports.items():
            args.extend(["-p", f"{host_port}:{container_port}"])

        # Image (must come last before command)
        args.append(spec.image)

        # Command or shell
        if spec.command is not None:
            args.extend(spec.command)
        elif spec.shell is not None:
            args.extend(["sh", "-c", spec.shell])

        return args
