#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
修正版：根据扁平 credentials.conf 打包并推送 safety-eval-service 镜像到张家口 ACR。
- 读取扁平 key=value 格式（MASTER_HOST / SSH_USER / SSH_PASSWD / ZJK_* / HZ_*）
- 将 jar 重命名为 start.jar 放入 build context 的 start/target/（匹配根 Dockerfile 的 COPY）
- 构建前登录杭州 ACR 以拉取基础镜像 pub/jdk:1.8.0_202
- 构建后登录张家口 ACR 并推送
- 推送后执行 docker pull + 容器内 jar 校验，验证可拉取 & 包正常

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

import paramiko

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"  # 根 Dockerfile COPY 期望的名字


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[:200]}")
    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 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)

    build_dir = f"/tmp/docker-build-{args.app}"
    print("=" * 64)
    print("  GBS build & push  ->  张家口 ACR (修正版)")
    print("=" * 64)
    print(f"  JAR : {jar_local}  ({os.path.getsize(jar_local)/1024/1024:.1f} MB)")
    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")

    # confirm=False: 跳过 put 结束时的 stat 确认（远端 OpenSSH sftp-server 对刚写完的文件 stat 会返回 ENOENT）
    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")

    # 大小校验：确认远端 jar 与本地一致（绕过 sftp.stat 的已知问题，用 ssh ls）
    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
    local_size = os.path.getsize(jar_local)
    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, _, _ = ssh_exec(client, f"docker push '{full_image}'", timeout=600)
    if rc != 0:
        print("ERROR: docker push failed")
        client.close()
        sys.exit(1)

    print("[6/6] 验证：可拉取 + 包正常")
    rc_pull, _, _ = ssh_exec(client, f"docker pull '{full_image}'", timeout=300)
    ssh_exec(
        client,
        f"docker run --rm --entrypoint sh '{full_image}' -c "
        f"'echo JAVA:; java -version 2>&1; echo JAR:; ls -l /opt/app.jar; "
        f"unzip -t /opt/app.jar >/dev/null 2>&1 && echo JAR_OK || echo JAR_BAD'",
        timeout=120,
    )

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

    print("\n" + "=" * 64)
    print("  DONE")
    print("=" * 64)
    print(f"  Image: {full_image}")
    print(f"  Pullable: {'YES (docker pull exit=0)' if rc_pull == 0 else 'NO'}")
    print("=" * 64)


if __name__ == "__main__":
    main()