#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
check_images.py — 检查 deploy.stack 常用 Docker 镜像最新版本

数据源：
  - Docker Hub  v2 API:  https://hub.docker.com/v2/repositories/<ns>/<repo>/tags/
  - 通用 registry v2:    https://<host>/v2/<repo>/tags/list

用法：
  ./check_images.py                          # 全量
  ./check_images.py gitea nginx postgres     # 名称包含关键字的镜像
  ./check_images.py -j 20                    # 并发数（默认 10）
  ./check_images.py --json                   # JSON 输出
  ./check_images.py --pre                    # 含 rc/beta/dev 等预发布
  ./check_images.py --outdated               # 只显示可升级的
  ./check_images.py --proxy                  # source ./proxy.sh 后开启代理
  ./check_images.py --timeout 20             # 单镜像超时（秒）
  ./check_images.py --transport tinyfish     # 强制 tinyfish（auto|curl|tinyfish）
  ./check_images.py -h | --help              # 帮助

退出码：
  0 = 全部最新或仅有 unknown；1 = 存在可升级；2 = 致命错误

注意：
  - 私有仓库（hub.tp229.com:3500、hub.wesais.cn、hub.6t7.net 等）大多不开
    匿名 tags API，这些会标 "unknown"，不会算作失败
  - Docker Hub 限流：匿名 100 req/IP/6h；并发 > 20 易触发 429
  - 默认排除预发布标记：-rc / -beta / -alpha / -dev / -pre / -edge /
    nightly / experimental / testing / -ea
  - 默认排除平台变体：nanoserver / windowsservercore / wincore /
    oraclelinux / oracle / pc / centos / rockylinux / alma /
    enterprise / fips / builder
  - 同版本（忽略后缀）视为已最新：1.36.0 ≡ 1.36.0-alpine

依赖：python3 标准库（urllib、json、subprocess、concurrent.futures）
可选：tinyfish CLI（auto/tinyfish 模式下使用，作为 curl 的回退）
"""

import argparse
import json
import os
import pathlib
import re
import subprocess
import sys
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass

# ===================== ANSI 颜色 =====================
class C:
    R   = "\033[0m"
    B   = "\033[1m"
    DIM = "\033[2m"
    RED = "\033[31m"
    GRN = "\033[32m"
    YEL = "\033[33m"
    CYN = "\033[36m"
    BC  = "\033[1;36m"


# ===================== 镜像清单发现 =====================
# 不再手工维护 IMAGES 列表，改从 <service>/env.cfg.example 自动扫描
# 扫描规则：
#   1. 递归找 **/env.cfg.example
#   2. 提取所有以 IMAGE* 开头的变量（如 IMAGE_TAG、IMAGE_NAME、IMAGE_TAG_V18、
#      IMAGE_TAG_REDISV8、Cadvisor_Image 等）；跳过纯版本号变量 IMAGE_TAG_VER
#   3. 解析 ${VAR} 变量引用（递归最多 5 层防环）
#   4. 只保留看起来像 image:tag 格式的值（含冒号、名字部分含 . / - _）
#   5. 按 (image, tag) 去重，保留首次出现位置
#
# 例外（无 env.cfg.example 或不走 IMAGE_TAG 约定的服务）：
#   - harbor: 官方安装器生成的 compose.yaml，不手工改
#   - mixapi/newapi: 在 env.cfg.example 中用 IMAGE_TAG 但版本号频繁迭代；
#     默认走自动扫描，同一镜像会同时列出 IMAGE_TAG 和 IMAGE_TAG_Vxx
# 手动补漏：MANUAL_IMAGES 在自动扫描结果上合并（一般留空）
MANUAL_IMAGES: list[tuple[str, str, str]] = [
    # ("some/image", "1.0.0", "手动补充：为什么自动扫不到的原因"),
]

_VAR_REF_RE = re.compile(r"\$\{(\w+)\}")
_IMAGE_NAME_OK = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9._/-]*$")


def parse_env_file(text: str) -> dict[str, str]:
    """解析 KEY=VALUE 形式的配置，跳过空行/注释/引号包裹。"""
    result = {}
    for line in text.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if "=" not in line:
            continue
        k, _, v = line.partition("=")
        k, v = k.strip(), v.strip()
        if len(v) >= 2 and v[0] == v[-1] and v[0] in ("'", '"'):
            v = v[1:-1]
        result[k] = v
    return result


def resolve_variables(value: str, env: dict[str, str], _depth: int = 0) -> str:
    """递归替换 ${VAR}；最多 5 层防循环引用。"""
    if _depth >= 5 or "${" not in value:
        return value

    def repl(m: re.Match) -> str:
        v = env.get(m.group(1), m.group(0))
        return resolve_variables(v, env, _depth + 1) if "${" in v else v

    return _VAR_REF_RE.sub(repl, value)


def split_image_tag(value: str) -> tuple[str, str]:
    """拆 'image:tag'。正确处理私有仓库含端口的场景。"""
    if ":" not in value:
        return value, ""
    # 冒号在最后一个斜杠之后 → tag 分隔符
    last_colon = value.rfind(":")
    last_slash = value.rfind("/")
    if last_colon <= last_slash:
        # 形如 redis （无 tag）或 host:5000 （无 repo）→ 不能用
        return value, ""
    return value[:last_colon], value[last_colon + 1:]


def looks_like_image_ref(value: str) -> bool:
    if ":" not in value:
        return False
    name, tag = split_image_tag(value)
    return bool(name and tag and _IMAGE_NAME_OK.match(name) and _IMAGE_NAME_OK.match(tag))


def discover_images(repo_root: str = ".") -> list[tuple[str, str, str]]:
    """扫描所有 env.cfg.example，生成 (image, current_tag, note) 列表。"""
    root = pathlib.Path(repo_root)
    found: list[tuple[str, str, str]] = []

    # 完全跳过的顶层目录（非服务）
    SKIP_TOPDIRS = {
        "shell", "etc", "config", "crontab", "apt.list",
        "i2c.py", ".trae", "$user",
    }
    # 需要下钻一层的容器目录（使用其子目录名作为 service）
    NESTED_DIRS = {"dbSer", "WireGuardVPN", "base"}

    for env_file in sorted(root.rglob("env.cfg.example")):
        rel = env_file.relative_to(root)
        parts = rel.parts
        if not parts or parts[0] in SKIP_TOPDIRS:
            continue

        # builder/ 是开发容器镜像（golang/node/alpine/debian），其版本故意锁定，不扫描
        if parts[0] == "builder":
            continue

        # 决定 service 名称
        if parts[0] in NESTED_DIRS and len(parts) >= 3:
            service = parts[1]
        else:
            service = parts[0] if len(parts) == 2 else parts[0]

        try:
            text = env_file.read_text(encoding="utf-8", errors="replace")
        except OSError:
            continue

        env = parse_env_file(text)
        for key in sorted(env.keys()):
            if not re.match(r"^IMAGE", key, re.IGNORECASE):
                continue
            if key.upper() == "IMAGE_TAG_VER":
                continue

            resolved = resolve_variables(env[key], env)
            if not looks_like_image_ref(resolved):
                continue

            image, tag = split_image_tag(resolved)
            note = f"{service}/{key}"
            found.append((image, tag, note))

    found.extend(MANUAL_IMAGES)

    # 去重：同 (image, tag) 只保留首次出现
    seen: set[tuple[str, str]] = set()
    unique: list[tuple[str, str, str]] = []
    for img, tag, note in found:
        key = (img.lower(), tag)
        if key in seen:
            continue
        seen.add(key)
        unique.append((img, tag, note))

    return unique


# ===================== 预发布 / 变体识别 =====================
PRE_MARKERS = {
    "rc", "beta", "alpha", "dev", "pre", "edge",
    "nightly", "experimental", "testing", "ea",
}
BAD_VARIANTS = {
    "nanoserver", "windowsservercore", "wincore",
    "oraclelinux", "oracle",
    "pc", "ea",
    "centos", "rockylinux", "alma",
    "enterprise", "fips", "builder",
}
VERSION_RE = re.compile(r"^v?(\d+(?:\.\d+)*)")


def is_pre_release(name: str) -> bool:
    parts = re.split(r"[.\-]", name.lower())
    parts = [re.sub(r"^\d+|\d+$", "", p) for p in parts]
    return any(p in PRE_MARKERS for p in parts)


def looks_like_version(name: str) -> bool:
    return bool(VERSION_RE.match(name))


def is_good_variant(name: str) -> bool:
    base = re.sub(r"^v?\d+(?:\.\d+)*", "", name.lower()).lstrip("-")
    return not any(bad in base for bad in BAD_VARIANTS)


def version_key(name: str):
    m = VERSION_RE.match(name)
    if not m:
        return (0,)
    return tuple(int(p) for p in m.group(1).split("."))


def pick_latest(data: dict, include_pre: bool = False):
    """从 Docker Hub / registry v2 的 tags 数据中选出"最新稳定 tag"。
    返回 (tag_str, None) 或 (None, err_msg)。"""
    tags = []
    if "results" in data and isinstance(data["results"], list):
        for t in data["results"]:
            if isinstance(t, dict):
                tags.append((t.get("name", ""), t.get("last_updated", "") or ""))
    elif "tags" in data and isinstance(data["tags"], list):
        for n in data["tags"]:
            tags.append((n, ""))
    else:
        if isinstance(data, dict):
            detail = data.get("detail") or data.get("message") or "unknown"
        else:
            detail = "unknown"
        return None, f"detail: {detail}"

    if not tags:
        return None, "no tags"

    # 主键：版本降序；次键：好变体优先；再次：last_updated 降序
    tags.sort(key=lambda x: (
        version_key(x[0]),
        1 if is_good_variant(x[0]) else 0,
        x[1], x[0],
    ), reverse=True)

    fallback_latest = None
    seen = set()

    # 第一遍：版本型 + 非预发布
    for name, _ in tags:
        if not name or name in seen:
            continue
        seen.add(name)
        if name.lower() == "latest":
            fallback_latest = name
            continue
        if not include_pre and is_pre_release(name):
            continue
        if not looks_like_version(name):
            continue
        return name, None

    # 第二遍：任意非预发布
    for name, _ in tags:
        if not name or name in seen:
            continue
        seen.add(name)
        if name.lower() == "latest":
            continue
        if not include_pre and is_pre_release(name):
            continue
        return name, None

    if fallback_latest:
        return fallback_latest, None
    return (tags[0][0], None) if tags else (None, "no usable tags")


# ===================== 网络层 =====================
def resolve_registry(image: str) -> tuple[str, str]:
    """返回 (api_type, url)。api_type: dockerhub | v2"""
    if (":" in image and "/" in image.split(":", 1)[1]) \
       or image.startswith(("hub.", "ghcr.io/", "quay.io/", "gcr.io/")):
        host_port = image.split("/", 1)[0]
        repo_path = image.split("/", 1)[1]
        return ("v2", f"https://{host_port}/v2/{repo_path}/tags/list")
    if "/" in image:
        ns, repo = image.split("/", 1)
    else:
        ns, repo = "library", image
    return ("dockerhub",
            f"https://hub.docker.com/v2/repositories/{ns}/{repo}/tags/?page_size=100&ordering=last_updated")


def fetch_via_curl(url: str, timeout: int) -> str:
    req = urllib.request.Request(url, headers={"User-Agent": "check_images.py/1.0"})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return resp.read().decode("utf-8", errors="replace")


def fetch_via_tinyfish(url: str, timeout: int) -> str:
    if not has_tinyfish():
        raise RuntimeError("tinyfish not installed")
    result = subprocess.run(
        ["tinyfish", "fetch", "content", "get", "--format", "json", url],
        capture_output=True, text=True, timeout=timeout,
    )
    if result.returncode != 0:
        raise RuntimeError(f"tinyfish exit {result.returncode}: {result.stderr[:200]}")
    return result.stdout


def has_tinyfish() -> bool:
    try:
        subprocess.run(["tinyfish", "--version"], capture_output=True, timeout=3, check=False)
        return True
    except (FileNotFoundError, subprocess.TimeoutExpired):
        return False


def fetch_json(url: str, transport: str, timeout: int) -> str:
    if transport == "curl":
        return fetch_via_curl(url, timeout)
    if transport == "tinyfish":
        return fetch_via_tinyfish(url, timeout)
    raise ValueError(f"unknown transport: {transport}")


def probe_transport(url: str, requested: str) -> str:
    """auto 模式：先试 curl，能通就用；不通且装了 tinyfish 就用 tinyfish。"""
    if requested != "auto":
        return requested
    try:
        fetch_via_curl(url, timeout=5)
        return "curl"
    except Exception:
        pass
    if has_tinyfish():
        print(f"{C.YEL}⚠️  curl 无法直连 Docker Hub，回退到 tinyfish{C.R}", file=sys.stderr)
        return "tinyfish"
    print(f"{C.RED}❌ curl 不通且未安装 tinyfish{C.R}", file=sys.stderr)
    return "curl"  # 维持 curl 让每个镜像标 (unreachable)


def unwrap_tinyfish(raw: str):
    """如 raw 是 tinyfish 包装 JSON，则解包返回内层 dict；否则原样解析。"""
    data = json.loads(raw)
    if isinstance(data, dict) and "results" in data and "errors" in data:
        results = data.get("results") or []
        if not results:
            return None, "tinyfish: empty results"
        text_node = (results[0] or {}).get("text") or {}
        inner = None
        for child in text_node.get("children") or []:
            if isinstance(child, dict) and child.get("type") == "code" and "text" in child:
                inner = child["text"]
                break
        if inner is None:
            return None, "tinyfish: no code block"
        try:
            return json.loads(inner), None
        except json.JSONDecodeError as e:
            return None, f"tinyfish inner parse: {e}"
    return data, None


# ===================== 业务逻辑 =====================
@dataclass
class Result:
    image: str
    current: str
    latest: str
    status: str   # up-to-date | outdated | unknown
    note: str = ""


def parse_version_base(tag: str) -> str:
    """1.36.0-alpine → 1.36.0；v1.2.3-rootless → 1.2.3"""
    m = re.match(r"^v?(\d+(?:\.\d+)*)", tag)
    return m.group(1) if m else ""


def check_one(image: str, current: str, note: str, transport: str,
              timeout: int, include_pre: bool) -> Result:
    api_type, url = resolve_registry(image)
    try:
        raw = fetch_json(url, transport, timeout)
    except Exception as e:
        return Result(image, current, "(unreachable)", "unknown", note)

    if not raw.strip():
        return Result(image, current, "(unreachable)", "unknown", note)

    try:
        data = json.loads(raw)
    except json.JSONDecodeError as e:
        return Result(image, current, f"(err: bad json: {e})", "unknown", note)

    # tinyfish 包装解包
    if isinstance(data, dict) and "results" in data and "errors" in data:
        data, err = unwrap_tinyfish(raw)
        if err:
            return Result(image, current, f"(err: {err})", "unknown", note)

    latest, err = pick_latest(data, include_pre)
    if err or latest is None:
        return Result(image, current, f"(err: {err or 'no tag'})", "unknown", note)

    # 状态判定：同版本（忽略后缀）视为已最新
    cur_base = parse_version_base(current)
    lat_base = parse_version_base(latest)
    if (current == "-" or current == "latest"
            or current == latest
            or (cur_base and cur_base == lat_base)):
        status = "up-to-date"
    else:
        status = "outdated"
    return Result(image, current, latest, status, note)


# ===================== 代理 =====================
def setup_proxy(use_proxy: bool):
    """调用 ./proxy.sh 的 proxy_on，把导出的环境变量拉进当前进程。"""
    if not use_proxy:
        return
    proxy_sh = os.path.join(os.path.dirname(os.path.abspath(__file__)), "proxy.sh")
    if not os.path.exists(proxy_sh):
        print(f"{C.YEL}⚠️  未找到 proxy.sh（{proxy_sh}），跳过代理{C.R}", file=sys.stderr)
        return
    try:
        result = subprocess.run(
            ["bash", "-c", f'. "{proxy_sh}" && proxy_on >/dev/null 2>&1 && env'],
            capture_output=True, text=True, timeout=10,
        )
        for line in result.stdout.splitlines():
            if "=" in line:
                k, v = line.split("=", 1)
                os.environ[k] = v
    except Exception as e:
        print(f"{C.YEL}⚠️  proxy_on 失败：{e}{C.R}", file=sys.stderr)


# ===================== 输出 =====================
def display_table(results: list[Result]) -> None:
    up_to_date = outdated = unknown = 0
    for r in results:
        if r.status == "up-to-date":   up_to_date += 1
        elif r.status == "outdated":   outdated   += 1
        else:                          unknown    += 1

    total = len(results)
    box_w = 96
    bar = "─" * box_w

    print(f"{C.CYN}┌{bar}┐{C.R}")
    print(f"{C.CYN}│{C.R}  {C.BC}{'IMAGE':46} {'CURRENT':14} {'LATEST':14} {'STATUS':14}{C.R}  {C.CYN}│{C.R}")
    print(f"{C.CYN}├{bar}┤{C.R}")

    for r in results:
        if r.status == "up-to-date":
            sd = f"{C.GRN}✓ up-to-date{C.R}"
            ld = r.latest
        elif r.status == "outdated":
            sd = f"{C.YEL}↑ outdated{C.R}"
            ld = f"{C.CYN}{r.latest}{C.R}"
        else:
            sd = f"{C.RED}✗ unknown{C.R}"
            ld = f"{C.DIM}{r.latest}{C.R}"

        img = r.image
        if len(img) > 46:
            img = img[:43] + "..."
        print(f"{C.CYN}│{C.R}  {img:46} {r.current:14} {ld:14} {sd:30}  {C.CYN}│{C.R}".rstrip())

    print(f"{C.CYN}└{bar}┘{C.R}")
    print()
    print(f"{C.BC}📊 汇总:{C.R} 总计 {total} | {C.GRN}✓ {up_to_date}{C.R} | "
          f"{C.YEL}↑ {outdated}{C.R} | {C.RED}? {unknown}{C.R}")


def display_json(results: list[Result]) -> None:
    print("[")
    for i, r in enumerate(results):
        sep = "," if i > 0 else ""
        print(f'{sep}\n  {{"image":"{r.image}","current":"{r.current}",'
              f'"latest":"{r.latest}","status":"{r.status}","note":"{r.note}"}}',
              end="")
    print("\n]")


# ===================== 入口 =====================
def main():
    parser = argparse.ArgumentParser(
        prog="check_images.py",
        description="检查 deploy.stack 常用 Docker 镜像的最新版本",
        add_help=False,
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="退出码：0=全最新/未知；1=存在可升级；2=致命错误",
    )
    parser.add_argument("filters", nargs="*", help="按关键字过滤镜像名（OR）")
    parser.add_argument("-j", "--parallel", type=int, default=10, help="并发数（默认 10）")
    parser.add_argument("--json",         action="store_true", help="JSON 输出")
    parser.add_argument("--pre",          action="store_true", help="含 rc/beta/dev 等预发布")
    parser.add_argument("--outdated",     action="store_true", help="只显示可升级的")
    parser.add_argument("--proxy",        action="store_true", help="source ./proxy.sh 后开启代理")
    parser.add_argument("--timeout",      type=int, default=15, help="单镜像超时（秒，默认 15）")
    parser.add_argument("--transport",
                        choices=["auto", "curl", "tinyfish"], default="auto",
                        help="传输方式：auto（默认）/ curl / tinyfish")
    parser.add_argument("--repo", default=".",
                        help="env.cfg.example 扫描根目录（默认当前目录）")
    parser.add_argument("-h", "--help", action="store_true", help="显示帮助")

    args = parser.parse_args()

    if args.help:
        # 打印文件顶部 docstring 的前 35 行作为帮助
        with open(__file__, encoding="utf-8") as f:
            for i, line in enumerate(f):
                if i >= 35: break
                if line.startswith("#!"): continue
                print(line.rstrip())
        sys.exit(0)

    setup_proxy(args.proxy)

    # 从 env.cfg.example 自动扫描镜像清单
    images = discover_images(args.repo)
    print(f"📦 discovered {len(images)} images from {args.repo}/**/env.cfg.example",
          file=sys.stderr)

    # 找一个 Docker Hub 镜像用于探测 transport
    probe_url = None
    for img, _, _ in images:
        api_type, url = resolve_registry(img)
        if api_type == "dockerhub":
            probe_url = url
            break
    transport = probe_transport(probe_url, args.transport) if probe_url else args.transport
    print(f"🔌 transport: {transport}", file=sys.stderr)

    # 过滤
    if args.filters:
        selected = [t for t in images if any(f in t[0] for f in args.filters)]
    else:
        selected = list(images)

    # 并发检查
    results: list[Result] = []
    with ThreadPoolExecutor(max_workers=max(1, args.parallel)) as pool:
        futures = {
            pool.submit(check_one, img, cur, note, transport, args.timeout, args.pre): img
            for img, cur, note in selected
        }
        for fut in as_completed(futures):
            results.append(fut.result())

    # 仅过期
    if args.outdated:
        results = [r for r in results if r.status == "outdated"]

    # 输出
    if args.json:
        display_json(results)
    else:
        display_table(results)

    # 退出码
    outdated_n = sum(1 for r in results if r.status == "outdated")
    sys.exit(1 if outdated_n > 0 else 0)


if __name__ == "__main__":
    main()