"""
Batch 4K AI Video Upscaler for Historical Series
Orchestrates sequential 4K Super-Resolution for all episodes (ep1 to ep8)
with progress tracking, resume capability, and safe VRAM/disk management.
"""

import os
import sys
import json
import time
import argparse
from pathlib import Path
from video_upscaler import upscale_video

# Force UTF-8 encoding on Windows console
if sys.platform.startswith("win"):
    try:
        sys.stdout.reconfigure(encoding="utf-8")
        sys.stderr.reconfigure(encoding="utf-8")
    except Exception:
        pass

STATE_FILE = Path("upscale_state.json")

def load_state():
    if STATE_FILE.exists():
        try:
            with open(STATE_FILE, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception:
            pass
    return {"completed": {}, "in_progress": None, "last_updated": None}

def save_state(state):
    state["last_updated"] = time.strftime("%Y-%m-%d %H:%M:%S")
    with open(STATE_FILE, "w", encoding="utf-8") as f:
        json.dump(state, f, indent=2, ensure_ascii=False)

def print_status():
    state = load_state()
    print("=== TRẠNG THÁI UPSCALE 4K CÁC TẬP PHIM ===")
    completed = state.get("completed", {})
    for i in range(1, 9):
        ep_key = f"ep{i}"
        if ep_key in completed:
            info = completed[ep_key]
            print(f" [x] Tập {i}: ĐÃ HOÀN THÀNH ({info.get('duration_min')} phút, {info.get('size_mb')} MB) - Lúc {info.get('finished_at')}")
        elif state.get("in_progress") == ep_key:
            print(f" [>] Tập {i}: ĐANG XỬ LÝ TRONG NỀN...")
        else:
            print(f" [ ] Tập {i}: Đang chờ trong hàng đợi")
    print("==========================================")

def run_batch_upscale(target_episodes=None, resume=True, tile_size=256, target_res="3840x2160"):
    source_dir = Path("downloads")
    out_dir = Path("content/videos_4k")
    out_dir.mkdir(parents=True, exist_ok=True)

    if not target_episodes:
        ep_files = sorted(list(source_dir.glob("ep[1-8].mp4")))
    else:
        ep_files = [source_dir / f"ep{i}.mp4" for i in target_episodes if (source_dir / f"ep{i}.mp4").exists()]

    state = load_state()
    print(f"\n=== KHỞI CHẠY BATCH UPSCALE 4K ({len(ep_files)} TẬP) ===")
    print(f"Card đồ họa: NVIDIA GeForce GTX 1660 Super (6GB VRAM)")
    print(f"Thư mục xuất 4K: {out_dir.resolve()}\n")

    for ep_file in ep_files:
        ep_name = ep_file.stem
        out_file = out_dir / f"{ep_name}_4k.mp4"

        # Kiểm tra nếu đã hoàn thành trước đó
        if resume and (ep_name in state["completed"] or (out_file.exists() and out_file.stat().st_size > 10 * 1024 * 1024)):
            print(f"[*] Bỏ qua {ep_name} (Đã có file 4K hoàn chỉnh: {out_file.name})")
            continue

        state["in_progress"] = ep_name
        save_state(state)

        start_ts = time.time()
        print(f"\n==================================================")
        print(f"BẮT ĐẦU UPSCALE: {ep_file.name} -> {out_file.name}")
        print(f"Bắt đầu lúc: {time.strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"==================================================")

        try:
            upscale_video(
                input_video=str(ep_file),
                output_video=str(out_file),
                tile_size=tile_size,
                model_name="realesrgan-x4plus",
                target_res=target_res,
                cq=19
            )

            elapsed_min = round((time.time() - start_ts) / 60, 2)
            size_mb = round(out_file.stat().st_size / (1024 * 1024), 2)

            state["completed"][ep_name] = {
                "file": str(out_file),
                "duration_min": elapsed_min,
                "size_mb": size_mb,
                "finished_at": time.strftime("%Y-%m-%d %H:%M:%S")
            }
            state["in_progress"] = None
            save_state(state)

            print(f"--> HOÀN THÀNH {ep_name}: {size_mb} MB trong {elapsed_min} phút.")

        except Exception as e:
            print(f"\n[!] LỖI khi xử lý {ep_name}: {e}")
            state["in_progress"] = None
            save_state(state)
            # Tiếp tục tập tiếp theo thay vì dừng toàn bộ
            continue

    print("\n=== TOÀN BỘ TIẾN TRÌNH BATCH ĐÃ KẾT THÚC ===")
    print_status()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Batch AI 4K Upscaler for Historical Series")
    parser.add_argument("--status", action="store_true", help="Xem báo cáo tiến độ hiện tại")
    parser.add_argument("--episodes", nargs="+", type=int, help="Chỉ định danh sách tập cần chạy (vd: --episodes 1 2 3)")
    parser.add_argument("--no-resume", action="store_true", help="Chạy lại từ đầu kể cả khi đã có file")
    parser.add_argument("--tile", type=int, default=256, help="Tile size (mặc định 256)")

    args = parser.parse_args()

    if args.status:
        print_status()
    else:
        run_batch_upscale(
            target_episodes=args.episodes,
            resume=not args.no_resume,
            tile_size=args.tile
        )