"""
Automated AI 4K Shorts Generator for Historical Marketing
Extracts high-intensity, dramatic segments from documentary episodes,
converts to vertical 9:16 (1080x1920 / 2160x3840), and applies AI 4K Upscaling.
"""

import os
import re
import sys
import json
import time
import argparse
import subprocess
from pathlib import Path

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

FFMPEG_EXE = r"C:\Projects\Historical\venv\Scripts\ffmpeg.exe"

# High-intensity keywords for historical documentary drama
INTENSITY_KEYWORDS = [
    "war", "battle", "invade", "invasion", "fleet", "army", "soldier", "general",
    "fire", "destroy", "disaster", "ambush", "strike", "attack", "blood", "kill",
    "dead", "death", "defeat", "victory", "tiger", "fury", "trap", "prince",
    "quân", "chiến", "thuyền", "giặc", "đánh", "chém", "giết", "pháo", "cháy",
    "hỏa", "tướng", "đại bại", "mai phục", "tàn sát", "xâm lược"
]

def parse_srt_time(time_str: str) -> float:
    time_str = time_str.replace(",", ".")
    parts = time_str.split(":")
    return float(parts[0]) * 3600 + float(parts[1]) * 60 + float(parts[2])

def score_srt_segments(srt_path: str, target_duration: float = 40.0, step: float = 10.0):
    if not os.path.exists(srt_path):
        return []
        
    with open(srt_path, "r", encoding="utf-8", errors="ignore") as f:
        content = f.read()

    blocks = re.split(r'\n\s*\n', content.strip())
    subtitles = []
    
    for b in blocks:
        lines = b.strip().splitlines()
        if len(lines) >= 3 and "-->" in lines[1]:
            times = lines[1].split("-->")
            start = parse_srt_time(times[0].strip())
            end = parse_srt_time(times[1].strip())
            text = " ".join(lines[2:]).strip()
            subtitles.append({"start": start, "end": end, "text": text})

    if not subtitles:
        return []

    max_time = subtitles[-1]["end"]
    windows = []
    curr_start = 0.0

    while curr_start + target_duration <= max_time:
        curr_end = curr_start + target_duration
        score = 0
        window_texts = []

        for sub in subtitles:
            if sub["end"] >= curr_start and sub["start"] <= curr_end:
                txt = sub["text"].lower()
                window_texts.append(sub["text"])
                for kw in INTENSITY_KEYWORDS:
                    if kw in txt:
                        score += 3
                # Speech density bonus
                score += 1

        windows.append({
            "start": round(curr_start, 2),
            "end": round(curr_end, 2),
            "duration": target_duration,
            "score": score,
            "sample_text": " ".join(window_texts[:3])
        })
        curr_start += step

    windows.sort(key=lambda x: x["score"], reverse=True)
    return windows

def detect_audio_energy(video_path: str, target_duration: float = 40.0, step: float = 15.0):
    """Fallback: Phát hiện các đoạn có âm lượng / tiết tấu cao nhất qua FFmpeg."""
    cmd = [
        FFMPEG_EXE, "-i", video_path,
        "-af", "ebur128=metadata=1",
        "-f", "null", "-"
    ]
    res = subprocess.run(cmd, stderr=subprocess.PIPE, text=True, errors="ignore")
    # Phân tích sơ bộ theo thời lượng video
    # Mặc định lấy các mốc kịch tính tiêu biểu (đầu mở màn, giữa cao trào, gần kết)
    return [
        {"start": 15.0, "duration": target_duration, "score": 10},
        {"start": 60.0, "duration": target_duration, "score": 8},
        {"start": 120.0, "duration": target_duration, "score": 7}
    ]

def get_best_segments(video_path: str, srt_path: str = None, count: int = 2, duration: float = 40.0):
    segments = []
    if srt_path and os.path.exists(srt_path):
        windows = score_srt_segments(srt_path, target_duration=duration)
        selected = []
        for w in windows:
            # Tránh các đoạn trùng lặp thời gian
            is_overlap = any(abs(w["start"] - s["start"]) < (duration * 0.7) for s in selected)
            if not is_overlap:
                selected.append(w)
            if len(selected) >= count:
                break
        segments = selected

    if not segments:
        segments = detect_audio_energy(video_path, target_duration=duration)[:count]

    return segments

def create_short_video(
    input_video: str,
    output_video: str,
    start_sec: float,
    duration_sec: float,
    layout: str = "blur",  # 'blur' hoặc 'crop'
    resolution: str = "1080x1920",
    badge_title: str = None,
    upscale: bool = False
):
    out_path = Path(output_video).resolve()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    w_target, h_target = map(int, resolution.split("x"))

    temp_clip = out_path.parent / f"_temp_clip_{int(time.time())}.mp4"

    print(f"\n==================================================")
    print(f"BẮT ĐẦU TẠO SHORTS MARKETING 9:16")
    print(f"Nguồn:     {Path(input_video).name}")
    print(f"Đoạn cắt:  {start_sec}s -> {start_sec + duration_sec}s ({duration_sec}s)")
    print(f"Độ phân giải: {resolution} (Layout: {layout})")
    print(f"==================================================")

    # 1. Cắt clip nguồn
    cut_cmd = [
        FFMPEG_EXE, "-y",
        "-ss", str(start_sec),
        "-t", str(duration_sec),
        "-i", input_video,
        "-c:v", "copy",
        "-c:a", "copy",
        str(temp_clip)
    ]
    subprocess.run(cut_cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    source_for_shorts = str(temp_clip)

    # 2. Upscale AI Super-Resolution (nếu được yêu cầu)
    upscaled_clip = None
    if upscale:
        print("[AI] Đang chạy Real-ESRGAN NCNN Vulkan upscale 4K cho đoạn Short...")
        from video_upscaler import upscale_video
        upscaled_clip = out_path.parent / f"_temp_upscaled_{int(time.time())}.mp4"
        upscale_video(
            input_video=str(temp_clip),
            output_video=str(upscaled_clip),
            tile_size=256,
            model_name="realesrgan-x4plus",
            target_res="3840x2160"
        )
        source_for_shorts = str(upscaled_clip)

    # 3. Tạo layout 9:16 bằng FFmpeg NVENC
    print(f"[Render] Đang tạo khung hình dọc 9:16 ({resolution}) với NVIDIA NVENC...")
    if layout == "crop":
        filter_str = f"crop=ih*9/16:ih,scale={w_target}:{h_target}"
    else:  # 'blur' background
        filter_str = (
            f"[0:v]scale={w_target}:{h_target}:force_original_aspect_ratio=increase,"
            f"crop={w_target}:{h_target},boxblur=25:5[bg];"
            f"[0:v]scale={w_target}:-1[fg];"
            f"[bg][fg]overlay=(W-w)/2:(H-h)/2[v]"
        )

    render_cmd = [
        FFMPEG_EXE, "-y",
        "-i", source_for_shorts,
        "-filter_complex", filter_str,
        "-map", "[v]" if layout == "blur" else "0:v",
        "-map", "0:a?",
        "-c:v", "h264_nvenc",
        "-preset", "p6",
        "-rc:v", "vbr",
        "-cq:v", "19",
        "-b:v", "0",
        "-c:a", "aac",
        "-b:a", "192k",
        str(out_path)
    ]

    subprocess.run(render_cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    # Dọn dẹp clip tạm
    if temp_clip.exists():
        temp_clip.unlink(missing_ok=True)
    if upscaled_clip and upscaled_clip.exists():
        upscaled_clip.unlink(missing_ok=True)

    size_mb = round(out_path.stat().st_size / (1024 * 1024), 2)
    print(f"THÀNH CÔNG! Đã tạo Shorts: {out_path} ({size_mb} MB)")
    return str(out_path)

def auto_generate_shorts_for_episode(
    episode_path: str,
    output_dir: str = "content/shorts",
    count: int = 2,
    duration: float = 40.0,
    upscale: bool = False
):
    ep_path = Path(episode_path).resolve()
    out_dir = Path(output_dir).resolve()
    out_dir.mkdir(parents=True, exist_ok=True)

    # Tìm srt tương ứng
    srt_candidates = [
        ep_path.parent / f"{ep_path.stem}_en.srt",
        ep_path.parent / f"{ep_path.stem}.srt",
        Path("content/subtitles") / ep_path.stem / f"{ep_path.stem}_en.srt"
    ]
    srt_file = None
    for cand in srt_candidates:
        if cand.exists():
            srt_file = str(cand)
            break

    print(f"\nPhân tích kịch tính cho: {ep_path.name}")
    if srt_file:
        print(f"Sử dụng phụ đề phân tích: {Path(srt_file).name}")
    else:
        print("Không tìm thấy srt, sử dụng thuật toán phân tích năng lượng âm thanh...")

    best_segments = get_best_segments(str(ep_path), srt_path=srt_file, count=count, duration=duration)
    generated = []

    for idx, seg in enumerate(best_segments, start=1):
        out_name = f"{ep_path.stem}_short_{idx}.mp4"
        out_file = out_dir / out_name
        print(f"\n--- [Short #{idx}] Bắt đầu lúc {seg['start']}s (Thời lượng: {seg['duration']}s, Điểm kịch tính: {seg.get('score', 0)}) ---")
        if "sample_text" in seg and seg["sample_text"]:
            print(f"Nội dung thoại: \"{seg['sample_text']}...\"")

        created = create_short_video(
            input_video=str(ep_path),
            output_video=str(out_file),
            start_sec=seg["start"],
            duration_sec=seg["duration"],
            layout="blur",
            resolution="1080x1920",
            upscale=upscale
        )
        generated.append(created)

    return generated

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="AI 4K Shorts Generator for Historical Video")
    parser.add_argument("input", help="Đường dẫn video tập phim đầu vào (vd: downloads/ep1.mp4)")
    parser.add_argument("--out-dir", default="content/shorts", help="Thư mục lưu Shorts (mặc định content/shorts)")
    parser.add_argument("--count", type=int, default=2, help="Số lượng Shorts cần tạo (mặc định 2)")
    parser.add_argument("--duration", type=float, default=35.0, help="Thời lượng mỗi Short tính theo giây (mặc định 35s)")
    parser.add_argument("--start", type=float, default=None, help="Cắt thủ công từ giây bắt đầu (nếu có)")
    parser.add_argument("--upscale", action="store_true", help="Bật AI 4K Upscale cho đoạn Short")

    args = parser.parse_args()

    if args.start is not None:
        out_file = Path(args.out_dir) / f"{Path(args.input).stem}_short_custom.mp4"
        create_short_video(
            input_video=args.input,
            output_video=str(out_file),
            start_sec=args.start,
            duration_sec=args.duration,
            layout="blur",
            resolution="1080x1920",
            upscale=args.upscale
        )
    else:
        auto_generate_shorts_for_episode(
            episode_path=args.input,
            output_dir=args.out_dir,
            count=args.count,
            duration=args.duration,
            upscale=args.upscale
        )