import os
import cv2
import numpy as np
import gradio as gr
import insightface
import urllib.request
from insightface.app import FaceAnalysis

# --- 0. モデルの自動ダウンロード ---
def download_model_if_missing():
    model_filename = 'inswapper_128.onnx'
    url = "https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx"
    if not os.path.exists(model_filename):
        print("Downloading model...")
        urllib.request.urlretrieve(url, model_filename)

download_model_if_missing()

# --- 1. モデルの初期化 ---
app = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
app.prepare(ctx_id=0, det_size=(640, 640))
swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=False, providers=['CPUExecutionProvider'])

# --- 2. 動画処理ロジック ---
def process_video(source_img, target_video, progress=gr.Progress()):
    if source_img is None or target_video is None:
        return None, "画像と動画を選択してください。"

    # ソース顔の抽出
    source_faces = app.get(source_img)
    if not source_faces:
        return None, "画像から顔が見つかりません。"
    source_face = source_faces[0]

    # 動画の読み込み
    cap = cv2.VideoCapture(target_video)
    if not cap.isOpened():
        return None, "動画ファイルを開けませんでした。"

    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    # 出力ファイル設定 (より汎用的な mp4v を使用)
    output_path = "output_result.mp4"
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    # フレーム処理ループ
    frame_idx = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        # 顔の検出と入れ替え
        faces = app.get(frame)
        res_frame = frame.copy()
        for face in faces:
            res_frame = swapper.get(res_frame, face, source_face, paste_back=True)
        
        out.write(res_frame)
        frame_idx += 1
        
        # 進捗をUIに反映
        if total_frames > 0:
            progress(frame_idx / total_frames, desc=f"処理中... {frame_idx}/{total_frames}")

    cap.release()
    out.release()
    
    # 最終的なファイルの存在確認
    if not os.path.exists(output_path):
        return None, "動画の生成に失敗しました。"

    return output_path, "完了しました！"

# --- 3. UI ---
with gr.Blocks() as demo:
    gr.Markdown("### 🎬 Stable Face Swapper")
    with gr.Row():
        with gr.Column():
            img_input = gr.Image(label="入れ替える顔画像", type="numpy")
            vid_input = gr.Video(label="ターゲット動画")
            submit_btn = gr.Button("変換開始")
        with gr.Column():
            vid_output = gr.Video(label="出力結果")
            status_text = gr.Textbox(label="ステータス")

    submit_btn.click(
        fn=process_video,
        inputs=[img_input, vid_input],
        outputs=[vid_output, status_text]
    )

demo.launch()