"""lumimaker-game-cover-maker 出图脚本：读取 02-封面方案.md 中的方案，调用用户自己的生图 API。

只依赖 Python 3.9+ 标准库。在用户游戏工程根目录执行：

    python <SKILL_DIR>/scripts/generate_image.py --check
    python <SKILL_DIR>/scripts/generate_image.py --plan 1 --dry-run
    python <SKILL_DIR>/scripts/generate_image.py --plan 1 --plan 3
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import hmac
import json
import mimetypes
import os
import re
import sys
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.parse import urlparse
from urllib.request import Request, urlopen

DEFAULT_PLANS_FILE = Path("cover-maker/02-封面方案.md")
DEFAULT_CONFIG_FILE = Path("cover-maker/config.local.json")
DEFAULT_OUT_DIR = Path("cover-maker/images")
DEFAULT_ASPECT = "16:9"
MAX_REFS = 3
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".gif"}

CN_NUMS = "零一二三四五六七八九十"


class CoverError(Exception):
    pass


# ---------- 方案文件解析 ----------

def cn_to_int(text: str) -> int:
    text = text.strip()
    if text.isdigit():
        return int(text)
    if text == "十":
        return 10
    if len(text) == 1 and text in CN_NUMS:
        return CN_NUMS.index(text)
    raise CoverError(f"无法识别的方案编号：{text}")


def int_to_cn(value: int) -> str:
    return CN_NUMS[value] if 0 < value <= 10 else str(value)


def ignored_ref(item: str) -> bool:
    text = item.strip()
    return (
        not text
        or text.startswith("无")
        or text.startswith("未找到")
        or text.startswith("可自行补充")
    )


def split_list(raw: str) -> list[str]:
    raw = raw.strip()
    if ignored_ref(raw):
        return []
    items = [item.strip().strip("`") for item in re.split(r"[,，、;；]", raw)]
    return [item for item in items if not ignored_ref(item)]


def parse_plans_file(path: Path) -> dict[str, Any]:
    if not path.is_file():
        raise CoverError(f"找不到方案文件：{path}。请先完成第 3 步，或用 --plans 指定路径。")
    text = path.read_text(encoding="utf-8-sig")

    aspect_match = re.search(r"\*\*画幅\*\*\s*[：:]\s*(\d+)\s*[:：]\s*(\d+)", text)
    aspect = f"{aspect_match.group(1)}:{aspect_match.group(2)}" if aspect_match else DEFAULT_ASPECT

    headings = list(re.finditer(r"^##\s*方案\s*([一二三四五六七八九十\d]+)\s*[：:]?\s*(.*)$", text, re.M))
    plans: dict[int, dict[str, Any]] = {}
    for index, heading in enumerate(headings):
        end = headings[index + 1].start() if index + 1 < len(headings) else len(text)
        section = text[heading.end():end]
        prompt_match = re.search(r"```prompt[^\n]*\n(.*?)\n```", section, re.S)
        if not prompt_match:
            continue
        prompt = prompt_match.group(1).strip()
        if not prompt or prompt in {"……", "..."}:
            continue
        refs_match = re.search(r"\*\*参考图文件\*\*\s*[：:]\s*(.+)", section)
        number = cn_to_int(heading.group(1))
        plans[number] = {
            "number": number,
            "title": heading.group(2).strip(),
            "prompt": prompt,
            "refs": split_list(refs_match.group(1)) if refs_match else [],
        }
    return {"aspect": aspect, "plans": plans}


# ---------- 配置 ----------

def mask(secret: str) -> str:
    if not secret:
        return "(未填写)"
    if len(secret) <= 8:
        return "****"
    return f"{secret[:4]}****{secret[-4:]}"


def load_config(path: Path) -> dict[str, Any]:
    if not path.is_file():
        raise CoverError(
            f"找不到配置文件：{path}。请把 scripts/config.example.json 复制为 {path} 并填写，"
            "或用 --config 指定路径。"
        )
    try:
        config = json.loads(path.read_text(encoding="utf-8-sig"))
    except json.JSONDecodeError as exc:
        raise CoverError(f"配置文件不是合法的 JSON：{exc}") from exc

    provider = str(config.get("provider") or "").strip()
    if provider not in {"openai", "ark", "tencent-vod"}:
        raise CoverError("配置里的 provider 必须是 openai / ark / tencent-vod 之一。")
    section = dict(config.get(provider) or {})

    if provider in {"openai", "ark"}:
        section["api_key"] = os.environ.get("COVER_API_KEY", "").strip() or str(section.get("api_key") or "").strip()
        if not section["api_key"]:
            raise CoverError(f"{provider} 缺少 api_key（可写在配置里，或设置环境变量 COVER_API_KEY）。")
        if not section.get("model"):
            raise CoverError(f"{provider} 缺少 model。")
        section["base_url"] = str(section.get("base_url") or "").rstrip("/")
    else:
        section["secret_id"] = os.environ.get("TENCENTCLOUD_SECRET_ID", "").strip() or str(section.get("secret_id") or "").strip()
        section["secret_key"] = os.environ.get("TENCENTCLOUD_SECRET_KEY", "").strip() or str(section.get("secret_key") or "").strip()
        missing = [key for key in ("secret_id", "secret_key", "sub_app_id") if not section.get(key)]
        if missing:
            raise CoverError(f"tencent-vod 缺少字段：{', '.join(missing)}")
        section["sub_app_id"] = int(section["sub_app_id"])

    section["provider"] = provider
    return section


# ---------- 参考图 ----------

def is_url(value: str) -> bool:
    return value.startswith(("http://", "https://"))


def resolve_refs(refs: list[str], provider: str) -> list[str]:
    if len(refs) > MAX_REFS:
        print(f"[提示] 参考图超过 {MAX_REFS} 张，只使用前 {MAX_REFS} 张：{refs[:MAX_REFS]}")
        refs = refs[:MAX_REFS]
    resolved: list[str] = []
    for ref in refs:
        if is_url(ref):
            if provider == "openai":
                raise CoverError(f"openai 的参考图只支持本地文件，不支持 URL：{ref}")
            resolved.append(ref)
            continue
        path = Path(ref)
        if path.suffix.lower() not in IMAGE_EXTS:
            raise CoverError(f"参考图只接受图片文件（png/jpg/jpeg/webp/gif）：{ref}")
        if not path.is_file():
            raise CoverError(f"参考图文件不存在：{ref}（路径应相对于工程根目录）。可以修改方案文件，或加 --no-refs。")
        if provider == "tencent-vod":
            raise CoverError(
                f"tencent-vod 的参考图只支持公网 URL，不支持本地文件：{ref}。"
                "请先把图片上传到图床，再用 --ref <URL> 传入，或加 --no-refs。"
            )
        resolved.append(str(path))
    return resolved


def file_to_data_url(path: str) -> str:
    mime = mimetypes.guess_type(path)[0] or "image/png"
    return f"data:{mime};base64,{base64.b64encode(Path(path).read_bytes()).decode('ascii')}"


# ---------- HTTP ----------

def http_json(url: str, body: bytes, headers: dict[str, str], timeout: int = 300) -> dict[str, Any]:
    request = Request(url, data=body, headers=headers, method="POST")
    try:
        with urlopen(request, timeout=timeout) as response:
            return json.loads(response.read().decode("utf-8"))
    except HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="ignore")[:1000]
        hint = {
            400: "请求参数有误：检查模型是否支持当前尺寸（可改配置里的 size_map）、Prompt 是否触发审核。",
            401: "鉴权失败：检查 api_key 是否正确、是否过期。",
            403: "无权限：检查账号是否开通了该模型、余额是否充足。",
            404: "接口不存在：检查 base_url 和 model 是否正确。",
            429: "请求太频繁或额度用尽：稍后再试。",
        }.get(exc.code, "")
        raise CoverError(f"HTTP {exc.code}：{detail}\n{hint}".strip()) from exc
    except URLError as exc:
        raise CoverError(f"网络连接失败：{exc.reason}。检查 base_url 和网络/代理。") from exc


def build_multipart(fields: dict[str, str], files: list[tuple[str, str]]) -> tuple[bytes, str]:
    boundary = f"----CoverMaker{uuid.uuid4().hex}"
    chunks: list[bytes] = []
    for name, value in fields.items():
        chunks.append(f'--{boundary}\r\nContent-Disposition: form-data; name="{name}"\r\n\r\n{value}\r\n'.encode("utf-8"))
    for name, path in files:
        mime = mimetypes.guess_type(path)[0] or "application/octet-stream"
        filename = Path(path).name
        chunks.append(
            f'--{boundary}\r\nContent-Disposition: form-data; name="{name}"; filename="{filename}"\r\n'
            f"Content-Type: {mime}\r\n\r\n".encode("utf-8")
        )
        chunks.append(Path(path).read_bytes())
        chunks.append(b"\r\n")
    chunks.append(f"--{boundary}--\r\n".encode("utf-8"))
    return b"".join(chunks), f"multipart/form-data; boundary={boundary}"


# ---------- 各服务商 ----------

def size_for(config: dict[str, Any], aspect: str) -> str:
    size_map = config.get("size_map") or {}
    if aspect not in size_map:
        raise CoverError(f"配置的 size_map 里没有画幅 {aspect}，请补充，例如 \"{aspect}\": \"1536x1024\"。")
    return str(size_map[aspect])


def openai_request(config: dict[str, Any], prompt: str, aspect: str, refs: list[str]) -> dict[str, Any]:
    fields = {"model": config["model"], "prompt": prompt, "size": size_for(config, aspect), "n": "1"}
    if config.get("quality"):
        fields["quality"] = str(config["quality"])
    headers = {"Authorization": f"Bearer {config['api_key']}"}
    if refs:
        url = f"{config['base_url']}/images/edits"
        files = [("image[]", ref) for ref in refs]
        return {"url": url, "headers": headers, "fields": fields, "files": files}
    url = f"{config['base_url']}/images/generations"
    return {"url": url, "headers": headers, "json": {**fields, "n": 1}}


def ark_request(config: dict[str, Any], prompt: str, aspect: str, refs: list[str]) -> dict[str, Any]:
    body: dict[str, Any] = {
        "model": config["model"],
        "prompt": prompt,
        "size": size_for(config, aspect),
        "response_format": "url",
        "watermark": bool(config.get("watermark", False)),
        "sequential_image_generation": "disabled",
    }
    if refs:
        images = [ref if is_url(ref) else file_to_data_url(ref) for ref in refs]
        body["image"] = images[0] if len(images) == 1 else images
    return {
        "url": f"{config['base_url']}/images/generations",
        "headers": {"Authorization": f"Bearer {config['api_key']}"},
        "json": body,
    }


def send_simple_request(req: dict[str, Any]) -> list[dict[str, str]]:
    headers = dict(req["headers"])
    if "files" in req:
        body, content_type = build_multipart(req["fields"], req["files"])
    else:
        body = json.dumps(req["json"], ensure_ascii=False).encode("utf-8")
        content_type = "application/json"
    headers["Content-Type"] = content_type
    data = http_json(req["url"], body, headers)
    if data.get("error"):
        raise CoverError(f"接口返回错误：{data['error']}")
    results = []
    for item in data.get("data") or []:
        if item.get("b64_json"):
            results.append({"b64": item["b64_json"]})
        elif item.get("url"):
            results.append({"url": item["url"]})
    if not results:
        raise CoverError(f"接口没有返回图片：{json.dumps(data, ensure_ascii=False)[:500]}")
    return results


def tencent_call(config: dict[str, Any], action: str, payload: dict[str, Any]) -> dict[str, Any]:
    host, service, version = "vod.tencentcloudapi.com", "vod", "2018-07-17"
    timestamp = int(time.time())
    date = datetime.fromtimestamp(timestamp, tz=timezone.utc).strftime("%Y-%m-%d")
    body = json.dumps(payload, ensure_ascii=False, separators=(",", ":"))

    canonical_headers = f"content-type:application/json; charset=utf-8\nhost:{host}\n"
    signed_headers = "content-type;host"
    canonical_request = "\n".join(
        ["POST", "/", "", canonical_headers, signed_headers, hashlib.sha256(body.encode("utf-8")).hexdigest()]
    )
    scope = f"{date}/{service}/tc3_request"
    string_to_sign = "\n".join(
        ["TC3-HMAC-SHA256", str(timestamp), scope, hashlib.sha256(canonical_request.encode("utf-8")).hexdigest()]
    )

    def sign(key: bytes, msg: str) -> bytes:
        return hmac.new(key, msg.encode("utf-8"), hashlib.sha256).digest()

    secret_signing = sign(sign(sign(("TC3" + config["secret_key"]).encode("utf-8"), date), service), "tc3_request")
    signature = hmac.new(secret_signing, string_to_sign.encode("utf-8"), hashlib.sha256).hexdigest()
    headers = {
        "Authorization": (
            f"TC3-HMAC-SHA256 Credential={config['secret_id']}/{scope}, "
            f"SignedHeaders={signed_headers}, Signature={signature}"
        ),
        "Content-Type": "application/json; charset=utf-8",
        "Host": host,
        "X-TC-Action": action,
        "X-TC-Timestamp": str(timestamp),
        "X-TC-Version": version,
    }
    data = http_json(f"https://{host}", body.encode("utf-8"), headers)
    response = data.get("Response") or {}
    if "Error" in response:
        error = response["Error"]
        raise CoverError(f"腾讯云错误 {error.get('Code')}：{error.get('Message')}")
    return response


def tencent_payload(config: dict[str, Any], prompt: str, aspect: str, refs: list[str]) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "SubAppId": config["sub_app_id"],
        "ModelName": str(config.get("model_name") or "OG"),
        "ModelVersion": str(config.get("model_version") or "image2_medium"),
        "Prompt": prompt,
        "OutputConfig": {
            "StorageMode": "Temporary",
            "Resolution": str(config.get("resolution") or "1K"),
            "AspectRatio": aspect,
        },
        "SessionId": uuid.uuid4().hex,
    }
    if refs:
        payload["FileInfos"] = [{"Type": "Url", "Url": ref} for ref in refs]
    return payload


def tencent_generate(config: dict[str, Any], payload: dict[str, Any]) -> list[dict[str, str]]:
    task_id = tencent_call(config, "CreateAigcImageTask", payload).get("TaskId")
    if not task_id:
        raise CoverError("腾讯云没有返回 TaskId。")
    print(f"  已提交任务 {task_id}，等待生成……")
    deadline = time.time() + 600
    while time.time() < deadline:
        time.sleep(5)
        response = tencent_call(config, "DescribeTaskDetail", {"TaskId": task_id, "SubAppId": config["sub_app_id"]})
        task = response.get("AigcImageTask") or {}
        status = response.get("Status") or task.get("Status")
        if status == "FINISH":
            if task.get("ErrCode"):
                raise CoverError(f"生成失败：{task.get('ErrCode')} {task.get('Message')}")
            files = (task.get("Output") or {}).get("FileInfos") or []
            urls = [f.get("FileUrl") or f.get("Url") for f in files if f.get("FileUrl") or f.get("Url")]
            if not urls:
                raise CoverError(f"任务完成但没有图片：{json.dumps(task, ensure_ascii=False)[:500]}")
            return [{"url": url} for url in urls]
        print(f"  状态：{status}，进度：{task.get('Progress', '-')}")
    raise CoverError(f"等待超时，任务 {task_id} 仍未完成，可稍后在腾讯云控制台查看。")


# ---------- 保存 ----------

def save_results(results: list[dict[str, str]], out_dir: Path, stem: str) -> list[Path]:
    out_dir.mkdir(parents=True, exist_ok=True)
    saved: list[Path] = []
    for index, item in enumerate(results, start=1):
        if "b64" in item:
            content, ext = base64.b64decode(item["b64"]), ".png"
        else:
            with urlopen(item["url"], timeout=120) as response:
                content = response.read()
                ext = mimetypes.guess_extension((response.headers.get_content_type() or "").strip()) or ""
            if ext not in IMAGE_EXTS:
                ext = Path(urlparse(item["url"]).path).suffix.lower()
            if ext not in IMAGE_EXTS:
                ext = ".png"
        suffix = f"_{index}" if len(results) > 1 else ""
        path = out_dir / f"{stem}{suffix}{ext}"
        path.write_bytes(content)
        saved.append(path)
    return saved


# ---------- 主流程 ----------

def describe_request(config: dict[str, Any], prompt: str, aspect: str, refs: list[str]) -> str:
    provider = config["provider"]
    if provider == "tencent-vod":
        payload = tencent_payload(config, prompt, aspect, refs)
        return f"POST https://vod.tencentcloudapi.com  Action=CreateAigcImageTask\n{json.dumps(payload, ensure_ascii=False, indent=2)}"
    req = (openai_request if provider == "openai" else ark_request)(config, prompt, aspect, refs)
    shown = dict(req.get("json") or req.get("fields") or {})
    if "image" in shown:
        shown["image"] = [r if is_url(r) else f"<base64: {r}>" for r in refs]
    if "files" in req:
        shown["image[]"] = [path for _, path in req["files"]]
    return f"POST {req['url']}\nAuthorization: Bearer {mask(config['api_key'])}\n{json.dumps(shown, ensure_ascii=False, indent=2)}"


def generate_once(config: dict[str, Any], prompt: str, aspect: str, refs: list[str]) -> list[dict[str, str]]:
    provider = config["provider"]
    if provider == "tencent-vod":
        return tencent_generate(config, tencent_payload(config, prompt, aspect, refs))
    req = (openai_request if provider == "openai" else ark_request)(config, prompt, aspect, refs)
    return send_simple_request(req)


def run_check(config_path: Path, plans_path: Path) -> None:
    config = load_config(config_path)
    print(f"配置文件：{config_path}")
    print(f"provider：{config['provider']}")
    if config["provider"] == "tencent-vod":
        print(f"secret_id：{mask(config['secret_id'])}  secret_key：{mask(config['secret_key'])}  sub_app_id：{config['sub_app_id']}")
        print(f"模型：{config.get('model_name')} / {config.get('model_version')} / {config.get('resolution')}")
    else:
        print(f"base_url：{config['base_url']}")
        print(f"api_key：{mask(config['api_key'])}")
        print(f"model：{config['model']}")
    if plans_path.is_file():
        parsed = parse_plans_file(plans_path)
        print(f"方案文件：{plans_path}（画幅 {parsed['aspect']}）")
        for number, plan in sorted(parsed["plans"].items()):
            print(f"  方案{int_to_cn(number)}：{plan['title']}  参考图文件：{plan['refs'] or '无'}")
    else:
        print(f"方案文件：{plans_path}（尚未生成）")
    print("配置检查通过（未发起任何网络请求）。")


def main() -> int:
    for stream in (sys.stdout, sys.stderr):
        if hasattr(stream, "reconfigure"):
            stream.reconfigure(encoding="utf-8")

    parser = argparse.ArgumentParser(description="lumimaker-game-cover-maker 出图脚本")
    parser.add_argument("--plans", type=Path, default=DEFAULT_PLANS_FILE, help="方案文件路径")
    parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG_FILE, help="配置文件路径")
    parser.add_argument("--out", type=Path, default=DEFAULT_OUT_DIR, help="图片输出目录")
    parser.add_argument("--plan", action="append", default=[], help="方案编号（1 或 一），可重复")
    parser.add_argument("--prompt-file", type=Path, help="直接使用该文件里的 Prompt（后期处理用），忽略 --plan")
    parser.add_argument("--ref", action="append", default=[], help="手动指定参考图（本地路径或 URL），可重复，会覆盖方案里的参考图")
    parser.add_argument("--no-refs", action="store_true", help="不使用任何参考图")
    parser.add_argument("--aspect", help="画幅，如 16:9、9:16，默认读取方案文件")
    parser.add_argument("--n", type=int, default=1, help="每个方案生成几张")
    parser.add_argument("--dry-run", action="store_true", help="只打印请求，不发送、不扣费")
    parser.add_argument("--check", action="store_true", help="检查配置和方案文件，不发送请求")
    args = parser.parse_args()

    try:
        if args.check:
            run_check(args.config, args.plans)
            return 0

        config = load_config(args.config)
        jobs: list[dict[str, Any]] = []
        if args.prompt_file:
            if not args.prompt_file.is_file():
                raise CoverError(f"找不到 Prompt 文件：{args.prompt_file}")
            jobs.append({
                "stem": "custom",
                "prompt": args.prompt_file.read_text(encoding="utf-8-sig").strip(),
                "refs": [],
                "aspect": DEFAULT_ASPECT,
            })
        else:
            if not args.plan:
                raise CoverError("请用 --plan 指定方案编号（如 --plan 1），或用 --prompt-file 指定 Prompt。")
            parsed = parse_plans_file(args.plans)
            for raw in args.plan:
                number = cn_to_int(raw)
                plan = parsed["plans"].get(number)
                if not plan:
                    available = "、".join(f"方案{int_to_cn(n)}" for n in sorted(parsed["plans"])) or "无"
                    raise CoverError(f"方案文件里没有方案{int_to_cn(number)}（或它的 Prompt 为空）。可用：{available}")
                jobs.append({
                    "stem": f"方案{int_to_cn(number)}",
                    "prompt": plan["prompt"],
                    "refs": plan["refs"],
                    "aspect": parsed["aspect"],
                })

        for job in jobs:
            aspect = args.aspect or job["aspect"]
            refs_raw = [] if args.no_refs else (args.ref or job["refs"])
            refs = resolve_refs(refs_raw, config["provider"])
            print(f"\n=== {job['stem']}  画幅 {aspect}  参考图 {len(refs)} 张 ===")
            if args.dry_run:
                print(describe_request(config, job["prompt"], aspect, refs))
                continue
            for _ in range(max(1, args.n)):
                results = generate_once(config, job["prompt"], aspect, refs)
                stem = f"{job['stem']}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
                for path in save_results(results, args.out, stem):
                    print(f"  已保存：{path}")
        if args.dry_run:
            print("\n（dry-run：未发送请求）")
        return 0
    except CoverError as exc:
        print(f"[错误] {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())