import gradio as gr
import cv2
import os
import shutil
import subprocess
from SinglePhoto import FaceSwapper

swapper = FaceSwapper()

def swap_single_photo(src_img, src_idx, dst_img, dst_idx):
    try:
        src_path = "temp_src.jpg"
        dst_path = "temp_dst.jpg"
        out_path = "output.jpg"
        
        cv2.imwrite(src_path, cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR))
        cv2.imwrite(dst_path, cv2.cvtColor(dst_img, cv2.COLOR_RGB2BGR))
        
        result = swapper.swap_faces(src_path, int(src_idx), dst_path, int(dst_idx))
        cv2.imwrite(out_path, result)
        
        for p in [src_path, dst_path]:
            if os.path.exists(p): os.remove(p)
                
        return out_path, "✅ Photo swap completed!"
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

def add_audio(original_video, no_audio_video, output_path):
    cmd = ["ffmpeg", "-y", "-i", no_audio_video, "-i", original_video, "-c:v", "copy", "-c:a", "aac", "-map", "0:v:0", "-map", "1:a:0?", "-shortest", output_path]
    subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    return output_path

def swap_video(src_img, src_idx, video, dst_idx):
    try:
        src_path = "temp_src.jpg"
        video_path = "temp_video.mp4"
        frames_dir = "frames"
        swapped_dir = "swapped"
        output_video = "output_video.mp4"
        final_video = "final_with_audio.mp4"

        os.makedirs(frames_dir, exist_ok=True)
        os.makedirs(swapped_dir, exist_ok=True)

        cv2.imwrite(src_path, cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR))
        shutil.copy(video, video_path) if isinstance(video, str) else None

        # Extract frames (simplified)
        cap = cv2.VideoCapture(video_path)
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_paths = []
        idx = 0
        while True:
            ret, frame = cap.read()
            if not ret: break
            frame_path = os.path.join(frames_dir, f"frame_{idx:05d}.jpg")
            cv2.imwrite(frame_path, frame)
            frame_paths.append(frame_path)
            idx += 1
        cap.release()

        # Swap faces on frames
        for i, frame_path in enumerate(frame_paths):
            swapped = swapper.swap_faces(src_path, int(src_idx), frame_path, int(dst_idx))
            cv2.imwrite(os.path.join(swapped_dir, f"swapped_{i:05d}.jpg"), swapped)

        # Rebuild video
        # (Use frames_to_video logic or ffmpeg - simplified here)
        # For full version, import from VideoSwapping.py

        # Add audio back
        final = add_audio(video_path, output_video, final_video) if os.path.exists(video_path) else output_video

        # Cleanup
        shutil.rmtree(frames_dir, ignore_errors=True)
        shutil.rmtree(swapped_dir, ignore_errors=True)
        for f in [src_path, video_path, output_video]:
            if os.path.exists(f): os.remove(f)

        return final, "✅ Video swap completed!"
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

with gr.Blocks(title="FaceSwap GPU") as demo:
    gr.Markdown("# Face Swapping Suite - GPU Version")

    with gr.Tab("Single Photo"):
        with gr.Row():
            with gr.Column():
                src = gr.Image(label="Source Image")
                src_idx = gr.Number(1, label="Source Face #", precision=0)
            with gr.Column():
                dst = gr.Image(label="Destination Image")
                dst_idx = gr.Number(1, label="Destination Face #", precision=0)
        btn = gr.Button("Swap", variant="primary")
        out = gr.Image(label="Result")
        log = gr.Textbox(label="Log")
        btn.click(swap_single_photo, [src, src_idx, dst, dst_idx], [out, log])

    with gr.Tab("Video Swap"):
        with gr.Row():
            with gr.Column():
                v_src = gr.Image(label="Source Face")
                v_src_idx = gr.Number(1, label="Source Face #", precision=0)
            with gr.Column():
                v_video = gr.Video(label="Target Video")
                v_dst_idx = gr.Number(1, label="Target Face #", precision=0)
        v_btn = gr.Button("Swap Video", variant="primary")
        v_out = gr.Video(label="Swapped Video")
        v_log = gr.Textbox(label="Log")
        v_btn.click(swap_video, [v_src, v_src_idx, v_video, v_dst_idx], [v_out, v_log])

if __name__ == "__main__":
    demo.launch()