import gradio as gr
import torch
import os
from PIL import Image
import numpy as np
from huggingface_hub import hf_hub_download, list_repo_files

# -----------------------------
# CONFIG
# -----------------------------
REPO_ID = "easygoing0114/AI_upscalers"
MODEL_DIR = "/tmp/upscalers"
os.makedirs(MODEL_DIR, exist_ok=True)

DEVICE = "cpu"

# -----------------------------
# LOAD MODEL LIST
# -----------------------------
def get_models():
    files = list_repo_files(REPO_ID)
    models = [f for f in files if f.endswith(".pth")]
    return models

MODEL_LIST = get_models()

# -----------------------------
# DOWNLOAD MODEL
# -----------------------------
def download_model(model_name):
    path = hf_hub_download(
        repo_id=REPO_ID,
        filename=model_name,
        local_dir=MODEL_DIR
    )
    return path

# -----------------------------
# LOAD UPSCALER (GENERIC ESRGAN STYLE)
# -----------------------------
def load_model(model_path):
    # Lazy import to avoid heavy startup
    from basicsr.archs.rrdbnet_arch import RRDBNet
    from realesrgan import RealESRGANer

    model = RRDBNet(
        num_in_ch=3,
        num_out_ch=3,
        num_feat=64,
        num_block=23,
        num_grow_ch=32,
        scale=4
    )

    upsampler = RealESRGANer(
        scale=4,
        model_path=model_path,
        model=model,
        tile=0,
        tile_pad=10,
        pre_pad=0,
        half=(DEVICE == "cuda"),
        device=DEVICE
    )

    return upsampler

# -----------------------------
# UPSCALE FUNCTION
# -----------------------------
def upscale_image(image, model_name):
    if image is None:
        return None

    # Download model
    model_path = download_model(model_name)

    # Load model
    upsampler = load_model(model_path)

    # Convert image
    img = np.array(image)

    # Upscale
    output, _ = upsampler.enhance(img, outscale=4)

    return Image.fromarray(output)

# -----------------------------
# UI
# -----------------------------
with gr.Blocks(title="AI Image Upscaler") as app:
    gr.Markdown("# 🔍 AI Image Upscaler (Multi-Model)")
    gr.Markdown("Select any model from the repository and upscale your image.")

    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Image")

    model_dropdown = gr.Dropdown(
        choices=MODEL_LIST,
        value=MODEL_LIST[0] if MODEL_LIST else None,
        label="Select Upscaler Model"
    )

    upscale_btn = gr.Button("✨ Upscale Image")

    output_image = gr.Image(label="Upscaled Image")

    upscale_btn.click(
        fn=upscale_image,
        inputs=[image_input, model_dropdown],
        outputs=output_image
    )

# -----------------------------
# LAUNCH
# -----------------------------
if __name__ == "__main__":
    app.launch()