import gradio as gr
import subprocess
from git import Repo
import os
import cv2
import numpy as np
from PIL import Image
import spaces
from diffusers import StableDiffusionImg2ImgPipeline
import torch


device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = None


def load_pipeline():
    global pipe
    if pipe is None:
        pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
            "nroggendorff/epicrealism",
            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
            safety_checker=None,
        )
        pipe = pipe.to(device)
        if hasattr(pipe, "enable_attention_slicing"):
            pipe.enable_attention_slicing()
    return pipe


@spaces.GPU
def refine_with_img2img(image_path, strength=0.3, steps=30, seed=42):
    pipeline = load_pipeline()
    img = Image.open(image_path).convert("RGB")
    generator = torch.Generator(device=device).manual_seed(seed)

    result = pipeline(
        prompt="high quality, detailed, photorealistic, natural texture",
        negative_prompt="blurry, low quality, distorted, deformed, watermark",
        image=img,
        strength=strength,
        num_inference_steps=steps,
        guidance_scale=7.5,
        generator=generator,
    ).images[0]

    result.save(image_path)


@spaces.GPU
def refine_video_with_img2img(
    video_path, strength=0.3, steps=30, seed=42, batch_size=4
):
    pipeline = load_pipeline()
    generator = torch.Generator(device=device).manual_seed(seed)

    cap = cv2.VideoCapture(video_path)
    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))

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

    frames_batch = []
    frame_count = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            if frames_batch:
                for pil_frame in frames_batch:
                    refined = pipeline(
                        prompt="high quality, detailed face, photorealistic, natural skin texture",
                        image=pil_frame,
                        strength=strength,
                        num_inference_steps=steps,
                        guidance_scale=7.5,
                        generator=generator,
                    ).images[0]
                    refined_cv = cv2.cvtColor(np.array(refined), cv2.COLOR_RGB2BGR)
                    out.write(refined_cv)
            break

        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        pil_frame = Image.fromarray(frame_rgb)
        frames_batch.append(pil_frame)

        if len(frames_batch) >= batch_size:
            for pil_frame in frames_batch:
                refined = pipeline(
                    prompt="high quality, detailed face, photorealistic, natural skin texture",
                    image=pil_frame,
                    strength=strength,
                    num_inference_steps=steps,
                    guidance_scale=7.5,
                    generator=generator,
                ).images[0]
                refined_cv = cv2.cvtColor(np.array(refined), cv2.COLOR_RGB2BGR)
                out.write(refined_cv)
                frame_count += 1
                print(f"Processed frame {frame_count}")
            frames_batch = []

    cap.release()
    out.release()
    os.replace(temp_output, video_path)


def denoise_image_gpu(image_path, strength=10):
    img = cv2.imread(image_path)
    img_gpu = cv2.cuda_GpuMat()
    img_gpu.upload(img)

    denoised_gpu = cv2.cuda.fastNlMeansDenoisingColored(
        img_gpu, strength, strength, 7, 21
    )
    denoised = denoised_gpu.download()
    cv2.imwrite(image_path, denoised)


def denoise_image(image_path, strength=10):
    try:
        denoise_image_gpu(image_path, strength)
    except:
        img = cv2.imread(image_path)
        denoised = cv2.fastNlMeansDenoisingColored(img, None, strength, strength, 7, 21)
        cv2.imwrite(image_path, denoised)


def denoise_video(video_path, strength=10):
    cap = cv2.VideoCapture(video_path)
    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))

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

    use_gpu = False
    try:
        test_gpu = cv2.cuda_GpuMat()
        use_gpu = True
    except:
        pass

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

        if use_gpu:
            try:
                frame_gpu = cv2.cuda_GpuMat()
                frame_gpu.upload(frame)
                denoised_gpu = cv2.cuda.fastNlMeansDenoisingColored(
                    frame_gpu, strength, strength, 7, 21
                )
                denoised_frame = denoised_gpu.download()
            except:
                denoised_frame = cv2.fastNlMeansDenoisingColored(
                    frame, None, strength, strength, 7, 21
                )
        else:
            denoised_frame = cv2.fastNlMeansDenoisingColored(
                frame, None, strength, strength, 7, 21
            )

        out.write(denoised_frame)

    cap.release()
    out.release()
    os.replace(temp_output, video_path)


def enhance_image(image_path):
    img = cv2.imread(image_path)
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    l = clahe.apply(l)
    enhanced = cv2.merge([l, a, b])
    enhanced = cv2.cvtColor(enhanced, cv2.COLOR_LAB2BGR)
    cv2.imwrite(image_path, enhanced)


@spaces.GPU
def process_media(
    image,
    image_or_video,
    denoise_strength,
    enhance,
    use_img2img,
    img2img_strength,
    img2img_steps,
    seed,
    execution_provider,
):
    if os.path.exists("output_video.mp4"):
        os.remove("output_video.mp4")
    if os.path.exists("output_image.png"):
        os.remove("output_image.png")
    if os.path.exists("source.png"):
        os.remove("source.png")

    image_output = None
    video_output = None

    if isinstance(image_or_video, str) and image_or_video.endswith(
        (".mp4", ".avi", ".mov")
    ):
        image.save("source.png")
        cmd = f"python3 roop/run.py -s source.png -t '{image_or_video}' -o output_video.mp4 --execution-provider {execution_provider}"
        subprocess.run(cmd, shell=True)

        if os.path.exists("output_video.mp4"):
            if use_img2img:
                refine_video_with_img2img(
                    "output_video.mp4", img2img_strength, img2img_steps, seed
                )
            denoise_video("output_video.mp4", denoise_strength)
            video_output = gr.Video(value="output_video.mp4", visible=True)

    elif isinstance(image_or_video, str) and image_or_video.endswith(
        (".png", ".jpg", ".jpeg")
    ):
        image.save("source.png")
        cmd = f"python3 roop/run.py -s source.png -t '{image_or_video}' -o output_image.png --execution-provider {execution_provider}"
        subprocess.run(cmd, shell=True)

        if os.path.exists("output_image.png"):
            if use_img2img:
                refine_with_img2img(
                    "output_image.png", img2img_strength, img2img_steps, seed
                )
            denoise_image("output_image.png", denoise_strength)
            if enhance:
                enhance_image("output_image.png")
            image_output = gr.Image(value="output_image.png", visible=True)

    return image_output, video_output


with gr.Blocks() as demo:
    with gr.Row():
        image = gr.Image(label="Image", type="pil")
        image_or_video = gr.File(label="Image or Video", type="filepath")

    with gr.Row():
        denoise_strength = gr.Slider(
            minimum=0, maximum=30, value=1, step=0.5, label="Denoise Strength"
        )
        enhance = gr.Checkbox(label="Enhance Contrast (Images Only)", value=True)

    with gr.Row():
        use_img2img = gr.Checkbox(
            label="Use EpicRealism Img2Img Refinement", value=False
        )
        img2img_strength = gr.Slider(
            minimum=0.1, maximum=0.8, value=0.15, step=0.05, label="Img2Img Strength"
        )
        img2img_steps = gr.Slider(
            minimum=10, maximum=50, value=20, step=5, label="Img2Img Steps"
        )
        seed = gr.Number(label="Seed", value=42, precision=0)

    with gr.Row():
        execution_provider = gr.Radio(
            choices=["cuda", "tensorrt"], value="cuda", label="Roop Execution Provider"
        )

    process_btn = gr.Button("Process")

    image_output = gr.Image(label="Output Image", visible=False)
    video_output = gr.Video(label="Output Video", visible=False)

    process_btn.click(
        fn=process_media,
        inputs=[
            image,
            image_or_video,
            denoise_strength,
            enhance,
            use_img2img,
            img2img_strength,
            img2img_steps,
            seed,
            execution_provider,
        ],
        outputs=[image_output, video_output],
    )

demo.queue()

if __name__ == "__main__":
    if not os.path.exists("roop"):
        Repo.clone_from("https://github.com/nroggendorff/roop.git", "roop")
        subprocess.run("pip install -r roop/requirements.txt", shell=True)

    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
    os.environ["TF_USE_LEGACY_KERAS"] = "1"
    os.environ["OMP_NUM_THREADS"] = "8"
    os.environ["MKL_NUM_THREADS"] = "8"

    demo.launch()