#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
打包 + 持续监听 + 推送后真实校验（一体化）。
- 读取扁平 credentials.conf
- jar 重命名为 start.jar 匹配根 Dockerfile
- SSH 上传 -> docker build -> 登录张家口 ACR -> docker push
- 推送后【强制真实拉取校验】：先 docker rmi 清掉本地缓存，再 docker pull（真正从 ACR 下载），
  然后容器内 sha256sum + jar -tf 校验包完整，与本地 jar 比对。
- 校验带重试（最多 3 次，应对 ACR 最终一致性），全部通过才输出最终可用地址。

用法:
  python -u build_verify.py --jar <jar> --app jjb-saas-safety-eval --version 78
"""
import argparse
import hashlib
import os
import sys
import time
from datetime import datetime

import paramiko

# Windows GBK 控制台无法编码 emoji，强制 stdout/stderr 为 UTF-8，避免结尾打印 PASS/FAIL 时崩溃
for _stream in (sys.stdout, sys.stderr):
    if hasattr(_stream, "reconfigure"):
        _stream.reconfigure(encoding="utf-8", errors="replace")

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
CONF = os.path.join(SCRIPT_DIR, "credentials.conf")
ROOT_DOCKERFILE = r"E:/works/projects/safety-eval-service/Dockerfile"
REMOTE_JAR_NAME = "start.jar"


def load_conf(path):
    conf = {}
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            s = line.strip()
            if not s or s.startswith("#") or "=" not in s:
                continue
            k, v = s.split("=", 1)
            conf[k.strip()] = v.strip()
    return conf


def ssh_exec(client, cmd, timeout=600):
    print(f"> {cmd[:160]}")
    stdin, stdout, stderr = client.exec_command(cmd, timeout=timeout)
    out = stdout.read().decode("utf-8", "replace")
    err = stderr.read().decode("utf-8", "replace")
    rc = stdout.channel.recv_exit_status()
    if out.strip():
        print(out.rstrip())
    if err.strip() and rc != 0:
        print("STDERR:", err.rstrip())
    print(f"[exit={rc}]\n")
    return rc, out, err


def local_sha256(p):
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for chunk in iter(lambda: f.read(8 * 1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def verify_image(client, image, local_sha, local_size, app, attempts=3):
    """强制真实校验：rmi 本地缓存 -> pull(真实从ACR下载) -> run 容器内 sha256 + jar -tf。
    返回 (ok: bool, detail: str)。"""
    for i in range(1, attempts + 1):
        print(f"\n[校验 尝试 {i}/{attempts}] 强制真实拉取并校验")
        # 清掉本地缓存，确保是从 ACR 真实下载
        ssh_exec(client, f"docker rmi -f '{image}' >/dev/null 2>&1; true")
        rc_pull, out_pull, _ = ssh_exec(client, f"docker pull '{image}'", timeout=300)
        if rc_pull != 0:
            print(f"  pull 失败，{i<attempts and '重试...' or '放弃'}")
            time.sleep(8)
            continue
        rc, out, _ = ssh_exec(
            client,
            f"docker run --rm --entrypoint sh '{image}' -c "
            f"'ls -l /opt/app.jar; sha256sum /opt/app.jar; "
            f"jar -tf /opt/app.jar >/dev/null 2>&1 && echo JAR_OK || echo JAR_BAD'",
            timeout=180,
        )
        remote_sha = None
        remote_size = None
        for ln in out.splitlines():
            ln = ln.strip()
            parts = ln.split()
            if len(parts) == 2 and len(parts[0]) == 64 and set(parts[0]) <= set("0123456789abcdef"):
                remote_sha = parts[0]
            if ln.endswith("app.jar") and "/" in ln:
                try:
                    remote_size = int(ln.split()[4])
                except Exception:
                    pass
        ok = (
            remote_sha == local_sha
            and remote_size == local_size
            and "JAR_OK" in out
        )
        print(f"  sha match: {'YES' if remote_sha==local_sha else 'NO'} ({remote_sha})")
        print(f"  size match: {'YES' if remote_size==local_size else 'NO'} ({remote_size})")
        print(f"  jar -tf  : {'OK' if 'JAR_OK' in out else 'BAD'}")
        if ok:
            return True, out_pull
        print(f"  校验未通过，{i<attempts and '重试...' or '放弃'}")
        time.sleep(8)
    return False, ""


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--jar", required=True)
    p.add_argument("--app", required=True)
    p.add_argument("--version", required=True)
    p.add_argument("--dockerfile", default=ROOT_DOCKERFILE)
    args = p.parse_args()

    c = load_conf(CONF)
    MASTER = c["MASTER_HOST"]
    SSH_USER = c["SSH_USER"]
    SSH_PASSWD = c["SSH_PASSWD"]
    ZJK_REGISTRY = c["ZJK_REGISTRY"]
    ZJK_REPO = c["ZJK_REPO"]
    ZJK_USER = c["ZJK_USER"]
    ZJK_PASS = c["ZJK_PASS"]
    HZ_REGISTRY = c["HZ_REGISTRY"]
    HZ_USER = c["HZ_USER"]
    HZ_PASS = c["HZ_PASS"]

    date_tag = datetime.now().strftime("%Y%m%d")
    tag = f"{args.app}-ota-{date_tag}-{args.version}"
    full_image = f"{ZJK_REGISTRY}/{ZJK_REPO}:{tag}"

    jar_local = args.jar
    if not os.path.exists(jar_local):
        print("ERROR jar not found:", jar_local); sys.exit(1)
    local_size = os.path.getsize(jar_local)
    local_sha = local_sha256(jar_local)

    build_dir = f"/tmp/docker-build-{args.app}"
    print("=" * 64)
    print("  GBS build & push -> 张家口 ACR  (build_verify 一体化)")
    print("=" * 64)
    print(f"  JAR : {jar_local}  ({local_size/1024/1024:.1f} MB)")
    print(f"  SHA : {local_sha}")
    print(f"  TAG : {tag}")
    print(f"  IMG : {full_image}")
    print(f"  HOST: {SSH_USER}@{MASTER}")
    print("=" * 64)

    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    client.connect(MASTER, username=SSH_USER, password=SSH_PASSWD,
                   timeout=30, look_for_keys=False, allow_agent=False,
                   banner_timeout=30, auth_timeout=30)
    print("[1/6] SSH connected")

    ssh_exec(client, f"rm -rf {build_dir} && mkdir -p {build_dir}/start/target")

    print("[2/6] Upload JAR -> start/target/start.jar")
    sftp = client.open_sftp()
    remote_jar = f"{build_dir}/start/target/{REMOTE_JAR_NAME}"
    start = time.time()
    def cb(trans, tot):
        pct = int(trans * 100 / tot)
        if pct % 20 == 0:
            print(f"    {pct}%  {trans/1024/1024:.1f} MB")
    sftp.put(jar_local, remote_jar, callback=cb, confirm=False)
    sftp.close()
    print(f"    done in {time.time()-start:.1f}s")

    print("[3/6] Upload Dockerfile + 大小校验")
    with open(args.dockerfile, "r", encoding="utf-8") as f:
        df = f.read()
    stdin, stdout, stderr = client.exec_command(f"cat > {build_dir}/Dockerfile")
    stdin.write(df); stdin.channel.shutdown_write(); stdout.read()
    ssh_exec(client, f"ls -l {build_dir}/start/target/ {build_dir}/Dockerfile")
    rc_ls, out_ls, _ = ssh_exec(client, f"stat -c %s {remote_jar}")
    try:
        remote_size = int(out_ls.strip().splitlines()[-1])
    except Exception:
        remote_size = -1
    if remote_size != local_size:
        print(f"ERROR: jar 大小不一致 local={local_size} remote={remote_size}")
        client.close(); sys.exit(1)
    print(f"    jar size OK: {local_size} bytes")

    print("[4/6] Login 杭州 ACR (基础镜像) & docker build")
    ssh_exec(client, f"echo '{HZ_PASS}' | docker login --username={HZ_USER} --password-stdin {HZ_REGISTRY}")
    rc, _, _ = ssh_exec(client, f"cd {build_dir} && docker build -t '{full_image}' .", timeout=300)
    if rc != 0:
        print("ERROR: docker build failed"); client.close(); sys.exit(1)

    print("[5/6] Login 张家口 ACR & docker push")
    ssh_exec(client, f"echo '{ZJK_PASS}' | docker login --username={ZJK_USER} --password-stdin {ZJK_REGISTRY}")
    rc, out_push, _ = ssh_exec(client, f"docker push '{full_image}'", timeout=600)
    if rc != 0:
        print("ERROR: docker push failed"); client.close(); sys.exit(1)
    # 提取 digest
    digest = ""
    for ln in out_push.splitlines():
        if ln.strip().startswith("jjb-") and "digest:" in ln:
            digest = ln.split("digest:")[-1].strip()

    print("[6/6] 推送后强制真实校验（监听+重试）")
    ok, _ = verify_image(client, full_image, local_sha, local_size, args.app, attempts=3)

    ssh_exec(client, f"rm -rf {build_dir}")
    client.close()

    print("\n" + "=" * 64)
    print("  RESULT")
    print("=" * 64)
    if ok:
        print(f"  STATUS : PASS ✅")
        print(f"  IMAGE  : {full_image}")
        if digest:
            print(f"  DIGEST : {digest}")
        print(f"  PULLABLE: YES (强制真实拉取校验通过)")
        print("=" * 64)
        sys.exit(0)
    else:
        print(f"  STATUS : FAIL ❌  镜像未通过真实校验，请勿使用！")
        print(f"  IMAGE  : {full_image}")
        print("=" * 64)
        sys.exit(2)


if __name__ == "__main__":
    main()