import io
import uuid
import os
from PIL import Image
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import FileResponse
from huggingface_hub import from_pretrained_fastai
import gradio as gr
import torch
import uvicorn

# ==========================================================
# 🔧 CONFIGURATION
# ==========================================================
MODEL_ID = "Hammad712/GAN-Colorization-Model"  # 👉 change this if you want another model

UPLOAD_DIR = "/tmp/uploads"
RESULT_DIR = "/tmp/results"
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(RESULT_DIR, exist_ok=True)

# ==========================================================
# 🚀 LOAD MODEL
# ==========================================================
print(f"Loading model: {MODEL_ID}")
learn = from_pretrained_fastai(MODEL_ID)
print("✅ Model loaded successfully!")

# ==========================================================
# 🧠 Colorization Function
# ==========================================================
def colorize_image(image: Image.Image):
    if image.mode != "RGB":
        image = image.convert("RGB")
    pred = learn.predict(image)[0]
    return pred

# ==========================================================
# 🌐 FASTAPI APP
# ==========================================================
app = FastAPI(title="Image Colorization API")

@app.post("/colorize")
async def colorize_endpoint(file: UploadFile = File(...)):
    """
    Upload a black & white image -> get colorized image
    """
    try:
        img_bytes = await file.read()
        image = Image.open(io.BytesIO(img_bytes)).convert("RGB")

        colorized = colorize_image(image)
        output_filename = f"{uuid.uuid4()}.png"
        output_path = os.path.join(RESULT_DIR, output_filename)
        colorized.save(output_path)

        return FileResponse(output_path, media_type="image/png")
    except Exception as e:
        return {"error": str(e)}

# ==========================================================
# 🎨 GRADIO INTERFACE
# ==========================================================
def gradio_interface(image):
    return colorize_image(image)

title = "🎨 FastAI / HuggingFace Image Colorizer"
description = "Upload a black & white photo to get a colorized version."

iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Image(type="pil", label="Colorized Output"),
    title=title,
    description=description,
)

gradio_app = gr.mount_gradio_app(app, iface, path="/")

# ==========================================================
# ▶️ RUN LOCALLY OR IN HUGGINGFACE SPACE
# ==========================================================
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)