#!/usr/bin/env python3
"""
绿幕抠像脚本 - 三步法流水线
从 Kling Avatar 生成的绿幕视频输出透明 MOV

使用方式：
    python green_screen_keying.py --input talking_head_green.mp4 --output pip_clean.mov

三步法：
    1. ffmpeg chromakey 提取 RGBA 帧
    2. Python PIL 后处理（硬阈值、平滑、Despill）
    3. 组装透明 MOV
"""

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

import numpy as np
from PIL import Image, ImageFilter

# ─────────────────────────── 配置 ───────────────────────────

ALPHA_THRESHOLD = 120  # 硬阈值：消除半透明像素（保留更多发丝）
GAUSSIAN_BLUR_RADIUS = 1  # 轻微平滑（不要用 2，过度模糊）
DESPILL_OFFSET = 5  # 绿色抑制强度（+5 比 +10 更激进）
CHROMAKEY_SIMILARITY = 0.20  # ffmpeg chromakey similarity
CHROMAKEY_BLEND = 0.05  # ffmpeg chromakey blend（保守参数，避免人物透明）
TARGET_FPS = 24

# ─────────────────────────── 辅助函数 ───────────────────────────


def sample_green_color(frame_path: str) -> str:
    """用 PIL 采样实际背景绿色均值（不要写死颜色值）"""
    img = Image.open(frame_path).convert("RGBA")
    arr = np.array(img)

    # 采样顶部区域（通常是纯绿幕背景）
    top = arr[:100, :, :]

    # 找出绿色像素（绿色分量大于红和蓝）
    mask = (
        (top[:, :, 1] > 100)
        & (top[:, :, 1] > top[:, :, 0])
        & (top[:, :, 1] > top[:, :, 2])
    )

    if mask.sum() == 0:
        print("⚠️ 自动采样失败，使用默认绿色 #00FF00")
        return "00FF00"

    avg = top[mask].mean(axis=0).astype(int)
    hex_color = f"{avg[0]:02X}{avg[1]:02X}{avg[2]:02X}"
    print(f"📊 采样绿色均值: #{hex_color} (R={avg[0]}, G={avg[1]}, B={avg[2]})")
    return hex_color


def process_frame(input_path: str, output_path: str) -> None:
    """
    单帧处理：硬阈值 + 轻微平滑 + Despill
    关键参数说明：
        - Alpha 硬阈值 > 120 → 255：消除半透明像素，保留更多发丝
        - 禁止 MinFilter：会严重损伤辫子/卷发等细节发丝
        - GaussianBlur(1)：轻微平滑，不过度模糊
        - Despill：g = min(g, avg(r,b) + 5)：激进绿色抑制
    """
    img = Image.open(input_path).convert("RGBA")
    arr = np.array(img, dtype=np.float32)
    r, g, b, a = arr[:, :, 0], arr[:, :, 1], arr[:, :, 2], arr[:, :, 3]

    # 硬阈值：消除半透明像素
    a_clean = np.where(a > ALPHA_THRESHOLD, 255.0, 0.0)
    alpha_img = Image.fromarray(a_clean.astype(np.uint8), "L")

    # 轻微平滑（避免锯齿，但不腐蚀）
    alpha_smooth = alpha_img.filter(ImageFilter.GaussianBlur(GAUSSIAN_BLUR_RADIUS))

    # 激进 Despill：钳制绿通道
    avg_rb = (r + b) / 2.0
    g_new = np.minimum(g, avg_rb + DESPILL_OFFSET)

    # 组装 RGBA
    result = np.stack(
        [r, g_new, b, np.array(alpha_smooth, dtype=np.float32)], axis=-1
    ).astype(np.uint8)

    Image.fromarray(result, "RGBA").save(output_path)


def run(args: argparse.Namespace) -> None:
    input_path = Path(args.input)
    output_path = Path(args.output)
    temp_dir = Path("_temp_keying")
    frames_dir = temp_dir / "frames"
    clean_dir = temp_dir / "clean_frames"

    # 清理临时目录
    for d in [frames_dir, clean_dir]:
        if d.exists():
            shutil.rmtree(d)
        d.mkdir(parents=True, exist_ok=True)

    # ── Step 1: 提取首帧用于采样绿色 ──
    print("\n📸 Step 1: 提取首帧...")
    subprocess.run(
        [
            "ffmpeg",
            "-i",
            str(input_path),
            "-vframes",
            "1",
            "-q:v",
            "2",
            str(temp_dir / "first_frame.png"),
        ],
        check=True,
        capture_output=True,
    )

    # 采样绿色
    green_hex = sample_green_color(str(temp_dir / "first_frame.png"))

    # ── Step 2: ffmpeg chromakey 提取 RGBA 帧 ──
    print(f"\n🎬 Step 2: ffmpeg chromakey 提取帧（green=#{green_hex}）...")
    subprocess.run(
        [
            "ffmpeg",
            "-i",
            str(input_path),
            "-vf",
            f"chromakey=0x{green_hex}:{CHROMAKEY_SIMILARITY}:{CHROMAKEY_BLEND},"
            f"format=rgba,fps={TARGET_FPS}",
            "-q:v",
            "2",
            str(frames_dir / "frame_%04d.png"),
        ],
        check=True,
        capture_output=True,
    )

    frame_files = sorted(frames_dir.glob("frame_*.png"))
    total = len(frame_files)
    print(f"   共提取 {total} 帧")

    # ── Step 3: Python PIL 后处理 + 组装 ──
    print(
        f"\n🧹 Step 3: Python 后处理（阈值={ALPHA_THRESHOLD}, "
        f"Blur={GAUSSIAN_BLUR_RADIUS}, Despill=+{DESPILL_OFFSET}）..."
    )

    for i, frame_file in enumerate(frame_files, 1):
        if i % 50 == 0 or i == 1 or i == total:
            print(f"   处理帧 {i}/{total}...")
        process_frame(str(frame_file), str(clean_dir / frame_file.name))

    # ── Step 4: 组装透明 MOV ──
    print("\n🎞️ Step 4: 组装透明 MOV...")
    subprocess.run(
        [
            "ffmpeg",
            "-framerate",
            str(TARGET_FPS),
            "-i",
            str(clean_dir / "frame_%04d.png"),
            "-c:v",
            "png",
            "-pix_fmt",
            "rgba",
            str(output_path),
        ],
        check=True,
        capture_output=True,
    )

    # 清理临时文件
    shutil.rmtree(temp_dir)

    print(f"\n✅ 完成！输出：{output_path}")
    duration = total / TARGET_FPS
    print(f"   时长: {duration:.2f}s | 帧数: {total} | 帧率: {TARGET_FPS}fps")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="绿幕抠像三步法流水线")
    parser.add_argument("--input", "-i", required=True, help="输入绿幕视频路径")
    parser.add_argument("--output", "-o", required=True, help="输出透明 MOV 路径")
    run(parser.parse_args())