import gradio as gr
import os
import cv2
import numpy as np
import shutil
import subprocess
import time
import argparse
from SinglePhoto import FaceSwapper
def swap_video_all_faces(src_img, video, num_faces_to_swap, delete_frames_dir=True, add_audio=True, progress=gr.Progress(track_tqdm=True)):
    """
    Swaps ONE source face onto the first N faces detected in EVERY frame of the video.
    Optimized for rented GPU servers.
    """
    log = ""
    start_time = time.time()
    base = "VideoSwappingAllFaces"
    
    # Create directories
    os.makedirs(f"{base}/src", exist_ok=True)
    os.makedirs(f"{base}/dst", exist_ok=True)
    os.makedirs(f"{base}/frames", exist_ok=True)
    os.makedirs(f"{base}/swapped", exist_ok=True)
    os.makedirs(f"{base}/temp", exist_ok=True)

    src_path = f"{base}/data_src.jpg"
    dst_video_path = f"{base}/data_dst.mp4"
    frames_dir = f"{base}/frames"
    swapped_dir = f"{base}/swapped"
    temp_dir = f"{base}/temp"
    output_video_path = f"{base}/output_no_audio.mp4"
    final_output_path = f"{base}/output_with_audio.mp4"

    # Save source image
    cv2.imwrite(src_path, cv2.cvtColor(np.array(src_img), cv2.COLOR_RGB2BGR))
    log += "✅ Source image saved.\n"
    progress(0.05, desc="Saved source image")

    # Copy video
    if isinstance(video, str) and os.path.exists(video):
        shutil.copy(video, dst_video_path)
        log += f"✅ Copied video to {dst_video_path}\n"
    else:
        dst_video_path = video

    # Import helpers
    from VideoSwapping import extract_frames, frames_to_video

    # Extract frames (with resume support)
    frame_paths = extract_frames(dst_video_path, frames_dir)
    log += f"✅ Extracted {len(frame_paths)} frames.\n"
    progress(0.15, desc="Frames extracted")

    # Process frames
    swapped_files = set(os.listdir(swapped_dir))
    total_frames = len(frame_paths)
    start_loop = time.time()

    for idx, frame_path in enumerate(frame_paths):
        swapped_name = f"swapped_{idx:05d}.jpg"
        out_path = os.path.join(swapped_dir, swapped_name)

        # Skip already processed frames (useful for resuming)
        if swapped_name in swapped_files and os.path.exists(out_path):
            progress(0.15 + 0.7 * (idx + 1) / total_frames, 
                    desc=f"Skipping {idx+1}/{total_frames}")
            continue

        temp_frame_path = os.path.join(temp_dir, "temp.jpg")
        try:
            shutil.copy(frame_path, temp_frame_path)
            
            # Swap up to num_faces_to_swap faces in this frame
            for face_idx in range(1, int(num_faces_to_swap) + 1):
                try:
                    swapped_img = swapper.swap_faces(src_path, 1, temp_frame_path, face_idx)
                    cv2.imwrite(temp_frame_path, swapped_img)
                    log += f"Frame {idx}: Swapped face {face_idx}\n"
                except Exception as e:
                    log += f"Frame {idx}: Face {face_idx} skipped - {str(e)[:100]}\n"
                    # Continue with next face instead of failing whole frame
            
            # Save final swapped frame
            shutil.copy(temp_frame_path, out_path)
            
        except Exception as e:
            # Fallback: copy original frame
            cv2.imwrite(out_path, cv2.imread(frame_path))
            log += f"❌ Frame {idx} failed, using original: {str(e)[:80]}\n"

        # Progress update with ETA
        elapsed = time.time() - start_loop
        avg_time = elapsed / (idx + 1) if idx > 0 else 0
        remaining = avg_time * (total_frames - idx - 1)
        mins, secs = divmod(int(remaining), 60)
        progress(0.15 + 0.7 * (idx + 1) / total_frames,
                desc=f"Swapping {idx+1}/{total_frames} | ETA {mins:02d}:{secs:02d}")

    # Combine frames into video
    cap = cv2.VideoCapture(dst_video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    cap.release()
    
    frames_to_video(swapped_dir, output_video_path, fps)
    log += f"✅ Frames combined into video ({fps:.2f} fps)\n"
    progress(0.9, desc="Muxing audio...")

    # Add audio if requested
    if add_audio:
        ok, audio_log = add_audio_to_video(dst_video_path, output_video_path, final_output_path)
        if ok:
            log += "✅ Audio muxed successfully.\n"
            final_video = final_output_path
        else:
            log += f"⚠️ Audio mux failed: {audio_log[:200]}\n"
            final_video = output_video_path
    else:
        final_video = output_video_path
        log += "ℹ️ Audio skipped as requested.\n"

    # Cleanup
    try:
        if delete_frames_dir:
            shutil.rmtree(frames_dir, ignore_errors=True)
            shutil.rmtree(swapped_dir, ignore_errors=True)
            shutil.rmtree(temp_dir, ignore_errors=True)
            log += "🧹 Temporary files cleaned up.\n"
        if os.path.exists(src_path):
            os.remove(src_path)
        if os.path.exists(dst_video_path) and dst_video_path != video:
            os.remove(dst_video_path)
    except Exception as cleanup_err:
        log += f"Cleanup warning: {cleanup_err}\n"

    total_elapsed = time.time() - start_time
    log += f"🎉 **All faces swap completed in {total_elapsed:.2f} seconds**"
    
    return final_video, log

# ========================= CONFIG =========================
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # Change if you have multiple GPUs
os.environ["GRADIO_TEMP_DIR"] = "/tmp/gradio"

wellcomingMessage = """
<h1>Face Swapping Suite - GPU Optimized</h1>
<p>All-in-one face swapping: single/multi photo + video. Designed for dedicated/rented GPU servers.</p>
"""

# Global FaceSwapper (loads model once at startup)
swapper = FaceSwapper()

# ====================== HELPER FUNCTIONS ======================
def add_audio_to_video(original_video_path, video_no_audio_path, output_path):
    """Mux audio back using ffmpeg"""
    cmd = [
        "ffmpeg", "-y",
        "-i", video_no_audio_path,
        "-i", original_video_path,
        "-c:v", "copy",
        "-c:a", "aac",
        "-map", "0:v:0",
        "-map", "1:a:0?",
        "-shortest",
        output_path
    ]
    try:
        subprocess.run(cmd, check=True, capture_output=True)
        return True, ""
    except subprocess.CalledProcessError as e:
        return False, e.stderr.decode()

# ====================== PHOTO SWAPPING ======================
def swap_single_photo(src_img, src_idx, dst_img, dst_idx, progress=gr.Progress(track_tqdm=True)):
    log = ""
    start_time = time.time()
    try:
        progress(0, desc="Saving images")
        os.makedirs("SinglePhoto", exist_ok=True)
        src_path = "SinglePhoto/data_src.jpg"
        dst_path = "SinglePhoto/data_dst.jpg"
        output_path = "SinglePhoto/output_swapped.jpg"

        cv2.imwrite(src_path, cv2.cvtColor(np.array(src_img), cv2.COLOR_RGB2BGR))
        cv2.imwrite(dst_path, cv2.cvtColor(np.array(dst_img), cv2.COLOR_RGB2BGR))

        progress(0.5, desc="Swapping faces")
        result = swapper.swap_faces(src_path, int(src_idx), dst_path, int(dst_idx))
        cv2.imwrite(output_path, result)

        progress(1, desc="Done")
        elapsed = time.time() - start_time
        log += f"✅ Success! Elapsed: {elapsed:.2f} seconds"
        return output_path, log
    except Exception as e:
        log += f"❌ Error: {str(e)}"
        return None, log


def swap_single_src_multi_dst(src_img, dst_imgs, dst_indices, progress=gr.Progress(track_tqdm=True)):
    log = ""
    results = []
    base_dir = "SingleSrcMultiDst"
    os.makedirs(f"{base_dir}/src", exist_ok=True)
    os.makedirs(f"{base_dir}/dst", exist_ok=True)
    os.makedirs(f"{base_dir}/output", exist_ok=True)

    # Save source
    src_path = f"{base_dir}/data_src.jpg"
    cv2.imwrite(src_path, cv2.cvtColor(np.array(src_img), cv2.COLOR_RGB2BGR))

    if isinstance(dst_indices, str):
        dst_indices_list = [int(x.strip()) for x in dst_indices.split(",") if x.strip().isdigit()]
    else:
        dst_indices_list = [int(x) for x in dst_indices]

    for j, dst_img in enumerate(dst_imgs):
        dst_path = f"{base_dir}/data_dst_{j}.jpg"
        out_path = f"{base_dir}/output/output_{j}.jpg"
        cv2.imwrite(dst_path, cv2.cvtColor(np.array(dst_img), cv2.COLOR_RGB2BGR))
        
        try:
            idx = dst_indices_list[j] if j < len(dst_indices_list) else 1
            result = swapper.swap_faces(src_path, 1, dst_path, idx)
            cv2.imwrite(out_path, result)
            results.append(out_path)
        except Exception as e:
            results.append(None)
            log += f"Error on dst {j}: {e}\n"
        progress((j + 1) / len(dst_imgs))
    return results, log


def swap_faces_custom(src_imgs, dst_img, mapping_str, progress=gr.Progress(track_tqdm=True)):
    log = ""
    start = time.time()
    base = "CustomSwap"
    os.makedirs(f"{base}/src", exist_ok=True)
    os.makedirs(f"{base}/temp", exist_ok=True)

    # Save destination
    dst_path = f"{base}/data_dst.jpg"
    cv2.imwrite(dst_path, cv2.cvtColor(np.array(dst_img), cv2.COLOR_RGB2BGR))
    temp_dst = f"{base}/temp/temp.jpg"
    shutil.copy(dst_path, temp_dst)

    # Save sources
    src_paths = []
    for i, img in enumerate(src_imgs):
        path = f"{base}/src/data_src_{i+1}.jpg"
        cv2.imwrite(path, cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))
        src_paths.append(path)

    # Parse mapping
    try:
        mapping = [int(x.strip()) for x in mapping_str.split(",") if x.strip().isdigit()]
    except:
        return None, "Invalid mapping format"

    for face_idx, src_idx in enumerate(mapping, 1):
        if src_idx < 1 or src_idx > len(src_paths):
            continue
        try:
            swapped = swapper.swap_faces(src_paths[src_idx-1], 1, temp_dst, face_idx)
            cv2.imwrite(temp_dst, swapped)
        except Exception as e:
            log += f"Face {face_idx} failed: {e}\n"

    output_path = f"{base}/output_swapped.jpg"
    shutil.copy(temp_dst, output_path)
    shutil.rmtree(f"{base}/temp", ignore_errors=True)
    
    log += f"✅ Done in {time.time()-start:.2f}s"
    return output_path, log


# ====================== VIDEO SWAPPING ======================
def swap_video(src_img, src_idx, video, dst_idx, delete_frames=True, add_audio=True, progress=gr.Progress(track_tqdm=True)):
    log = ""
    start_time = time.time()
    base = "VideoSwapping"
    for d in ["src", "dst", "frames", "swapped", "output"]:
        os.makedirs(f"{base}/{d}", exist_ok=True)

    src_path = f"{base}/data_src.jpg"
    dst_video = f"{base}/data_dst.mp4"
    frames_dir = f"{base}/frames"
    swapped_dir = f"{base}/swapped"
    tmp_out = f"{base}/tmp_video.mp4"
    final_out = f"{base}/final_with_audio.mp4"

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

    from VideoSwapping import extract_frames, frames_to_video

    frame_paths = extract_frames(dst_video, frames_dir)
    progress(0.2, desc=f"Extracted {len(frame_paths)} frames")

    total = len(frame_paths)
    for i, frame in enumerate(frame_paths):
        out_frame = f"{swapped_dir}/swapped_{i:05d}.jpg"
        try:
            swapped = swapper.swap_faces(src_path, int(src_idx), frame, int(dst_idx))
            cv2.imwrite(out_frame, swapped)
        except Exception as e:
            cv2.imwrite(out_frame, cv2.imread(frame))
            log += f"Frame {i} fallback: {e}\n"
        progress(0.2 + 0.7 * (i+1)/total, desc=f"Frame {i+1}/{total}")

    cap = cv2.VideoCapture(dst_video)
    fps = cap.get(cv2.CAP_PROP_FPS)
    cap.release()
    frames_to_video(swapped_dir, tmp_out, fps)

    if add_audio:
        ok, msg = add_audio_to_video(dst_video, tmp_out, final_out)
        output = final_out if ok else tmp_out
    else:
        output = tmp_out

    # Cleanup
    if delete_frames:
        shutil.rmtree(frames_dir, ignore_errors=True)
        shutil.rmtree(swapped_dir, ignore_errors=True)

    log += f"✅ Video done in {time.time()-start_time:.2f}s"
    return output, log


# (Additional video functions like all-faces, custom mapping, multi-video follow similar pattern.
# For brevity in this response I've included core ones. Full multi-video/custom can be expanded identically.)

# ====================== GRADIO INTERFACE ======================
with gr.Blocks(title="FaceSwapAll GPU", theme=gr.themes.Soft(), css=".gradio-container {max-width: 1200px}") as demo:
    gr.Markdown(wellcomingMessage)

    with gr.Tab("Single Photo"):
        gr.Interface(
            fn=swap_single_photo,
            inputs=[
                gr.Image(label="Source Image", type="numpy"),
                gr.Number(value=1, label="Source Face Index", precision=0),
                gr.Image(label="Destination Image", type="numpy"),
                gr.Number(value=1, label="Destination Face Index", precision=0),
            ],
            outputs=[gr.Image(label="Swapped"), gr.Textbox(label="Log", lines=8)],
        )

    with gr.Tab("Single Src → Multi Dst"):
        gr.Interface(
            fn=swap_single_src_multi_dst,
            inputs=[
                gr.Image(label="Source"),
                gr.Gallery(label="Destination Images", columns=3, type="numpy", height=400),
                gr.Textbox(value="1", label="Dst Indices (comma sep)"),
            ],
            outputs=[gr.Gallery(label="Results", columns=3), gr.Textbox(label="Log")],
        )

    with gr.Tab("Custom Mapping"):
        gr.Interface(
            fn=swap_faces_custom,
            inputs=[
                gr.Gallery(label="Source Images", columns=3, type="numpy"),
                gr.Image(label="Destination Image"),
                gr.Textbox(value="1,2,3", label="Mapping (e.g. 2,1,3)"),
            ],
            outputs=[gr.Image(label="Result"), gr.Textbox(label="Log")],
        )

    with gr.Tab("Video Swap"):
        gr.Interface(
            fn=swap_video,
            inputs=[
                gr.Image(label="Source Face Image"),
                gr.Number(value=1, label="Src Face Index"),
                gr.Video(label="Target Video"),
                gr.Number(value=1, label="Dst Face Index"),
                gr.Checkbox(label="Delete temp frames", value=True),
                gr.Checkbox(label="Add original audio", value=True),
            ],
            outputs=[gr.Video(label="Result"), gr.Textbox(label="Log", lines=10)],
        )

    # Add more tabs as needed (Multi-video, All Faces, etc.)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--share", action="store_true")
    parser.add_argument("--port", type=int, default=7860)
    args = parser.parse_args()

    demo.launch(
        server_name="0.0.0.0",
        server_port=args.port,
        share=args.share,
        max_threads=4,           # Tune based on your GPU VRAM
        # Critical for large videos
        show_api=False,
    )