# retoor <retoor@molodetz.nl>

from __future__ import annotations

import asyncio
import json
import logging
import re
import tempfile
from pathlib import Path
from typing import Optional

from devplacepy.services.containers.backend.base import (
    WORKSPACE_MOUNT,
    Backend,
    BuildResult,
    ExecResult,
    LogCallback,
    PsRow,
    RunSpec,
    StatsSample,
)

logger = logging.getLogger("containers.docker")

DOCKER = "docker"
DEFAULT_TIMEOUT = 600
BUILD_TIMEOUT = 3600
_EXIT_RE = re.compile(r"Exited \((\d+)\)")
_SIZE_RE = re.compile(r"([0-9.]+)\s*([kKmMgGtT]?i?)B")


def build_run_argv(spec: RunSpec) -> list:
    argv = [DOCKER, "run", "-d", "--name", spec.name]
    for key, value in spec.labels.items():
        argv += ["--label", f"{key}={value}"]
    for key, value in spec.env.items():
        argv += ["-e", f"{key}={value}"]
    if spec.cpu_limit:
        argv += ["--cpus", str(spec.cpu_limit)]
    if spec.mem_limit:
        argv += ["--memory", str(spec.mem_limit)]
    for port in spec.ports:
        argv += ["-p", f"{port.host}:{port.container}/{port.proto}"]
    for mount in spec.mounts:
        argv += ["-v", f"{mount.host}:{mount.container}:{mount.mode}"]
    if spec.restart_policy and spec.restart_policy != "never":
        policy = "no" if spec.restart_policy in ("no", "never") else spec.restart_policy
        argv += ["--restart", policy]
    argv.append(spec.image)
    argv += list(spec.command)
    return argv


def build_image_argv(
    context_dir: str, dockerfile_path: str, tags: list, iidfile: str, network: str = ""
) -> list:
    argv = [
        DOCKER,
        "build",
        "--label",
        "devplace.build=1",
        "--iidfile",
        iidfile,
        "-f",
        dockerfile_path,
    ]
    if network:
        argv += ["--network", network]
    for tag in tags:
        argv += ["-t", tag]
    argv.append(context_dir)
    return argv


def parse_size(text: str) -> int:
    match = _SIZE_RE.search(text or "")
    if not match:
        return 0
    value = float(match.group(1))
    unit = match.group(2).lower().rstrip("i")
    factor = {"": 1, "k": 1024, "m": 1024**2, "g": 1024**3, "t": 1024**4}.get(unit, 1)
    return int(value * factor)


def parse_stats_line(row: dict) -> StatsSample:
    def _pair(text: str) -> tuple:
        parts = (text or "").split("/")
        left = parse_size(parts[0]) if parts else 0
        right = parse_size(parts[1]) if len(parts) > 1 else 0
        return left, right

    rx, tx = _pair(row.get("NetIO", ""))
    rd, wr = _pair(row.get("BlockIO", ""))
    return StatsSample(
        cpu_pct=float((row.get("CPUPerc", "0%") or "0%").rstrip("%") or 0),
        mem_bytes=_pair(row.get("MemUsage", ""))[0],
        mem_pct=float((row.get("MemPerc", "0%") or "0%").rstrip("%") or 0),
        net_rx=rx,
        net_tx=tx,
        blk_read=rd,
        blk_write=wr,
    )


class DockerCliBackend(Backend):
    def __init__(
        self,
        binary: str = DOCKER,
        timeout: int = DEFAULT_TIMEOUT,
        build_timeout: int = BUILD_TIMEOUT,
    ) -> None:
        self._binary = binary
        self._timeout = timeout
        self._build_timeout = build_timeout

    async def _run(self, argv: list, *, timeout: Optional[int] = None) -> tuple:
        proc = await asyncio.create_subprocess_exec(
            *argv, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
        )
        try:
            out, err = await asyncio.wait_for(
                proc.communicate(), timeout=timeout or self._timeout
            )
        except asyncio.TimeoutError:
            proc.kill()
            raise RuntimeError(f"docker command timed out: {' '.join(argv[:3])}")
        return (
            proc.returncode,
            out.decode("utf-8", "replace"),
            err.decode("utf-8", "replace"),
        )

    async def _stream(
        self,
        argv: list,
        on_log: Optional[LogCallback],
        *,
        timeout: Optional[int] = None,
    ) -> int:
        proc = await asyncio.create_subprocess_exec(
            *argv, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT
        )

        async def pump():
            assert proc.stdout is not None
            async for raw in proc.stdout:
                line = raw.decode("utf-8", "replace").rstrip("\n")
                if on_log is not None:
                    await on_log(line)
            return await proc.wait()

        try:
            return await asyncio.wait_for(pump(), timeout=timeout or self._timeout)
        except asyncio.TimeoutError:
            proc.kill()
            raise RuntimeError(f"docker stream timed out: {' '.join(argv[:3])}")

    async def _check(self, argv: list) -> str:
        code, out, err = await self._run(argv)
        if code != 0:
            raise RuntimeError(
                f"{' '.join(argv[:2])} failed ({code}): {(err or out).strip()[:500]}"
            )
        return out.strip()

    async def build(
        self, *, context_dir, dockerfile_text, tags, on_log=None, network=""
    ) -> BuildResult:
        ctx = Path(context_dir)
        dockerfile_path = ctx / "Dockerfile.devplace"
        iid_path = await asyncio.to_thread(
            self._prepare_context, dockerfile_path, dockerfile_text
        )
        argv = build_image_argv(str(ctx), str(dockerfile_path), tags, iid_path, network)
        try:
            code = await self._stream(argv, on_log, timeout=self._build_timeout)
            if code != 0:
                return BuildResult(success=False, error=f"docker build exited {code}")
            image_id = (await asyncio.to_thread(Path(iid_path).read_text)).strip()
            return BuildResult(success=True, image_id=image_id)
        finally:
            await asyncio.to_thread(self._cleanup_context, iid_path, dockerfile_path)

    @staticmethod
    def _prepare_context(dockerfile_path: Path, dockerfile_text: str) -> str:
        dockerfile_path.write_text(dockerfile_text, encoding="utf-8")
        iid = tempfile.NamedTemporaryFile(suffix=".iid", delete=False)
        iid.close()
        return iid.name

    @staticmethod
    def _cleanup_context(iid_path: str, dockerfile_path: Path) -> None:
        Path(iid_path).unlink(missing_ok=True)
        dockerfile_path.unlink(missing_ok=True)

    async def run(self, spec: RunSpec) -> str:
        return await self._check(build_run_argv(spec))

    async def start(self, cid: str) -> None:
        await self._check([self._binary, "start", cid])

    async def stop(self, cid: str, timeout: int = 10) -> None:
        await self._check([self._binary, "stop", "-t", str(timeout), cid])

    async def restart(self, cid: str) -> None:
        await self._check([self._binary, "restart", cid])

    async def pause(self, cid: str) -> None:
        await self._check([self._binary, "pause", cid])

    async def unpause(self, cid: str) -> None:
        await self._check([self._binary, "unpause", cid])

    async def rm(self, cid: str, force: bool = False) -> None:
        argv = [self._binary, "rm"]
        if force:
            argv.append("-f")
        argv.append(cid)
        await self._check(argv)

    async def exec(self, cid, cmd, *, tty=False, on_log=None) -> ExecResult:
        argv = [self._binary, "exec", "-w", WORKSPACE_MOUNT]
        if tty:
            argv.append("-t")
        argv.append(cid)
        argv += list(cmd)
        if on_log is not None:
            code = await self._stream(argv, on_log)
            return ExecResult(exit_code=code)
        code, out, err = await self._run(argv)
        return ExecResult(exit_code=code or 0, output=(out + err))

    async def logs(self, cid, *, follow=False, tail=200, on_log=None) -> None:
        argv = [self._binary, "logs", "--tail", str(tail)]
        if follow:
            argv.append("--follow")
        argv.append(cid)
        await self._stream(argv, on_log)

    async def stats_once(self, cids: list) -> dict:
        if not cids:
            return {}
        argv = [self._binary, "stats", "--no-stream", "--format", "{{json .}}", *cids]
        code, out, _ = await self._run(argv, timeout=30)
        samples = {}
        if code != 0:
            return samples
        for line in out.splitlines():
            line = line.strip()
            if not line:
                continue
            try:
                row = json.loads(line)
            except ValueError:
                continue
            name = row.get("Name") or row.get("ID") or ""
            samples[name] = parse_stats_line(row)
        return samples

    async def ps(self, *, label_filter: str = "devplace.instance") -> list:
        argv = [
            self._binary,
            "ps",
            "-a",
            "--filter",
            f"label={label_filter}",
            "--format",
            "{{json .}}",
        ]
        code, out, _ = await self._run(argv, timeout=30)
        rows = []
        if code != 0:
            return rows
        for line in out.splitlines():
            line = line.strip()
            if not line:
                continue
            try:
                row = json.loads(line)
            except ValueError:
                continue
            labels = {}
            for pair in (row.get("Labels", "") or "").split(","):
                if "=" in pair:
                    key, value = pair.split("=", 1)
                    labels[key] = value
            status = row.get("Status", "") or ""
            match = _EXIT_RE.search(status)
            rows.append(
                PsRow(
                    container_id=row.get("ID", ""),
                    name=row.get("Names", ""),
                    state=(row.get("State", "") or "").lower(),
                    status=status,
                    exit_code=int(match.group(1)) if match else None,
                    labels=labels,
                )
            )
        return rows

    async def inspect(self, cid: str) -> dict:
        code, out, _ = await self._run([self._binary, "inspect", cid], timeout=30)
        if code != 0:
            return {}
        try:
            data = json.loads(out)
            return data[0] if isinstance(data, list) and data else {}
        except ValueError:
            return {}

    async def image_prune(self) -> None:
        try:
            await self._run(
                [
                    self._binary,
                    "image",
                    "prune",
                    "-f",
                    "--filter",
                    "label=devplace.build",
                ],
                timeout=120,
            )
        except RuntimeError as exc:
            logger.warning("image prune failed: %s", exc)

    async def remove_image(self, ref: str, force: bool = False) -> None:
        argv = [self._binary, "rmi"]
        if force:
            argv.append("-f")
        argv.append(ref)
        code, _, err = await self._run(argv, timeout=60)
        if code != 0 and "No such image" not in (err or ""):
            logger.warning("rmi %s failed: %s", ref, (err or "").strip()[:200])

    async def image_exists(self, ref: str) -> bool:
        code, _, _ = await self._run(
            [self._binary, "image", "inspect", ref], timeout=30
        )
        return code == 0
