"""Resolve verified public-image mirrors without changing Docker's global registry settings."""

import json
import os
from pathlib import Path
import re
import subprocess
import sys
from urllib.parse import urlsplit

from docker_command import safe_line

ROOT = Path(__file__).resolve().parents[1]
DEFAULT_MIRROR_PREFIX = "m.daocloud.io"
MANIFEST_TIMEOUT = 30
PULL_TIMEOUT = 180
PUBLIC_REPOSITORIES = {
    "docker.io/library/python", "docker.io/library/postgres", "docker.io/library/redis",
    "ghcr.io/astral-sh/uv",
}


class MirrorError(RuntimeError):
    pass


def build_images(environment: dict[str, str]) -> dict[str, str]:
    defaults = dict(re.findall(
        r"^ARG (PYTHON_IMAGE|UV_IMAGE)=(.+)$", (ROOT / "aiprovider/Dockerfile").read_text(), re.M
    ))
    return {key: environment.get(key) or value for key, value in defaults.items()}


def mirror_prefix(environment: dict[str, str]) -> str:
    prefix = environment.get("PLANET_DOCKER_MIRROR_PREFIX", DEFAULT_MIRROR_PREFIX).rstrip("/")
    if not re.fullmatch(r"[a-zA-Z0-9.-]+(?::[0-9]+)?(?:/[a-zA-Z0-9._-]+)*", prefix):
        raise MirrorError("备用镜像地址须为 registry 或 registry/path，不得包含凭据或 URL 参数")
    return prefix


def canonical_image(image: str) -> str:
    first = image.split("/")[0]
    if "/" not in image or not ("." in first or ":" in first or first == "localhost"):
        image = "docker.io/" + ("library/" if "/" not in image else "") + image
    if image.startswith("index.docker.io/"):
        image = "docker.io/" + image.removeprefix("index.docker.io/")
    if image.startswith("docker.io/") and image.count("/") == 1:
        image = "docker.io/library/" + image.removeprefix("docker.io/")
    return image


def mirror_image(image: str, environment: dict[str, str]) -> str:
    canonical = canonical_image(image)
    repository = canonical.split("@")[0].rsplit("/", 1)
    repository = repository[0] + "/" + repository[1].split(":")[0]
    if repository not in PUBLIC_REPOSITORIES:
        raise MirrorError("仅为 Planet 的公开基础镜像回退；自定义或私有镜像须配置自身仓库")
    return mirror_prefix(environment) + "/" + canonical


def mirror_registry_urls(environment: dict[str, str]) -> list[str]:
    for image in build_images(environment).values():
        mirror_image(image, environment)
    return ["https://" + urlsplit("https://" + mirror_prefix(environment)).netloc + "/v2/"]


def run(command: list[str], timeout: int = MANIFEST_TIMEOUT) -> str:
    result = subprocess.run(command, capture_output=True, text=True, timeout=timeout, check=False)
    if result.returncode:
        detail = safe_line(result.stderr.strip() or result.stdout.strip())
        raise MirrorError(detail[-600:] or "备用镜像命令失败")
    return result.stdout


def resolve_image(image: str, environment: dict[str, str]) -> str:
    mirrored = mirror_image(image, environment)
    manifest = json.loads(run([
        "docker", "buildx", "imagetools", "inspect", mirrored, "--format", "{{json .Manifest}}"
    ]))
    digest = manifest.get("digest") if isinstance(manifest, dict) else None
    if not isinstance(digest, str) or not re.fullmatch(r"sha256:[0-9a-f]{64}", digest):
        raise MirrorError("备用镜像未返回有效的 manifest digest")
    if "@" in image and digest != image.rsplit("@", 1)[1]:
        raise MirrorError("备用镜像 digest 与指定镜像不一致")
    repository = mirrored.split("@")[0]
    if ":" in repository.rsplit("/", 1)[1]:
        repository = repository.rsplit(":", 1)[0]
    return repository + "@" + digest


def pull_services(services: list[str], environment: dict[str, str]) -> None:
    if not services or any(service not in {"postgres", "redis"} for service in services):
        raise MirrorError("备用拉取只处理本次请求的 PostgreSQL / Redis 服务")
    configuration = json.loads(run(["docker", "compose", "config", "--format", "json"]))
    for service in services:
        image = configuration["services"][service]["image"]
        local = subprocess.run(
            ["docker", "image", "inspect", image], capture_output=True, timeout=MANIFEST_TIMEOUT
        )
        if local.returncode == 0:
            continue
        resolved = resolve_image(image, environment)
        print(f"备用源拉取 {service}", flush=True)
        # Docker verifies layer digests while pulling the immutable manifest reference.
        run(["docker", "pull", resolved], PULL_TIMEOUT)
        run(["docker", "tag", resolved, image])


def main() -> None:
    environment = dict(os.environ)
    if sys.argv[1:] == ["build"]:
        resolved = {key: resolve_image(value, environment)
                    for key, value in build_images(environment).items()}
        for key, value in resolved.items():
            print(f"{key}={value}")
    elif sys.argv[1:2] == ["pull"]:
        pull_services(sys.argv[2:], environment)
    else:
        raise MirrorError("未知备用镜像操作")


if __name__ == "__main__":
    try:
        main()
    except (MirrorError, OSError, ValueError, KeyError, subprocess.TimeoutExpired) as error:
        detail = str(error) if isinstance(error, MirrorError) else "镜像验证或拉取未完成"
        sys.stderr.write("PLANET_DOCKER_MIRROR_FAILED: " + detail + "\n")
        sys.exit(1)