#!/usr/bin/env python3
from __future__ import annotations

import argparse
import hashlib
import json
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from services.ai.qwen_client import QwenImageGenerationClient, _ensure_qwen_image_model, _to_data_url  # noqa: E402


WALK_STRIP_PROMPT = """基于这张参考图，生成一张横向排列的 4 帧连续走路动画图。只输出生成后的 4 个动作帧，不要把参考图放进画面。

四帧从左到右依次排列，间距均匀，大小一致，都是同一角色的右侧侧面正在往右走路。身体横向、面部朝右，保留主要外貌特征。纯白底图片干净无杂物，不要文字、编号、边框、分割线。

第 1 帧：前方一只前脚向右前方伸出，另一只前脚在身体下方支撑，后脚自然跟随。
第 2 帧：伸出的前脚落地支撑，另一只前脚开始向右前方抬起，身体略向右移动。
第 3 帧：另一只前脚向右前方伸出，后脚交替前进，保持自然四脚走路姿态。
第 4 帧：前脚回到接近第 1 帧的循环姿态，身体仍朝右，准备接回下一轮走路。"""


def main() -> int:
    parser = argparse.ArgumentParser(description="Probe Qwen with a one-shot four-frame walk strip prompt.")
    parser.add_argument("--pet-name", required=True)
    parser.add_argument("--canonical", required=True, help="Canonical PNG reference image.")
    parser.add_argument("--output", default=str(ROOT / "outputs" / "pet_builds" / "walk_strip_prompt_probes"))
    parser.add_argument("--size", default="1280*720")
    parser.add_argument("--prompt", default=WALK_STRIP_PROMPT, help="Exact one-shot strip prompt.")
    args = parser.parse_args()

    canonical_path = Path(args.canonical)
    if not canonical_path.exists():
        raise FileNotFoundError(canonical_path)
    canonical_bytes = canonical_path.read_bytes()
    if not canonical_bytes:
        raise ValueError("canonical image is empty")

    output_dir = Path(args.output) / f"walk-strip-{_slugify(args.pet_name)}-{_utc_stamp()}"
    raw_dir = output_dir / "raw"
    prompt_dir = output_dir / "prompts"
    qa_dir = output_dir / "qa"
    for path in (raw_dir, prompt_dir, qa_dir):
        path.mkdir(parents=True, exist_ok=True)

    prompt = args.prompt.strip()
    (prompt_dir / "walk-strip.txt").write_text(prompt, encoding="utf-8")

    client = QwenImageGenerationClient.from_env()
    _ensure_qwen_image_model(client.edit_model)
    content = [
        {"image": _to_data_url(canonical_bytes, "image/png")},
        {"text": prompt},
    ]
    result = client._generate_image(
        content=content,
        prompt=prompt,
        size=args.size,
        model=client.edit_model,
        negative_prompt="",
    )

    image_bytes = result["image_bytes"]
    raw_path = raw_dir / "walk-strip.png"
    raw_path.write_bytes(image_bytes)
    (qa_dir / "walk-strip-preview.png").write_bytes(image_bytes)

    manifest: dict[str, Any] = {
        "kind": "qwen_walk_strip_prompt_probe",
        "pet_name": args.pet_name,
        "canonical_path": str(canonical_path),
        "model": client.edit_model,
        "size": args.size,
        "prompt_sha256": _sha256(prompt),
        "request_id": result.get("request_id"),
        "usage": result.get("usage", {}),
        "transparency_mode": "disabled",
        "outputs": {
            "raw": "raw/walk-strip.png",
            "preview": "qa/walk-strip-preview.png",
            "prompt": "prompts/walk-strip.txt",
        },
    }
    (output_dir / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
    print(
        json.dumps(
            {
                "status": "ok",
                "output_dir": str(output_dir.resolve()),
                "raw": str(raw_path.resolve()),
                "preview": str((qa_dir / "walk-strip-preview.png").resolve()),
                "request_id": result.get("request_id"),
            },
            ensure_ascii=False,
            indent=2,
        )
    )
    return 0


def _slugify(value: str) -> str:
    safe = "".join(char if char.isalnum() or char in "-_" else "-" for char in value.strip().lower())
    return "-".join(part for part in safe.split("-") if part)[:80] or "pet"


def _utc_stamp() -> str:
    return datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S")


def _sha256(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


if __name__ == "__main__":
    raise SystemExit(main())