#!/usr/bin/env python3
"""七牛云对象存储的轻量封装（零 SDK 依赖）。

仅用 Python 标准库（hmac/hashlib/base64/json）+ curl 实现七牛云的上传/
下载/列举。worker 端无需 pip 装任何东西。

七牛上传签名机制：
  PutPolicy = {"scope":"<bucket>:<key>", "deadline":<unix秒>}
  编码后用 AccessKey/SecretKey 做 HMAC-SHA1 签名，得到 uploadToken。

用途（本分布式网格）：
  - 七牛云当"种子库"：每个收敛模型的 .7 大气（393KB）上传为
    <seed_prefix>/<model_name>.7，供异地 worker 在 seed_step 回退时拉取。
  - 仅 master 持有 AK/SK 生成 uptoken；worker 只拿 token 上传（无需密钥）。
  - 大产物（.spec 等 13MB）不上传，留各机本地。

环境变量：
  QINIU_ACCESS_KEY / QINIU_SECRET_KEY / QINIU_BUCKET / QINIU_DOMAIN
  （可被 dist_config.yaml 的 qiniu 段覆盖；master 在派发时把 token 传给 worker）
"""
import base64
import hashlib
import hmac
import json
import os
import subprocess
import time

try:
    import urllib.request as urlreq
except ImportError:  # py2 兜底（不会触发）
    urlreq = None


class QiniuStore:
    """七牛云存储客户端。master 与 worker 共用；worker 只用 download/upload(token)。"""

    def __init__(self, access_key=None, secret_key=None, bucket=None,
                 domain=None, seed_prefix="seeds"):
        self.ak = access_key or os.environ.get("QINIU_ACCESS_KEY", "")
        self.sk = secret_key or os.environ.get("QINIU_SECRET_KEY", "")
        self.bucket = bucket or os.environ.get("QINIU_BUCKET", "")
        # 绑定域名（如 http:// 或 https:// + 域名，无尾斜杠）。公开读下载用。
        self.domain = (domain or os.environ.get("QINIU_DOMAIN", "")).rstrip("/")
        self.seed_prefix = seed_prefix

    # ---- base64 url-safe（七牛要求）----
    @staticmethod
    def _b64u(data):
        """bytes -> url-safe base64 字符串（去 padding 也可，七牛两边都接受）。"""
        if isinstance(data, str):
            data = data.encode("utf-8")
        return base64.urlsafe_b64encode(data).decode("ascii")

    @staticmethod
    def _b64u_json(obj):
        return QiniuStore._b64u(json.dumps(obj, separators=(",", ":")))

    def _sign(self, data):
        """HMAC-SHA1(data, sk) -> url-safe base64。"""
        if isinstance(data, str):
            data = data.encode("utf-8")
        return self._b64u(hmac.new(self.sk.encode("utf-8"), data,
                                   hashlib.sha1).digest())

    # ---- uptoken 生成（master 持有 AK/SK 时调用）----
    def gen_uptoken(self, key, expires=3600):
        """为指定 key 生成上传 token。

        key: 对象名（不含 bucket 前缀），如 "seeds/t30000_g5.0_he0_c-1_n-1_o-1.7"
        返回 uploadToken 字符串，交给 worker 用 curl 上传。
        """
        scope = "{}:{}".format(self.bucket, key) if key else self.bucket
        policy = {"scope": scope, "deadline": int(time.time()) + expires}
        encoded = self._b64u_json(policy)
        sign = self._sign(encoded)
        return "{}:{}".format(self.ak, sign) + ":" + encoded

    def gen_private_url(self, key, expires=3600):
        """生成私有空间的临时下载 URL（带 e/token 签名）。

        若 bucket 为公开读，直接用 domain/key 即可，不必调这个。
        """
        if not self.domain:
            raise ValueError("QINIU_DOMAIN 未配置")
        url = "{}/{}".format(self.domain, key)
        e = int(time.time()) + expires
        to_sign = "{}?e={}".format(url, e)
        token = "{}:{}".format(self.ak, self._sign(to_sign))
        return "{}?e={}&token={}".format(url, e, token)

    def public_url(self, key):
        """公开空间的下载 URL（无签名）。"""
        if not self.domain:
            raise ValueError("QINIU_DOMAIN 未配置")
        return "{}/{}".format(self.domain, key)

    # ---- 上传 / 下载（通过 curl，worker 无需 pip）----
    def upload(self, local_path, key, uptoken=None):
        """上传本地文件到七牛。uptoken 可由 master 预生成传入。

        用 curl 表单 POST（七牛上传接口）：
          POST http://upload.qiniup.com  multipart/form-data
          field "token"=uptoken, "key"=key, "file"=@local_path
        返回 (ok:bool, resp_text)。
        """
        if uptoken is None:
            uptoken = self.gen_uptoken(key)
        # 七牛上传域名：华东 upload.qiniup.com，其它区域用对应域名。
        upload_host = os.environ.get("QINIU_UPLOAD_HOST", "upload.qiniup.com")
        url = "http://{}/".format(upload_host)
        cmd = ["curl", "-sS", "-m", "300", "-X", "POST", url,
               "-F", "token={}".format(uptoken),
               "-F", "key={}".format(key),
               "-F", "file=@{}".format(local_path)]
        try:
            out = subprocess.run(cmd, capture_output=True, text=True,
                                 timeout=320)
        except subprocess.TimeoutExpired:
            return False, "upload timeout (>300s)"
        if out.returncode != 0:
            return False, out.stderr.strip() or out.stdout.strip()
        try:
            resp = json.loads(out.stdout)
        except ValueError:
            return False, out.stdout.strip()
        # 成功响应含 "key" 与 "hash"；失败含 "error"
        if "error" in resp:
            return False, resp["error"]
        return True, resp.get("key", "")

    def download(self, key, local_path, timeout=120):
        """下载对象到本地。公开空间用 public_url，私有空间用 private_url。

        返回 True/False。
        """
        if self.ak and self.sk:
            url = self.gen_private_url(key)
        else:
            url = self.public_url(key)
        cmd = ["curl", "-sS", "-m", str(timeout), "-f", "-o", local_path, url]
        try:
            r = subprocess.run(cmd, capture_output=True, timeout=timeout + 10)
            return r.returncode == 0
        except subprocess.TimeoutExpired:
            return False

    # ---- 种子库语义封装 ----
    def seed_key(self, model_name):
        """规范种子对象名：seeds/<model>.7"""
        return "{}/{}.7".format(self.seed_prefix, model_name)

    def upload_seed(self, local_path, model_name, uptoken=None):
        """上传某模型的最终大气 .7 作为种子。返回 (ok, msg)。"""
        return self.upload(local_path, self.seed_key(model_name), uptoken)

    def download_seed(self, model_name, local_path):
        """下载某模型的种子 .7 到 local_path。返回 True/False。"""
        return self.download(self.seed_key(model_name), local_path)

    # ---- 列举（用于 find_seed 跨网查最近邻种子）----
    def list_seeds(self):
        """返回七牛种子库中所有 <model_name>.7 的模型名列表。

        调用七牛 RS API: GET /list?bucket=<b>&prefix=<seed_prefix>
        （需 AK/SK 签名）。返回 [] 当未配置或失败。
        结果缓存到本地文件避免每次都拉（见 _seeds_cache_path）。
        """
        if not (self.ak and self.sk and self.bucket):
            return []
        cache = self._seeds_cache_path()
        # 缓存 5 分钟
        if os.path.isfile(cache) and time.time() - os.path.getmtime(cache) < 300:
            try:
                return json.load(open(cache))
            except Exception:
                pass
        names = self._fetch_seed_list()
        try:
            with open(cache, "w") as f:
                json.dump(names, f)
        except Exception:
            pass
        return names

    def _seeds_cache_path(self):
        return os.path.join(os.path.dirname(os.path.abspath(__file__)),
                            "..", ".qiniu_seeds_cache.json")

    def _fetch_seed_list(self):
        """分页拉取七牛种子列表，返回 model_name 列表。"""
        rs_host = os.environ.get("QINIU_RS_HOST", "http://rsf.qiniuapi.com")
        names = []
        marker = ""
        prefix = self.seed_prefix + "/"
        while True:
            path = "/list?bucket={}&prefix={}".format(self.bucket, prefix)
            if marker:
                path += "&marker={}".format(marker)
            # 签名：AccessToken = AK:urlsafe_b64(hmac(sk, path)):<path>
            encoded_path = path  # 七牛签名用原始 path
            sign = self._sign(encoded_path)
            access_token = "{}:{}".format(self.ak, sign) + encoded_path
            url = rs_host + path
            cmd = ["curl", "-sS", "-m", "60",
                   "-H", "Authorization: QBox " + access_token, url]
            try:
                r = subprocess.run(cmd, capture_output=True, text=True,
                                   timeout=70)
                if r.returncode != 0:
                    break
                data = json.loads(r.stdout)
            except (ValueError, subprocess.TimeoutExpired):
                break
            for item in data.get("items", []):
                k = item.get("key", "")
                # seeds/<model>.7 -> <model>
                if k.startswith(prefix) and k.endswith(".7"):
                    names.append(k[len(prefix):-2])
            marker = data.get("marker", "")
            if not marker:
                break
        return names


def load_from_dist_config(cfg):
    """从 dist_config.yaml 的 qiniu 段构造 QiniuStore。"""
    q = (cfg or {}).get("qiniu", {}) if isinstance(cfg, dict) else {}
    return QiniuStore(
        access_key=q.get("access_key"),
        secret_key=q.get("secret_key"),
        bucket=q.get("bucket"),
        domain=q.get("domain"),
        seed_prefix=q.get("seed_prefix", "seeds"),
    )


if __name__ == "__main__":
    # 自检：打印 token（需配置 AK/SK/Bucket）
    import argparse
    ap = argparse.ArgumentParser(description="qiniu store 自检")
    ap.add_argument("--gen-token", help="为某 key 生成 uptoken")
    ap.add_argument("--list", action="store_true", help="列出种子")
    ap.add_argument("--config", help="dist_config.yaml 路径")
    args = ap.parse_args()
    if args.config:
        try:
            import yaml
            cfg = yaml.safe_load(open(args.config))
        except Exception:
            # 兜底：不装 yaml 时只从环境变量读
            cfg = {"qiniu": {}}
        store = load_from_dist_config(cfg)
    else:
        store = QiniuStore()
    if args.gen_token:
        print(store.gen_uptoken(args.gen_token))
    elif args.list:
        for n in store.list_seeds():
            print(n)
    else:
        print("AK set:", bool(store.ak), "Bucket:", store.bucket,
              "Domain:", store.domain or "(未配置)")