#!/usr/bin/env python3
"""Download a curated Mixamo animation set for Y Bot (male) and X Bot (female)."""

from __future__ import annotations

import argparse
import json
import re
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path

API = "https://www.mixamo.com/api/v1"

# Local filenames -> Mixamo search query + preferred exact name
ANIMATIONS = [
    ("Idle", "Idle", "Idle"),
    ("Walking", "Walking", "Walking"),
    ("Running", "Running", "Running"),
    ("Punching", "Punching", "Punching"),
    ("Kick", "Kicking", "Kicking"),
    ("SwordSlash", "Stable Sword Outward Slash", "Stable Sword Outward Slash"),
    ("HitReaction", "Hit Reaction", "Hit Reaction"),
    ("FallingFlatImpact", "Falling Flat Impact", "Falling Flat Impact"),
]

CHARACTERS = {
    "Male": "Y Bot",
    "Female": "X Bot",
}


def headers(token: str) -> dict[str, str]:
    return {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": f"Bearer {token}",
        "X-Api-Key": "mixamo2",
        "X-Requested-With": "XMLHttpRequest",
    }


def api_request(token: str, method: str, path: str, body: dict | None = None, query: dict | None = None):
    url = f"{API}{path}"
    if query:
        url += "?" + urllib.parse.urlencode(query)
    data = None if body is None else json.dumps(body).encode("utf-8")
    req = urllib.request.Request(url, data=data, headers=headers(token), method=method)
    try:
        with urllib.request.urlopen(req, timeout=60) as resp:
            raw = resp.read()
            if not raw:
                return resp.status, None
            try:
                return resp.status, json.loads(raw.decode("utf-8"))
            except json.JSONDecodeError:
                return resp.status, raw
    except urllib.error.HTTPError as e:
        raw = e.read()
        try:
            payload = json.loads(raw.decode("utf-8")) if raw else None
        except json.JSONDecodeError:
            payload = raw.decode("utf-8", errors="replace") if raw else None
        return e.code, payload


def search_products(token: str, query: str, ptype: str, pages: int = 3, limit: int = 96) -> list[dict]:
    seen: set[str] = set()
    items: list[dict] = []
    for page in range(1, pages + 1):
        status, data = api_request(
            token,
            "GET",
            "/products",
            query={
                "page": str(page),
                "limit": str(limit),
                "order": "",
                "type": ptype,
                "query": query,
            },
        )
        if status != 200 or not isinstance(data, dict):
            raise RuntimeError(f"search failed ({status}): {data}")
        for item in data.get("results") or []:
            pid = item.get("id")
            if not pid or pid in seen:
                continue
            seen.add(pid)
            items.append(item)
        num_pages = (data.get("pagination") or {}).get("num_pages") or page
        if page >= int(num_pages):
            break
    return items


def product_label(item: dict) -> str:
    return (item.get("name") or item.get("description") or "").strip()


def pick_by_name(results: list[dict], preferred: str) -> dict:
    preferred_l = preferred.lower().strip()
    exact = [item for item in results if product_label(item).lower() == preferred_l]
    if exact:
        return exact[0]
    contains = [item for item in results if preferred_l in product_label(item).lower()]
    if contains:
        # Prefer shortest label (closest to exact phrase)
        contains.sort(key=lambda item: len(product_label(item)))
        return contains[0]
    if not results:
        raise RuntimeError(f"no products for '{preferred}'")
    return results[0]


def resolve_character(token: str, name: str) -> dict:
    results = search_products(token, name, "Character", pages=1, limit=24)
    item = pick_by_name(results, name)
    return {"id": item["id"], "name": product_label(item) or name}


def resolve_animation(token: str, query: str, preferred: str) -> dict:
    results = search_products(token, query, "Motion,MotionPack", pages=3)
    motions = [r for r in results if (r.get("type") or "").lower() == "motion"]
    pool = motions or results
    item = pick_by_name(pool, preferred)
    return {
        "id": item["id"],
        "name": product_label(item) or preferred,
        "type": item.get("type"),
    }


def get_product(token: str, anim_id: str, character_id: str) -> dict:
    status, data = api_request(
        token,
        "GET",
        f"/products/{anim_id}",
        query={"similar": "0", "character_id": character_id},
    )
    if status != 200 or not isinstance(data, dict):
        raise RuntimeError(f"product {anim_id} failed ({status}): {data}")
    return data


def normalize_gms_hash(gms_hash: dict) -> dict:
    """Bake default params into the string form Mixamo export expects."""
    out = dict(gms_hash)
    params = out.get("params")
    if isinstance(params, list):
        vals: list[str] = []
        for item in params:
            if isinstance(item, (list, tuple)) and len(item) >= 2:
                vals.append(str(item[1]))
            else:
                vals.append(str(item))
        out["params"] = ",".join(vals)
    elif isinstance(params, str) and "=" in params:
        vals = []
        for part in params.split(","):
            part = part.strip()
            if not part:
                continue
            vals.append(part.split("=", 1)[1] if "=" in part else part)
        out["params"] = ",".join(vals)
    return out


def export_animation(token: str, character_id: str, gms_hash: dict, product_name: str) -> None:
    body = {
        "character_id": character_id,
        "gms_hash": [gms_hash],
        "preferences": {
            "format": "fbx7_unity",
            "skin": "false",
            "fps": "30",
            "reducekf": "0",
        },
        "product_name": product_name,
        "type": "Motion",
    }
    for attempt in range(8):
        status, data = api_request(token, "POST", "/animations/export", body=body)
        if status in (200, 202):
            return
        if status == 429:
            time.sleep(1.5 * (attempt + 1))
            continue
        raise RuntimeError(f"export failed ({status}): {data}")
    raise RuntimeError("export rate-limited too many times")


def monitor_export(token: str, character_id: str, timeout_s: float = 120.0) -> str:
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        status, data = api_request(token, "GET", f"/characters/{character_id}/monitor")
        if status == 429:
            time.sleep(2)
            continue
        if status != 200 or not isinstance(data, dict):
            time.sleep(1.5)
            continue
        job_status = (data.get("status") or "").lower()
        if job_status == "completed":
            result = data.get("job_result")
            if isinstance(result, str) and result.startswith("http"):
                return result
            if isinstance(result, dict):
                url = result.get("url") or result.get("result") or result.get("message")
                if isinstance(url, str) and url.startswith("http"):
                    return url
            maybe = data.get("message")
            if isinstance(maybe, str) and maybe.startswith("http"):
                return maybe
            raise RuntimeError(f"completed but no url: {data}")
        if job_status == "failed":
            raise RuntimeError(f"export monitor failed: {data}")
        time.sleep(1.2)
    raise TimeoutError(f"export monitor timeout for {character_id}")


def download_file(url: str, dest: Path) -> None:
    req = urllib.request.Request(url, method="GET")
    with urllib.request.urlopen(req, timeout=120) as resp:
        data = resp.read()
    dest.parent.mkdir(parents=True, exist_ok=True)
    dest.write_bytes(data)


def safe_filename(name: str) -> str:
    return re.sub(r"[^A-Za-z0-9_.-]+", "_", name).strip("_")


def download_one(
    token: str,
    character: dict,
    anim: dict,
    local_name: str,
    out_dir: Path,
    skip_existing: bool,
) -> Path:
    dest = out_dir / f"{local_name}.fbx"
    if skip_existing and dest.exists() and dest.stat().st_size > 1000:
        print(f"  skip existing {dest}")
        return dest

    product = get_product(token, anim["id"], character["id"])
    details = product.get("details") or {}
    gms = details.get("gms_hash")
    if not gms:
        raise RuntimeError(f"no gms_hash for {anim['name']}: {product}")
    gms = normalize_gms_hash(gms)
    product_name = product.get("name") or anim["name"]

    print(f"  export {product_name} -> {character['name']}")
    export_animation(token, character["id"], gms, product_name)
    url = monitor_export(token, character["id"])
    print(f"  download {url[:80]}...")
    download_file(url, dest)
    print(f"  saved {dest} ({dest.stat().st_size} bytes)")
    return dest


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--token-file", type=Path, required=True)
    parser.add_argument("--out", type=Path, default=Path("Tools/output/mixamo"))
    parser.add_argument("--skip-existing", action="store_true", default=True)
    args = parser.parse_args()

    token = args.token_file.read_text(encoding="utf-8").strip()
    if not token:
        print("empty token", file=sys.stderr)
        return 1

    print("Resolving characters...")
    chars = {}
    for gender, name in CHARACTERS.items():
        chars[gender] = resolve_character(token, name)
        print(f"  {gender}: {chars[gender]['name']} ({chars[gender]['id']})")

    print("Resolving animations...")
    anims = []
    for local_name, query, preferred in ANIMATIONS:
        anim = resolve_animation(token, query, preferred)
        anims.append((local_name, anim))
        print(f"  {local_name}: {anim['name']} ({anim['id']}) type={anim.get('type')}")

    results = []
    for gender, character in chars.items():
        out_dir = args.out / gender
        out_dir.mkdir(parents=True, exist_ok=True)
        print(f"\n=== {gender} / {character['name']} ===")
        for local_name, anim in anims:
            try:
                path = download_one(
                    token, character, anim, local_name, out_dir, args.skip_existing
                )
                results.append({"gender": gender, "anim": local_name, "path": str(path), "ok": True})
            except Exception as e:
                print(f"  ERROR {local_name}: {e}")
                results.append({"gender": gender, "anim": local_name, "error": str(e), "ok": False})
            time.sleep(0.8)

    manifest = args.out / "manifest.json"
    manifest.write_text(json.dumps(results, indent=2, ensure_ascii=False), encoding="utf-8")
    ok = sum(1 for r in results if r["ok"])
    print(f"\nDone: {ok}/{len(results)} ok. Manifest: {manifest}")
    return 0 if ok == len(results) else 2


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