#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
腾讯云 COS 上传 —— 纯标准库, q-sign(SHA1) 签名(与官方 qcloud_cos CosS3Auth 一致)
把本地产品图传到 COS, 返回可直接喂给 autodl ref_image_0 的公网 URL。
配置: 同目录 cos_upload.json (secret_id / secret_key / bucket / region)
用法: python3 cos_upload.py 本地图.png [更多...]  -> 每行输出一个公网 URL
"""
import datetime, hashlib, hmac, json, os, sys, time, urllib.request, urllib.error
from urllib.parse import quote

CONFIG = os.path.join(os.path.dirname(os.path.abspath(__file__)), "cos_upload.json")
DEFAULT_CT = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp",".mp3": "audio/mpeg", ".wav": "audio/wav", ".m4a": "audio/mp4", ".aac": "audio/aac", ".flac": "audio/flac", ".ogg": "audio/ogg"}

def load_config():
    if not os.path.exists(CONFIG):
        sys.exit("缺配置: 请创建 cos_upload.json 填入 secret_id / secret_key / bucket / region")
    with open(CONFIG) as f:
        c = json.load(f)
    for k in ("secret_id", "secret_key", "bucket", "region"):
        if not c.get(k):
            sys.exit("配置缺少字段: " + k)
    return c

def _qsign(secret_id, secret_key, method, path, headers, key_time):
    # 与官方 qcloud_cos CosS3Auth 逐行一致
    ch = {}
    for k, v in headers.items():
        k2 = quote(k.encode("utf-8"), "-_.~").lower()
        v2 = quote(v.encode("utf-8"), "-_.~")
        ch[k2] = v2
    header_line = "&".join("%s=%s" % (k, ch[k]) for k in sorted(ch))
    format_str = "%s\n%s\n%s\n%s\n" % (method.lower(), path, "", header_line)
    sha1 = hashlib.sha1(format_str.encode()).hexdigest()
    str_to_sign = "sha1\n%s\n%s\n" % (key_time, sha1)
    sign_key = hmac.new(secret_key.encode(), key_time.encode(), hashlib.sha1).hexdigest()
    sig = hmac.new(sign_key.encode(), str_to_sign.encode(), hashlib.sha1).hexdigest()
    header_list = ";".join(sorted(ch))
    return ("q-sign-algorithm=sha1&q-ak=%s&q-sign-time=%s&q-key-time=%s"
            "&q-header-list=%s&q-url-param-list=&q-signature=%s"
            % (secret_id, key_time, key_time, header_list, sig))

def put_object(cfg, local_path, key):
    data = open(local_path, "rb").read()
    ext = os.path.splitext(local_path)[1].lower()
    content_type = DEFAULT_CT.get(ext, "application/octet-stream")
    host = cfg["bucket"] + ".cos." + cfg["region"] + ".myqcloud.com"
    now = int(time.time())
    key_time = "%d;%d" % (now - 60, now + 10000)
    auth = _qsign(cfg["secret_id"], cfg["secret_key"], "PUT", "/" + key,
                  {"Host": host, "Content-Type": content_type}, key_time)
    url = "https://" + host + "/" + key
    req = urllib.request.Request(url, data=data, method="PUT")
    req.add_header("Host", host)
    req.add_header("Content-Type", content_type)
    req.add_header("Authorization", auth)
    try:
        with urllib.request.urlopen(req, timeout=120) as r:
            return r.status
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8", "ignore")
        raise RuntimeError("COS 上传失败 HTTP %s: %s" % (e.code, body[:500]))

def main():
    if len(sys.argv) < 2:
        sys.exit(__doc__)
    cfg = load_config()
    for p in sys.argv[1:]:
        if not os.path.exists(p):
            print("[skip] 文件不存在: " + p, file=sys.stderr)
            continue
        base = os.path.basename(p)
        t = time.strftime("%Y%m%d")
        import uuid
        key = "douyin/" + t + "/" + uuid.uuid4().hex[:8] + "_" + base
        put_object(cfg, p, key)
        print("https://" + cfg["bucket"] + ".cos." + cfg["region"] + ".myqcloud.com/" + key)

if __name__ == "__main__":
    main()