#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
4K Cinema Movie Concatenator & Chapter Injector
==============================================
Ghép 8 tập phim ngắn 4K UHD thành 1 bộ phim hoàn chỉnh (Full Feature Movie)
- Bảo toàn 100% chất lượng video 4K HEVC gốc (Zero Re-encoding / Stream Copy)
- Chuẩn hóa đồng bộ luồng âm thanh AAC Stereo 48,000 Hz 192kbps (tránh lệch/mất tiếng ở EP8)
- Đồng bộ hóa Video Track Timescale (12288) triệt tiêu lỗi giật/nhảy khung hình ở EP7
- Tự động nhúng Chapter Markers (phân đoạn từng tập) trực tiếp vào Metadata MP4
- Tự động xuất file Timeline YouTube Chapters (sẵn sàng copy-paste vào YouTube Description)
"""

import os
import sys
import re
import shutil
import argparse
import subprocess
from pathlib import Path

# Đảm bảo in tiếng Việt chuẩn trên Windows Terminal
if sys.platform == "win32":
    try:
        sys.stdout.reconfigure(encoding="utf-8")
        sys.stderr.reconfigure(encoding="utf-8")
    except Exception:
        pass

DEFAULT_TITLES = [
    "The Exiled Prince & The Desperate Alliance (1785)",
    "50,000 Invaders: The Siamese Fleet Dominates the River",
    "A Nation in Flames: The Fall of Gia Dinh",
    "The Tiger Awakens: General Nguyen Hue Strikes Back",
    "The Deadliest River Trap in Military History: Rach Gam (1785)",
    "Night of Fire: The Fire Ships Unleashed",
    "The Annihilation: Blood and Ash on the River",
    "Strategic Legacy: Sovereignty and Geopolitics in 1785"
]

def find_ffmpeg():
    """Tìm đường dẫn ffmpeg.exe ưu tiên venv."""
    local_ffmpeg = Path(__file__).parent / "venv" / "Scripts" / "ffmpeg.exe"
    if local_ffmpeg.exists():
        return str(local_ffmpeg)
    system_ffmpeg = shutil.which("ffmpeg")
    if system_ffmpeg:
        return system_ffmpeg
    raise FileNotFoundError("Không tìm thấy ffmpeg.exe trong venv/Scripts/ hoặc hệ thống!")

def get_video_info(ffmpeg_bin, filepath):
    """Lấy thông tin duration, timescale, audio stream qua ffmpeg."""
    cmd = [ffmpeg_bin, "-i", filepath]
    res = subprocess.run(cmd, stderr=subprocess.PIPE, text=True, encoding="utf-8", errors="ignore")
    
    # Parse Duration
    dur_match = re.search(r"Duration:\s*(\d+):(\d+):(\d+\.\d+)", res.stderr)
    if dur_match:
        h, m, s = map(float, dur_match.groups())
        duration_sec = h * 3600 + m * 60 + s
    else:
        duration_sec = 0.0

    # Parse Audio info
    audio_info = "Unknown"
    for line in res.stderr.splitlines():
        if "Stream #0:1" in line or "Audio:" in line:
            audio_info = line.strip()
            break
            
    return {
        "duration_sec": duration_sec,
        "audio_info": audio_info,
        "raw_stderr": res.stderr
    }

def format_timestamp(seconds):
    """Định dạng giây thành HH:MM:SS."""
    s_int = int(seconds)
    hours = s_int // 3600
    minutes = (s_int % 3600) // 60
    secs = s_int % 60
    return f"{hours:02d}:{minutes:02d}:{secs:02d}"

def normalize_segment(ffmpeg_bin, in_path, out_ts_path):
    """
    Chuẩn hóa phân đoạn sang MPEG-TS với bộ lọc hevc_mp4toannexb:
    - Bắt buộc chèn in-band parameter sets (VPS, SPS, PPS) trước mỗi IRAP/IDR/CRA
      để triệt tiêu lỗi POC mismatch, QP delta out of range, và hiện tượng vỡ hình/green screen
      khi chuyển giao giữa các tập phim 4K được render độc lập.
    - Chuẩn hóa luồng âm thanh sang AAC Stereo 48,000 Hz 192kbps.
    - Đồng bộ hóa timebase 90kHz tự nhiên của MPEG-TS.
    """
    cmd = [
        ffmpeg_bin, "-y",
        "-i", str(in_path),
        "-c:v", "copy",
        "-bsf:v", "hevc_mp4toannexb",
        "-c:a", "aac",
        "-ar", "48000",
        "-ac", "2",
        "-b:a", "192k",
        str(out_ts_path)
    ]
    res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding="utf-8", errors="ignore")
    if res.returncode != 0:
        raise RuntimeError(f"Lỗi chuẩn hóa file {in_path}:\n{res.stderr[-500:]}")

def generate_ffmetadata(movie_title, chapters, out_meta_path):
    """Tạo file FFMETADATA1 cho FFmpeg để nhúng chapter markers."""
    lines = [
        ";FFMETADATA1",
        f"title={movie_title}",
        "artist=The Historical Vault",
        "genre=Historical Documentary",
        ""
    ]
    for ch in chapters:
        lines.append("[CHAPTER]")
        lines.append("TIMEBASE=1/1000")
        lines.append(f"START={ch['start_ms']}")
        lines.append(f"END={ch['end_ms']}")
        lines.append(f"title=Tập {ch['ep']}: {ch['title']}")
        lines.append("")

    with open(out_meta_path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))

def main():
    parser = argparse.ArgumentParser(description="Ghép 8 tập 4K thành 1 bộ phim hoàn chỉnh.")
    parser.add_argument("--input-dir", default="content/videos_4k", help="Thư mục chứa 8 tập phim 4K")
    parser.add_argument("--output", default="content/videos_4k/Forgotten_Battles_1785_Full_4K.mp4", help="Đường dẫn file xuất")
    parser.add_argument("--title", default="Forgotten Battles (1785) - The Complete Series (4K UHD)", help="Tên phim")
    parser.add_argument("--dry-run", action="store_true", help="Chỉ kiểm tra và tính toán timeline, không ghép file")
    parser.add_argument("--keep-temp", action="store_true", help="Giữ lại các file trung gian")
    args = parser.parse_args()

    ffmpeg_bin = find_ffmpeg()
    print("=" * 70)
    print("🎬 4K CINEMA OMNIBUS CONCATENATOR (THE HISTORICAL VAULT)")
    print("=" * 70)
    print(f"FFmpeg Binary : {ffmpeg_bin}")
    print(f"Input Dir     : {args.input_dir}")
    print(f"Output File   : {args.output}")
    print(f"Movie Title   : {args.title}\n")

    input_dir = Path(args.input_dir)
    ep_files = []
    for i in range(1, 9):
        f = input_dir / f"ep{i}_4k.mp4"
        if not f.exists():
            raise FileNotFoundError(f"Không tìm thấy tập {f.name} tại {input_dir}")
        ep_files.append(f)

    # 1. Tính toán Timeline & Chapters
    print("🔍 [Bước 1/4] Khảo sát luồng và tính toán timeline chi tiết:")
    current_time = 0.0
    chapters = []
    yt_chapters = []

    for i, f in enumerate(ep_files, 1):
        info = get_video_info(ffmpeg_bin, str(f))
        dur = info["duration_sec"]
        start_sec = current_time
        end_sec = current_time + dur
        current_time = end_sec
        title = DEFAULT_TITLES[i - 1] if i <= len(DEFAULT_TITLES) else f"Phần {i}"
        
        start_str = format_timestamp(start_sec)
        end_str = format_timestamp(end_sec)
        
        chapters.append({
            "ep": i,
            "title": title,
            "start_ms": int(start_sec * 1000),
            "end_ms": int(end_sec * 1000),
            "start_str": start_str,
            "end_str": end_str,
            "duration": dur,
            "file": f
        })
        yt_chapters.append(f"{start_str} - Tập {i}: {title}")
        print(f"  • EP{i}: {start_str} -> {end_str} ({dur:.1f}s) | Audio: {info['audio_info'][:45]}...")

    total_duration_str = format_timestamp(current_time)
    print(f"\n⏱️ Tổng thời lượng phim: {total_duration_str} ({current_time:.2f} giây)")

    # Xuất YouTube Timestamps ra file text
    yt_file = input_dir / "YOUTUBE_CHAPTERS.txt"
    with open(yt_file, "w", encoding="utf-8") as f:
        f.write(f"=== YOUTUBE CHAPTERS: {args.title} ===\n")
        f.write(f"Tổng thời lượng: {total_duration_str}\n\n")
        f.write("\n".join(yt_chapters) + "\n")
    print(f"📝 Đã tạo file mốc thời gian YouTube: {yt_file}")

    if args.dry_run:
        print("\n[DRY RUN] Đã hoàn thành bước phân tích và tạo danh sách timeline. Không ghép file.")
        return

    # 2. Chuẩn hóa âm thanh & Parameter Sets (Annex B TS)
    temp_dir = input_dir / "temp_normalized"
    temp_dir.mkdir(parents=True, exist_ok=True)
    norm_files = []

    print("\n⚡ [Bước 2/4] Chuẩn hóa đồng bộ Audio & Chèn In-band VPS/SPS/PPS (MPEG-TS):")
    for ch in chapters:
        i = ch["ep"]
        in_f = ch["file"]
        out_f = temp_dir / f"ep{i}_norm.ts"
        norm_files.append(out_f)
        print(f"  -> Chuẩn hóa EP{i} (Video Stream Copy + hevc_mp4toannexb)...", end="", flush=True)
        normalize_segment(ffmpeg_bin, in_f, out_f)
        sz_mb = out_f.stat().st_size / (1024 * 1024)
        print(f" Xong! ({sz_mb:.1f} MB)")

    # 3. Tạo danh sách Concat và chuẩn bị Metadata
    print("\n🔗 [Bước 3/4] Chuẩn bị danh sách Concat và nhúng Metadata Chapters...")
    concat_list_file = temp_dir / "concat_list.txt"
    with open(concat_list_file, "w", encoding="utf-8") as f:
        for nf in norm_files:
            f_escaped = str(nf.resolve()).replace("\\", "/")
            f.write(f"file '{f_escaped}'\n")

    ffmeta_file = temp_dir / "ffmetadata.txt"
    generate_ffmetadata(args.title, chapters, ffmeta_file)

    # 4. Ghép nối bitstream trực tiếp & xuất Master 4K
    print("\n🎬 [Bước 4/4] Ghép nối bitstream trực tiếp & Hoàn thiện file Master 4K...")
    out_file = Path(args.output)
    out_file.parent.mkdir(parents=True, exist_ok=True)

    cmd_concat = [
        ffmpeg_bin, "-y",
        "-f", "concat",
        "-safe", "0",
        "-i", str(concat_list_file),
        "-i", str(ffmeta_file),
        "-map_metadata", "1",
        "-c", "copy",
        "-movflags", "+faststart",
        str(out_file)
    ]
    res_concat = subprocess.run(cmd_concat, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding="utf-8", errors="ignore")
    if res_concat.returncode != 0:
        raise RuntimeError(f"Lỗi khi ghép concat demuxer:\n{res_concat.stderr[-600:]}")
    print("  -> Ghép nối 8 tập & nhúng Chapters hoàn tất không lỗi!")

    # 5. Hàng rào kiểm thử chất lượng tự động (Autonomous Quality Gate)
    print("\n🛡️ [Kiểm chứng] Kiểm tra tính toàn vẹn luồng decode tại tất cả các điểm giao tập:")
    all_clean = True
    for ch in chapters:
        chk_sec = ch["start_ms"] / 1000.0 + 2.0
        cmd_chk = [
            ffmpeg_bin, "-v", "error",
            "-ss", str(chk_sec),
            "-i", str(out_file),
            "-vframes", "1",
            "-f", "null", "-"
        ]
        r_chk = subprocess.run(cmd_chk, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding="utf-8", errors="ignore")
        if r_chk.stderr.strip():
            print(f"  ❌ Cảnh báo tại Tập {ch['ep']} ({format_timestamp(chk_sec)}): {r_chk.stderr.strip()[:100]}")
            all_clean = False
        else:
            print(f"  ✅ Tập {ch['ep']} ({format_timestamp(chk_sec)}): Khung hình chuẩn xác, 0 lỗi decode (CLEAN)")

    if not all_clean:
        print("⚠️ CẢNH BÁO: Phát hiện lỗi decode ở một số phân đoạn!")
    else:
        print("🎉 TOÀN BỘ 8 TẬP ĐÃ VƯỢT QUA KIỂM ĐỊNH DECODE 100% SẠCH SẼ!")

    # Dọn dẹp thư mục tạm
    if not args.keep_temp:
        shutil.rmtree(temp_dir, ignore_errors=True)
        print("\n  -> Đã dọn dẹp các file trung gian.")

    final_size_mb = out_file.stat().st_size / (1024 * 1024)
    print("\n" + "=" * 70)
    print("🎉 XUẤT BẢN THÀNH CÔNG BỘ PHIM HOÀN CHỈNH 4K UHD!")
    print("=" * 70)
    print(f"📁 Đường dẫn Master : {out_file}")
    print(f"📊 Dung lượng       : {final_size_mb:.2f} MB ({final_size_mb/1024:.2f} GB)")
    print(f"⏱️ Tổng thời lượng  : {total_duration_str}")
    print(f"📌 Danh sách Chapter: Có 8 chương phân đoạn đã được nhúng sẵn")
    print(f"📋 Mốc mô tả YouTube: Đã lưu tại {yt_file}")
    print("=" * 70)

if __name__ == "__main__":
    main()