import os
import warnings
from pathlib import Path

import gradio as gr
from gradio_imageslider import ImageSlider  # Import ImageSlider
import torch
import fastai
from deoldify import device
from deoldify.device_id import DeviceId
from deoldify.visualize import *
from huggingface_hub import HfApi

# Suppress warnings
warnings.filterwarnings("ignore", category=UserWarning, message=".*?Your .*? set is empty.*?")

# Hugging Face model repo details
repo_id = "afondiel/image-colorizer-deoldify"
repo_type = "space"

# Set device
device.set(device=DeviceId.CPU)  # Use CPU first
api = HfApi()
snapshot_folder = api.snapshot_download(repo_id=repo_id, repo_type=repo_type)
device.set(device=DeviceId.GPU0)  # If GPU available

# Load model
_colorizer = get_image_colorizer(root_folder=Path(snapshot_folder), artistic=True)

def colorizer_fn(input_img):
    """Colorize grayscale images with fixed render factor 40."""
    if input_img:
        output_img = _colorizer.get_transformed_image(
            path=input_img,
            render_factor=40,  # Fixed render factor
            watermarked=False,
            post_process=True,
        )
        return (input_img, output_img)  # Return tuple for ImageSlider
    else:
        return None

# ✅ Custom CSS for simple white "Colorizer" heading
custom_css = """
body, .gradio-container {
    background-color: black !important;       /* Black background */
    color: white !important;                  /* White text */
}
.gr-box {
    width: 100% !important;                   /* Full-width box */
    height: 100% !important;                  /* Full-height box */
}
img {
    width: 100% !important;                   /* Stretch image */
    height: 100% !important;
    object-fit: cover !important;             /* Fit image perfectly */
}
footer { 
    display: none !important;                 /* Hide footer */
}
button[title="Share"] {
    display: none !important;                 /* Hide Share button */
}

/* Simple white "Colorizer" heading */
h1 {
    font-size: 36px !important;               /* Standard font size */
    color: white !important;                  /* Plain white color */
    text-align: center !important;            /* Center alignment */
}
"""

# ✅ Gradio Interface (Image uploader + Slider output)
with gr.Blocks(css=custom_css) as demo:
    gr.Markdown("# Colorizer")

    # File uploader
    image_input = gr.File(label="Select Image", file_types=[".png", ".jpg", ".jpeg"])
    
    # Image Slider for comparison with full-size images
    output = ImageSlider(label="Old vs Colorized Image", type="filepath")

    # Process image when uploaded
    image_input.change(colorizer_fn, inputs=image_input, outputs=output)

# ✅ Launch the app
if __name__ == "__main__":
    demo.launch(share=True)