#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
tools/enrich_prompt_start_frame_strategies.py
Công cụ làm giàu metadata định tuyến Start Frame cho toàn bộ các file Prompt JSON:
- start_frame_strategy: CHARACTER_TURNAROUND_VIEW | GEMINI_BANANA_KEYFRAME | HEAD_TAIL_PREV_FRAME | TEXT_TO_VIDEO
- turnaround_view: portrait | front | profile (cho CHARACTER_TURNAROUND_VIEW)
"""

import sys
import json
import argparse
from pathlib import Path
from collections import Counter

sys.stdout.reconfigure(encoding='utf-8')

PROJECT_ROOT = Path(__file__).resolve().parent.parent
PIPELINE_DIR = PROJECT_ROOT / "05_Production_Pipeline"
if str(PIPELINE_DIR) not in sys.path:
    sys.path.insert(0, str(PIPELINE_DIR))

import production_orchestrator as po


def enrich_episode_prompts(ep_id: str = "ep01", apply: bool = False):
    prompt_file = PROJECT_ROOT / "episodes" / ep_id / "prompts" / "muse_prompts.json"
    if not prompt_file.exists():
        print(f"[!] Không tìm thấy file prompt cho {ep_id}: {prompt_file}")
        return

    data = json.loads(prompt_file.read_text(encoding="utf-8"))
    motion_prompts = data.get("motion_prompts", {})

    print(f"\n=======================================================")
    print(f"🎬 LÀM GIÀU METADATA START FRAME CHO TẬP {ep_id.upper()}")
    print(f"   • File: {prompt_file}")
    print(f"   • Tổng số shots: {len(motion_prompts)}")
    print(f"   • Chế độ: {'APPLY (Ghi đĩa)' if apply else 'DRY-RUN (Xem trước)'}")
    print(f"=======================================================\n")

    counter = Counter()
    enriched_count = 0

    for sid, sdata in motion_prompts.items():
        take_type = po.classify_shot_take(sid, shot_data=sdata)
        strat, vtype = po.classify_start_frame_strategy(sid, sdata, take_type=take_type, resolved_path="clean_frame_239")
        counter[strat] += 1

        # Gán metadata nếu chưa có hoặc cập nhật
        sdata["start_frame_strategy"] = strat
        if vtype:
            sdata["turnaround_view"] = vtype
        elif "turnaround_view" in sdata and strat != "CHARACTER_TURNAROUND_VIEW":
            del sdata["turnaround_view"]

        enriched_count += 1

    print("📊 Phân bố chiến lược Start Frame:")
    for strat, count in counter.most_common():
        pct = (count / len(motion_prompts)) * 100
        print(f"   • {strat:<26}: {count:>3} shots ({pct:>5.1f}%)")

    if apply:
        prompt_file.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
        print(f"\n[✓] Đã cập nhật thành công {enriched_count} shots vào {prompt_file.name}!")
    else:
        print(f"\n[i] Chạy thử nghiệm hoàn tất. Thêm cờ --apply để lưu vào file.")


def main():
    parser = argparse.ArgumentParser(description="Làm giàu metadata Start Frame cho Prompt JSON")
    parser.add_argument("--episode", type=str, default="ep01", help="Mã tập (ep01 .. ep06)")
    parser.add_argument("--all", action="store_true", help="Làm giàu toàn bộ 6 tập")
    parser.add_argument("--apply", action="store_true", help="Áp dụng ghi đè vào tệp prompt")

    args = parser.parse_args()
    if args.all:
        for i in range(1, 7):
            enrich_episode_prompts(f"ep{i:02d}", apply=args.apply)
    else:
        enrich_episode_prompts(args.episode, apply=args.apply)


if __name__ == "__main__":
    main()