#!/usr/bin/env python3
"""本地视频转步态序列、调用银河水滴步态 API 提取特征，并计算点积相似度。"""

from __future__ import annotations

import argparse
import json
import mimetypes
import os
import re
import sys
import time
from pathlib import Path
from typing import Any
from urllib.parse import urljoin

import requests

from persondet_numpy import Config, run_video


IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}
TIMEOUT = 1800
MAX_BATCH_UPLOAD_BYTES = 32 << 20


def main() -> int:
    parser = argparse.ArgumentParser(description="银河水滴步态特征提取完整示例")
    parser.add_argument("video", help="本地视频路径")
    parser.add_argument("--base-url", default=os.environ.get("WATRIX_GAIT_BASE_URL", "https://www.watrixgait.com"))
    parser.add_argument("--api-key", default=os.environ.get("WATRIX_GAIT_API_KEY", ""))
    parser.add_argument("--output", default="watrix_gait_demo_output", help="输出目录")
    parser.add_argument("--min-frames", type=int, default=20)
    parser.add_argument("--max-frames", type=int, default=120)
    args = parser.parse_args()

    video_path = Path(args.video).expanduser().resolve()
    if not video_path.is_file():
        raise FileNotFoundError(video_path)
    if not args.api_key.strip():
        raise RuntimeError("请通过 --api-key 或环境变量 WATRIX_GAIT_API_KEY 设置 API Key")

    output_root = Path(args.output).expanduser().resolve()
    sequence_output = output_root / "sequences"
    result_output = output_root / "api_results"
    sequence_output.mkdir(parents=True, exist_ok=True)
    result_output.mkdir(parents=True, exist_ok=True)

    old_cwd = Path.cwd()
    try:
        os.chdir(sequence_output)
        run_video(video_path, Config(min_frames=args.min_frames, max_frames=args.max_frames))
    finally:
        os.chdir(old_cwd)

    sequence_root = sequence_output / f"{video_path.stem}_gait_sequences"
    sequence_dirs = collect_leaf_sequence_dirs(sequence_root)
    if not sequence_dirs:
        raise RuntimeError(f"没有从视频中导出有效序列：{sequence_root}")

    headers = {"Authorization": f"Bearer {args.api_key.strip()}"}
    summary: dict[str, Any] = {
        "started_at": iso_now(),
        "base_url": args.base_url,
        "video": str(video_path),
        "sequence_root": str(sequence_root),
        "sequence_count": len(sequence_dirs),
        "items": [],
    }

    with requests.Session() as session:
        check_health(session, args.base_url)
        for index, seq_dir in enumerate(sequence_dirs, start=1):
            print(f"[{index}/{len(sequence_dirs)}] 上传并解析序列：{seq_dir}")
            item = run_one_sequence(session, args.base_url, headers, seq_dir)
            out_path = result_output / f"{safe_filename(seq_dir.name)}.json"
            write_json(out_path, item)
            summary["items"].append({
                "sequence_dir": str(seq_dir),
                "task_id": item.get("task_id", ""),
                "ok": item.get("ok", False),
                "result_file": str(out_path),
                "error": item.get("error", ""),
            })
            write_json(output_root / "summary.json", summary)

    similarity = build_similarity_report(result_output)
    write_json(output_root / "similarity.json", similarity)
    summary["finished_at"] = iso_now()
    summary["similarity_file"] = str(output_root / "similarity.json")
    write_json(output_root / "summary.json", summary)
    print(f"结果目录：{output_root}")
    return 0


def run_one_sequence(session: requests.Session, base_url: str, headers: dict[str, str], seq_dir: Path) -> dict[str, Any]:
    started_at = iso_now()
    try:
        frames = collect_frames(seq_dir)
        created = request_json(session, base_url, "POST", "/v1/sequences", headers, {"frame_count": len(frames)})
        task_id = created["task_id"]
        upload = created["upload"]
        if upload.get("frame_count") != len(frames):
            raise RuntimeError(f"上传帧数量不一致：api={upload.get('frame_count')} local={len(frames)}")
        upload_batch(session, base_url, upload["upload_url"], upload["upload_token"], frames)

        parsed = request_json(
            session,
            base_url,
            "POST",
            f"/v1/sequences/{task_id}/parse",
            headers,
            {},
        )
        return {
            "ok": True,
            "started_at": started_at,
            "finished_at": iso_now(),
            "sequence_dir": str(seq_dir),
            "task_id": task_id,
            "frame_count": len(frames),
            "result": parsed,
        }
    except Exception as exc:
        return {
            "ok": False,
            "started_at": started_at,
            "finished_at": iso_now(),
            "sequence_dir": str(seq_dir),
            "error": str(exc),
        }


def collect_leaf_sequence_dirs(root: Path) -> list[Path]:
    dirs: list[Path] = []
    if not root.exists():
        return dirs
    for current, child_dirs, files in os.walk(root):
        if child_dirs:
            continue
        if any(Path(name).suffix.lower() in IMAGE_EXTS for name in files):
            dirs.append(Path(current))
    return sorted(dirs, key=lambda p: str(p))


def collect_frames(seq_dir: Path) -> list[Path]:
    frames = [p for p in seq_dir.iterdir() if p.is_file() and p.suffix.lower() in IMAGE_EXTS]
    frames.sort(key=lambda p: p.name)
    if not frames:
        raise RuntimeError(f"序列目录没有图片：{seq_dir}")
    return frames


def upload_batch(session: requests.Session, base_url: str, upload_url: str, upload_token: str, paths: list[Path]) -> None:
    # The API accepts one multipart request with a total size of at most 32 MB.
    total_bytes = 0
    files: list[tuple[str, tuple[str, bytes, str]]] = []
    for path in paths:
        body = path.read_bytes()
        total_bytes += len(body)
        if total_bytes > MAX_BATCH_UPLOAD_BYTES:
            raise RuntimeError("当前序列图片总大小超过 32 MB，请减少帧数或压缩图片")
        mime = mimetypes.guess_type(path.name)[0] or "application/octet-stream"
        # Repeated frames fields preserve the sequence order; filenames are ignored by the API.
        files.append(("frames", (path.name, body, mime)))
    url = urljoin(base_url.rstrip("/") + "/", upload_url.lstrip("/"))
    resp = session.post(url, data={"upload_token": upload_token}, files=files, timeout=TIMEOUT)
    if resp.status_code >= 400:
        raise_http(resp)


def request_json(
    session: requests.Session,
    base_url: str,
    method: str,
    path: str,
    headers: dict[str, str] | None = None,
    payload: Any | None = None,
) -> dict[str, Any]:
    req_headers = {"Accept": "application/json"}
    if headers:
        req_headers.update(headers)
    url = urljoin(base_url.rstrip("/") + "/", path.lstrip("/"))
    resp = session.request(method, url, headers=req_headers, json=payload, timeout=TIMEOUT)
    if resp.status_code >= 400:
        raise_http(resp)
    return resp.json() if resp.text else {}


def check_health(session: requests.Session, base_url: str) -> None:
    url = urljoin(base_url.rstrip("/") + "/", "healthz")
    resp = session.get(url, headers={"Accept": "application/json"}, timeout=10)
    if resp.status_code != 200:
        raise RuntimeError(f"服务不可用：HTTP {resp.status_code} {url}\n{resp.text}")


def build_similarity_report(result_dir: Path) -> dict[str, Any]:
    sequences: list[dict[str, Any]] = []
    for path in sorted(result_dir.glob("*.json")):
        item = read_json(path)
        if not item.get("ok"):
            continue
        for seq in (item.get("result") or {}).get("sequences") or []:
            if not isinstance(seq, dict):
                continue
            sequences.append({
                "source": str(path),
                "sequence_id": seq.get("sequence_id", ""),
                "frame_count": seq.get("frame_count", 0),
                "features": {
                    "gait_feature": numeric_vector(seq.get("gait_feature")),
                    "reid_feature": numeric_vector(seq.get("reid_feature")),
                    "face_feature": numeric_vector(seq.get("face_feature")),
                },
            })

    pairs = []
    for left_index in range(len(sequences)):
        for right_index in range(left_index + 1, len(sequences)):
            left = sequences[left_index]
            right = sequences[right_index]
            pairs.append({
                "left": identity(left),
                "right": identity(right),
                "gait_similarity": dot(left["features"]["gait_feature"], right["features"]["gait_feature"]),
                "reid_similarity": dot(left["features"]["reid_feature"], right["features"]["reid_feature"]),
                "face_similarity": dot(left["features"]["face_feature"], right["features"]["face_feature"]),
            })
    return {
        "generated_at": iso_now(),
        "sequence_count": len(sequences),
        "pair_count": len(pairs),
        "pairs": pairs,
    }


def dot(left: list[float], right: list[float]) -> dict[str, Any]:
    if not left or not right:
        return {"score": None, "reason": "empty_feature", "left_dim": len(left), "right_dim": len(right)}
    used_dim = min(len(left), len(right))
    return {
        "score": sum(left[i] * right[i] for i in range(used_dim)),
        "left_dim": len(left),
        "right_dim": len(right),
        "used_dim": used_dim,
        "dimension_mismatch": len(left) != len(right),
    }


def numeric_vector(value: Any) -> list[float]:
    if not isinstance(value, list):
        return []
    return [float(v) for v in value if isinstance(v, (int, float))]


def identity(item: dict[str, Any]) -> dict[str, Any]:
    return {
        "source": item.get("source", ""),
        "sequence_id": item.get("sequence_id", ""),
        "frame_count": item.get("frame_count", 0),
    }


def raise_http(resp: requests.Response) -> None:
    raise RuntimeError(f"HTTP {resp.status_code} {resp.request.method} {resp.request.url}\n{resp.text}")


def write_json(path: Path, data: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def read_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


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


def iso_now() -> str:
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except KeyboardInterrupt:
        raise SystemExit(130)
    except Exception as exc:
        print(f"错误：{exc}", file=sys.stderr)
        raise
