import gradio as gr
import os
import cv2
import numpy as np
import shutil
from SinglePhoto import FaceSwapper

wellcomingMessage = """
    <h1>Face Swapping Suite</h1>
    <p>All-in-one face swapping: single photo, video, multi-source, and multi-destination!</p>
"""

swapper = FaceSwapper()

def swap_single_photo(src_img, src_idx, dst_img, dst_idx):
    log = ""
    try:
        src_path = "SinglePhoto/data_src.jpg"
        dst_path = "SinglePhoto/data_dst.jpg"
        output_path = "SinglePhoto/output_swapped.jpg"
        os.makedirs(os.path.dirname(src_path), exist_ok=True)
        os.makedirs(os.path.dirname(dst_path), exist_ok=True)
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        src_img_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
        dst_img_bgr = cv2.cvtColor(dst_img, cv2.COLOR_RGB2BGR)
        cv2.imwrite(src_path, src_img_bgr)
        cv2.imwrite(dst_path, dst_img_bgr)
        log += f"Saved source to {src_path}, destination to {dst_path}\n"
        result = swapper.swap_faces(src_path, int(src_idx), dst_path, int(dst_idx))
        cv2.imwrite(output_path, result)
        log += f"Swapped and saved result to {output_path}\n"
        try:
            if os.path.exists(src_path):
                os.remove(src_path)
            if os.path.exists(dst_path):
                os.remove(dst_path)
            log += "Cleaned up temp files.\n"
        except Exception as cleanup_error:
            log += f"Cleanup error: {cleanup_error}\n"
        return output_path, log
    except Exception as e:
        log += f"Error: {e}\n"
        return None, log

def swap_video(src_img, src_idx, video, dst_idx):
    log = ""
    src_path = "VideoSwapping/data_src.jpg"
    dst_video_path = "VideoSwapping/data_dst.mp4"
    frames_dir = "VideoSwapping/video_frames"
    swapped_dir = "VideoSwapping/swapped_frames"
    output_video_path = "VideoSwapping/output_tmp_output_video.mp4"

    os.makedirs(os.path.dirname(src_path), exist_ok=True)
    os.makedirs(os.path.dirname(dst_video_path), exist_ok=True)
    os.makedirs(frames_dir, exist_ok=True)
    os.makedirs(swapped_dir, exist_ok=True)

    src_img_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
    cv2.imwrite(src_path, src_img_bgr)
    log += f"Saved source image to {src_path}\n"

    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

    from VideoSwapping import extract_frames, frames_to_video
    frame_paths = extract_frames(dst_video_path, frames_dir)
    log += f"Extracted {len(frame_paths)} frames to {frames_dir}\n"
    for idx, frame_path in enumerate(frame_paths):
        out_path = os.path.join(swapped_dir, f"swapped_{idx:05d}.jpg")
        try:
            swapped = swapper.swap_faces(src_path, int(src_idx), frame_path, int(dst_idx))
            cv2.imwrite(out_path, swapped)
            log += f"Swapped frame {idx} and saved to {out_path}\n"
        except Exception as e:
            cv2.imwrite(out_path, cv2.imread(frame_path))
            log += f"Failed to swap frame {idx}: {e}\n"
    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"Combined swapped frames into video {output_video_path}\n"

    try:
        if os.path.exists(src_path):
            os.remove(src_path)
        if os.path.exists(dst_video_path):
            os.remove(dst_video_path)
        if os.path.exists(frames_dir):
            shutil.rmtree(frames_dir)
        if os.path.exists(swapped_dir):
            shutil.rmtree(swapped_dir)
        log += "Cleaned up temp files and folders.\n"
    except Exception as cleanup_error:
        log += f"Cleanup error: {cleanup_error}\n"

    return output_video_path, log

def swap_multi_src_single_dst(src_imgs, dst_img, dst_idx):
    log = ""
    results = []
    src_dir = "MultiSrcSingleDst/src"
    dst_dir = "MultiSrcSingleDst/dst"
    output_dir = "MultiSrcSingleDst/output"
    os.makedirs(src_dir, exist_ok=True)
    os.makedirs(dst_dir, exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)

    if isinstance(dst_img, tuple):
        dst_img = dst_img[0]
    dst_img_bgr = cv2.cvtColor(dst_img, cv2.COLOR_RGB2BGR)
    dst_path = os.path.join(dst_dir, "data_dst.jpg")
    cv2.imwrite(dst_path, dst_img_bgr)
    log += f"Saved destination image to {dst_path}\n"

    for i, src_img in enumerate(src_imgs):
        if isinstance(src_img, tuple):
            src_img = src_img[0]
        src_path = os.path.join(src_dir, f"data_src_{i}.jpg")
        output_path = os.path.join(output_dir, f"output_swapped_{i}.jpg")
        src_img_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
        cv2.imwrite(src_path, src_img_bgr)
        log += f"Saved source image {i} to {src_path}\n"
        try:
            result = swapper.swap_faces(src_path, 1, dst_path, int(dst_idx))
            cv2.imwrite(output_path, result)
            results.append(output_path)
            log += f"Swapped and saved result to {output_path}\n"
        except Exception as e:
            results.append(f"Error: {e}")
            log += f"Error swapping source {i}: {e}\n"
    return results, log

def swap_multi_src_multi_dst(src_imgs, dst_imgs, dst_indices):
    log = ""
    results = []
    src_dir = "MultiSrcMultiDst/src"
    dst_dir = "MultiSrcMultiDst/dst"
    output_dir = "MultiSrcMultiDst/output"
    os.makedirs(src_dir, exist_ok=True)
    os.makedirs(dst_dir, exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)

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

    for i, src_img in enumerate(src_imgs):
        if isinstance(src_img, tuple):
            src_img = src_img[0]
        if src_img is None:
            results.append(f"Error: Source image at index {i} is None")
            log += f"Source image at index {i} is None\n"
            continue
        src_path = os.path.join(src_dir, f"data_src_{i}.jpg")
        if isinstance(src_img, np.ndarray):
            src_img_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
            cv2.imwrite(src_path, src_img_bgr)
            log += f"Saved source image {i} to {src_path}\n"
        elif isinstance(src_img, str) and os.path.exists(src_img):
            shutil.copy(src_img, src_path)
            log += f"Copied source image {i} from {src_img} to {src_path}\n"
        else:
            results.append(f"Error: Invalid source image at index {i}")
            log += f"Invalid source image at index {i}\n"
            continue
        for j, dst_img in enumerate(dst_imgs):
            if isinstance(dst_img, tuple):
                dst_img = dst_img[0]
            if dst_img is None:
                results.append(f"Error: Destination image at index {j} is None")
                log += f"Destination image at index {j} is None\n"
                continue
            dst_path = os.path.join(dst_dir, f"data_dst_{j}.jpg")
            output_path = os.path.join(output_dir, f"output_swapped_{i}_{j}.jpg")
            if isinstance(dst_img, np.ndarray):
                dst_img_bgr = cv2.cvtColor(dst_img, cv2.COLOR_RGB2BGR)
                cv2.imwrite(dst_path, dst_img_bgr)
                log += f"Saved destination image {j} to {dst_path}\n"
            elif isinstance(dst_img, str) and os.path.exists(dst_img):
                shutil.copy(dst_img, dst_path)
                log += f"Copied destination image {j} from {dst_img} to {dst_path}\n"
            else:
                results.append(f"Error: Invalid destination image at index {j}")
                log += f"Invalid destination image at index {j}\n"
                continue
            try:
                dst_idx = dst_indices_list[j] if j < len(dst_indices_list) else 1
                result = swapper.swap_faces(src_path, 1, dst_path, int(dst_idx))
                cv2.imwrite(output_path, result)
                results.append(output_path)
                log += f"Swapped src {i} with dst {j} and saved to {output_path}\n"
            except Exception as e:
                results.append(f"Error: {e}")
                log += f"Error swapping src {i} with dst {j}: {e}\n"
    return results, log

def swap_single_src_multi_dst(src_img, dst_imgs, dst_indices):
    log = ""
    results = []
    src_dir = "SingleSrcMultiDst/src"
    dst_dir = "SingleSrcMultiDst/dst"
    output_dir = "SingleSrcMultiDst/output"
    os.makedirs(src_dir, exist_ok=True)
    os.makedirs(dst_dir, exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)

    if isinstance(src_img, tuple):
        src_img = src_img[0]
    src_path = os.path.join(src_dir, "data_src.jpg")
    src_img_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
    cv2.imwrite(src_path, src_img_bgr)
    log += f"Saved source image to {src_path}\n"

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

    for j, dst_img in enumerate(dst_imgs):
        if isinstance(dst_img, tuple):
            dst_img = dst_img[0]
        dst_path = os.path.join(dst_dir, f"data_dst_{j}.jpg")
        output_path = os.path.join(output_dir, f"output_swapped_{j}.jpg")
        dst_img_bgr = cv2.cvtColor(dst_img, cv2.COLOR_RGB2BGR)
        cv2.imwrite(dst_path, dst_img_bgr)
        log += f"Saved destination image {j} to {dst_path}\n"
        try:
            dst_idx = dst_indices_list[j] if j < len(dst_indices_list) else 1
            result = swapper.swap_faces(src_path, 1, dst_path, int(dst_idx))
            cv2.imwrite(output_path, result)
            results.append(output_path)
            log += f"Swapped and saved result to {output_path}\n"
        except Exception as e:
            results.append(f"Error: {e}")
            log += f"Error swapping with destination {j}: {e}\n"
    return results, log

with gr.Blocks() as demo:
    gr.Markdown(wellcomingMessage)
    with gr.Tab("Single Photo Swapping"):
        gr.Interface(
            fn=swap_single_photo,
            inputs=[
                gr.Image(label="Source Image"),
                gr.Number(value=1, label="Source Face Index"),
                gr.Image(label="Destination Image"),
                gr.Number(value=1, label="Destination Face Index"),
            ],
            outputs=[
                gr.Image(label="Swapped Image"),
                gr.Textbox(label="Log Output", lines=8, interactive=False)
            ],
        )
    with gr.Tab("Video Swapping"):
        gr.Interface(
            fn=swap_video,
            inputs=[
                gr.Image(label="Source Image"),
                gr.Number(value=1, label="Source Face Index"),
                gr.Video(label="Target Video"),
                gr.Number(value=1, label="Destination Face Index"),
            ],
            outputs=[
                gr.Video(label="Swapped Video"),
                gr.Textbox(label="Log Output", lines=8, interactive=False)
            ],
        )
    with gr.Tab("SingleSrc MultiDst"):
        gr.Interface(
            fn=swap_single_src_multi_dst,
            inputs=[
                gr.Image(label="Source Image"),
                gr.Gallery(label="Destination Images", type="numpy", columns=3),
                gr.Textbox(label="Destination Face Indices (comma-separated, e.g. 1,1,2)"),
            ],
            outputs=[
                gr.Gallery(label="Swapped Images"),
                gr.Textbox(label="Log Output", lines=8, interactive=False)
            ],
        )
    with gr.Tab("MultiSrc SingleDst"):
        gr.Interface(
            fn=swap_multi_src_single_dst,
            inputs=[
                gr.Gallery(label="Source Images", type="numpy", columns=3),
                gr.Image(label="Destination Image"),
                gr.Number(value=1, label="Destination Face Index"),
            ],
            outputs=[
                gr.Gallery(label="Swapped Images"),
                gr.Textbox(label="Log Output", lines=8, interactive=False)
            ],
        )
    with gr.Tab("MultiSrc MultiDst"):
        gr.Interface(
            fn=swap_multi_src_multi_dst,
            inputs=[
                gr.Gallery(label="Source Images", type="numpy", columns=3),
                gr.Gallery(label="Destination Images", type="numpy", columns=3),
                gr.Textbox(label="Destination Face Indices (comma-separated, e.g. 1,1,2)"),
            ],
            outputs=[
                gr.Gallery(label="Swapped Images"),
                gr.Textbox(label="Log Output", lines=8, interactive=False)
            ],
        )

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