# -*- coding: UTF-8 -*-
#!/usr/bin/env python

import os
import json
import shutil
from datetime import datetime
from dotenv import load_dotenv
import numpy as np
import cv2
from PIL import Image
import gradio as gr
from huggingface_hub import HfApi, login
from roop.globals import (
    start,
    decode_execution_providers,
    suggest_max_memory,
    suggest_execution_threads,
)
from roop.core import normalize_output_path
from roop.processors.frame.core import get_frame_processors_modules
from insightface.app import FaceAnalysis

# Load environment variables
load_dotenv()

# Cosine similarity function
def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-6)

# Dataset handler class
class FaceIntegrDataset:
    def __init__(self, repo_id="Arrcttacsrks/face_integrData"):
        self.token = os.getenv("hf_token")
        if not self.token:
            raise ValueError("HF_TOKEN environment variable is not set")
        self.repo_id = repo_id
        self.api = HfApi()
        login(self.token)
        self.temp_dir = "temp_dataset"
        os.makedirs(self.temp_dir, exist_ok=True)

    def create_date_folder(self):
        current_date = datetime.now().strftime("%Y-%m-%d")
        folder_path = os.path.join(self.temp_dir, current_date)
        os.makedirs(folder_path, exist_ok=True)
        return folder_path, current_date

    def save_metadata(self, source_path, target_path, output_path, timestamp):
        metadata = {
            "timestamp": timestamp,
            "source_image": source_path,
            "target_image": target_path,
            "output_image": output_path,
            "date_created": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        }
        return metadata

    def upload_to_hf(self, local_folder, date_folder):
        try:
            self.api.upload_folder(
                folder_path=local_folder,
                repo_id=self.repo_id,
                repo_type="dataset",
                path_in_repo=date_folder,
            )
            return True
        except Exception as e:
            print(f"Error uploading to Hugging Face: {str(e)}")
            return False

# Image face swap function
def swap_face(source_file, target_file, doFaceEnhancer):
    dataset_handler = FaceIntegrDataset()
    folder_path, date_folder = dataset_handler.create_date_folder()
    timestamp = datetime.now().strftime("%S-%M-%H-%d-%m-%Y")

    try:
        # Save source and target images
        source_path = os.path.join(folder_path, f"source_{timestamp}.jpg")
        target_path = os.path.join(folder_path, f"target_{timestamp}.jpg")
        output_path = os.path.join(folder_path, f"OutputImage_{timestamp}.jpg")

        if source_file is None or target_file is None:
            raise ValueError("Source and target images are required")

        Image.fromarray(source_file).save(source_path)
        Image.fromarray(target_file).save(target_path)

        # Configure Roop globals
        roop.globals.source_path = source_path
        roop.globals.target_path = target_path
        roop.globals.output_path = normalize_output_path(
            roop.globals.source_path, roop.globals.target_path, output_path
        )

        roop.globals.frame_processors = (
            ["face_swapper", "face_enhancer"] if doFaceEnhancer else ["face_swapper"]
        )
        roop.globals.headless = True
        roop.globals.keep_fps = True
        roop.globals.keep_audio = True
        roop.globals.keep_frames = False
        roop.globals.many_faces = False
        roop.globals.video_encoder = "libx264"
        roop.globals.video_quality = 18
        roop.globals.max_memory = suggest_max_memory()
        roop.globals.execution_providers = decode_execution_providers(["cuda"])
        roop.globals.execution_threads = suggest_execution_threads()

        # Pre-check frame processors
        for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
            if not frame_processor.pre_check():
                raise RuntimeError("Frame processor pre-check failed")

        # Start face swap process
        start()

        # Save metadata
        metadata = dataset_handler.save_metadata(
            f"source_{timestamp}.jpg",
            f"target_{timestamp}.jpg",
            f"OutputImage_{timestamp}.jpg",
            timestamp,
        )
        metadata_path = os.path.join(folder_path, f"metadata_{timestamp}.json")
        with open(metadata_path, "w") as f:
            json.dump(metadata, f, indent=4)

        # Upload to Hugging Face
        upload_success = dataset_handler.upload_to_hf(folder_path, date_folder)
        if not upload_success:
            print("Failed to upload files to Hugging Face dataset")

        # Return output image
        if os.path.exists(output_path):
            output_image = Image.open(output_path)
            return np.array(output_image)
        else:
            raise FileNotFoundError("Output image not found")

    except Exception as e:
        print(f"Error in face swap process: {str(e)}")
        raise gr.Error(f"Face swap failed: {str(e)}")
    finally:
        if folder_path and os.path.exists(folder_path):
            shutil.rmtree(folder_path)

# Video face swap helper function
def swap_face_frame(frame_bgr, replacement_face_rgb, doFaceEnhancer):
    frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
    temp_dir = "temp_faceswap_frame"
    os.makedirs(temp_dir, exist_ok=True)
    timestamp = datetime.now().strftime("%S-%M-%H-%d-%m-%Y")

    try:
        source_path = os.path.join(temp_dir, f"source_{timestamp}.jpg")
        target_path = os.path.join(temp_dir, f"target_{timestamp}.jpg")
        output_path = os.path.join(temp_dir, f"OutputImage_{timestamp}.jpg")

        Image.fromarray(frame_rgb).save(source_path)
        Image.fromarray(replacement_face_rgb).save(target_path)

        # Configure Roop globals
        roop.globals.source_path = source_path
        roop.globals.target_path = target_path
        roop.globals.output_path = normalize_output_path(
            source_path, target_path, output_path
        )
        roop.globals.frame_processors = (
            ["face_swapper", "face_enhancer"] if doFaceEnhancer else ["face_swapper"]
        )
        roop.globals.headless = True
        roop.globals.keep_fps = True
        roop.globals.keep_audio = True
        roop.globals.keep_frames = False
        roop.globals.many_faces = False
        roop.globals.video_encoder = "libx264"
        roop.globals.video_quality = 18
        roop.globals.max_memory = suggest_max_memory()
        roop.globals.execution_providers = decode_execution_providers(["cuda"])
        roop.globals.execution_threads = suggest_execution_threads()

        start()

        # Return swapped frame
        if os.path.exists(output_path):
            return np.array(Image.open(output_path))
        else:
            return frame_rgb
    finally:
        shutil.rmtree(temp_dir)

# Video face swap function
def swap_face_video(reference_face, replacement_face, video_input, similarity_threshold, doFaceEnhancer):
    fa = FaceAnalysis()
    fa.prepare(ctx_id=0)

    ref_detections = fa.get(reference_face)
    if not ref_detections:
        raise gr.Error("No face detected in the reference image!")
    ref_embedding = ref_detections[0].embedding

    cap = cv2.VideoCapture(video_input)
    if not cap.isOpened():
        raise gr.Error("Failed to open input video!")

    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))

    output_video_path = "temp_faceswap_video.mp4"
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))

    frame_index = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        detections = fa.get(frame_rgb)
        swap_this_frame = any(
            cosine_similarity(det.embedding, ref_embedding) >= similarity_threshold
            for det in detections
        )

        if swap_this_frame:
            swapped_frame_rgb = swap_face_frame(frame, replacement_face, doFaceEnhancer)
            swapped_frame = cv2.cvtColor(swapped_frame_rgb, cv2.COLOR_RGB2BGR)
        else:
            swapped_frame = frame

        out.write(swapped_frame)
        frame_index += 1
        print(f"Processed frame {frame_index}")

    cap.release()
    out.release()
    return output_video_path

# Gradio interface
def create_interface():
    custom_css = """
    .container {
        max-width: 1200px;
        margin: auto;
        padding: 20px;
    }
    .output-image {
        min-height: 400px;
        border: 1px solid #ccc;
        border-radius: 8px;
        padding: 10px;
    }
    """
    title = "Face - Integrator"
    description = "Upload source and target images to perform face swap."
    article = """
    <div style="text-align: center; max-width: 650px; margin: 40px auto;">
        <p>This tool performs face swapping with optional enhancement.</p>
    </div>
    """

    with gr.Blocks(title=title, css=custom_css) as app:
        gr.Markdown(f"<h1 style='text-align: center;'>{title}</h1>")
        gr.Markdown(description)

        with gr.Tabs():
            with gr.TabItem("FaceSwap Image"):
                with gr.Row():
                    source_image = gr.Image(label="Source Image", type="numpy", sources=["upload"])
                    target_image = gr.Image(label="Target Image", type="numpy", sources=["upload"])
                    output_image = gr.Image(label="Output Image", type="numpy", interactive=False, elem_classes="output-image")

                enhance_checkbox = gr.Checkbox(label="Apply Enhancement?", info="Improve image quality", value=False)
                process_btn = gr.Button("Process Face Swap", variant="primary", size="lg")

                process_btn.click(
                    fn=swap_face,
                    inputs=[source_image, target_image, enhance_checkbox],
                    outputs=output_image,
                    api_name="swap_face",
                )

            with gr.TabItem("FaceSwap Video"):
                gr.Markdown("<h2 style='text-align:center;'>FaceSwap Video</h2>")
                with gr.Row():
                    ref_image = gr.Image(label="Reference Face", type="numpy", sources=["upload"])
                    swap_image = gr.Image(label="Replacement Face", type="numpy", sources=["upload"])
                    video_input = gr.Video(label="Input Video")
                    similarity_threshold = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.7, label="Similarity Threshold")
                    enhance_checkbox_video = gr.Checkbox(label="Apply Enhancement?", info="Improve image quality", value=False)
                    video_output = gr.Video(label="Output Video")

                process_video_btn = gr.Button("Process FaceSwap Video", variant="primary", size="lg")

                process_video_btn.click(
                    fn=swap_face_video,
                    inputs=[ref_image, swap_image, video_input, similarity_threshold, enhance_checkbox_video],
                    outputs=video_output,
                    api_name="swap_face_video",
                )

        gr.Markdown(article)

    return app

# Main function
def main():
    app = create_interface()
    app.launch(share=False)

if __name__ == "__main__":
    main()