# -*- coding: utf-8 -*-
"""①前処理: 全動画の濃度解決（湯気除外→スープ抽出→濃度推定）をキャッシュ化する。

LLM送信(②)で使い回すために、各動画について
  - density_resolver.resolve() の結果（濃度値・可読性・assessable・steam等）
  - LLMに送る「最良フレーム画像」（恒久フォルダに保持）
を density_cache.jsonl に1行ずつ保存する（resume対応）。

明日 batch_process.py を density.cache_path 指定で回せば、①を再計算せず②だけ実行できる。
"""

from __future__ import annotations

import argparse
import csv
import glob
import json
import shutil
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "local_batch_api"))
from density_resolver import DensityResolver, normalize_density  # noqa: E402

VIDEO_EXTS = [".MOV", ".mov", ".mp4", ".MP4", ".m4v", ".avi", ".mkv"]


def find_video(video_dir: Path, vid: str) -> Path | None:
    for e in VIDEO_EXTS:
        c = video_dir / f"{vid}{e}"
        if c.exists():
            return c
    return None


def load_human(label_csv: str) -> dict[str, str]:
    out: dict[str, str] = {}
    fs = glob.glob(label_csv)
    if fs:
        for r in csv.DictReader(open(fs[0], encoding="utf-8-sig")):
            vid = (r.get("ID") or "").strip()
            if vid:
                out[vid] = normalize_density(r.get("スープ濃度"))
    return out


def main() -> int:
    ap = argparse.ArgumentParser(description="①濃度前処理キャッシュ生成（LLM送信で使い回す）")
    ap.add_argument("--video-dir", default=str(ROOT / "学習データ"))
    ap.add_argument("--label-csv", default=str(ROOT / "学習データ" / "*グリッド 10.csv"))
    ap.add_argument("--frames", type=int, default=16)
    ap.add_argument("--cache", default=str(ROOT / "local_batch_api" / "density_cache.jsonl"))
    ap.add_argument("--best-frames-dir", default=str(ROOT / "local_batch_api" / "prepared_best_frames"))
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--resume", action="store_true")
    args = ap.parse_args()

    video_dir = Path(args.video_dir)
    human = load_human(args.label_csv)
    ids = sorted(human.keys(), key=lambda x: int(x) if x.isdigit() else 0)
    if args.limit > 0:
        ids = ids[: args.limit]

    cache_path = Path(args.cache)
    best_dir = Path(args.best_frames_dir)
    best_dir.mkdir(parents=True, exist_ok=True)
    scratch = ROOT / "local_batch_api" / "prepare_scratch"
    scratch.mkdir(parents=True, exist_ok=True)

    done: set[str] = set()
    if args.resume and cache_path.exists():
        for line in cache_path.open("r", encoding="utf-8"):
            line = line.strip()
            if line:
                try:
                    done.add(json.loads(line).get("stem"))
                except json.JSONDecodeError:
                    pass
        print(f"resume: 既存 {len(done)} 件をスキップ", flush=True)

    resolver = DensityResolver(frames=args.frames, crops_dir=scratch)
    mode = "a" if (args.resume and cache_path.exists()) else "w"
    fh = cache_path.open(mode, encoding="utf-8")

    total = len(ids)
    done_n = 0
    for vid in ids:
        if vid in done:
            continue
        vpath = find_video(video_dir, vid)
        if vpath is None:
            continue
        meta = {"スープ濃度": human.get(vid, "")}  # 提供濃度あり経路（本番と同条件）
        res = resolver.resolve(vpath, meta)

        # 最良フレームを恒久フォルダへ退避し、パスを差し替え
        bf = res.get("best_frame_path") or ""
        if bf and Path(bf).exists():
            dest = best_dir / f"{vpath.stem}_best.jpg"
            shutil.copyfile(bf, dest)
            res["best_frame_path"] = str(dest)

        # scratch をこの動画分だけ掃除（ディスク節約）
        for p in scratch.glob(f"{vpath.stem}_*"):
            try:
                p.unlink()
            except OSError:
                pass

        rec = {"stem": vpath.stem, "video_path": str(vpath), "human": human.get(vid, ""), **res}
        fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
        fh.flush()
        done_n += 1
        if done_n % 10 == 0 or done_n == total:
            print(f"prepared {done_n} (last id={vid} -> {res.get('value')} {res.get('confidence')} "
                  f"readable={res.get('readable')})", flush=True)

    fh.close()
    print(f"\ndone. cache -> {cache_path}\nbest frames -> {best_dir}", flush=True)
    print("明日は batch_process.py を density.cache_path 指定で実行すれば①を再計算しません。", flush=True)
    return 0


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