# app.py

import os
import io
import base64
import numpy as np
from PIL import Image
import torch
from realesrgan import RealESRGANer

# FastAPI Libraries
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse, StreamingResponse
import uvicorn
import gradio as gr

# --- 1. Model Loading (Free Tier Optimized) ---

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Model will run on: {DEVICE}")

try:
    # Real-ESRGAN ka lightweight, optimized model use kar rahe hain
    model_path = RealESRGANer.model_path_from_name('RealESRGAN_x4plus')
    
    # Model ko load karna (yeh memory mein rahega)
    UPSCALER = RealESRGANer(
        scale=4, 
        model_path=model_path,
        dni_weight=None,
        model_name='RealESRGAN_x4plus',
        device=DEVICE
    )
    print("Real-ESRGAN model loaded successfully.")

except Exception as e:
    print(f"ERROR: Model load nahi ho paya. Error: {e}")
    UPSCALER = None

def run_upscaler(img_np: np.ndarray):
    """Core upscaling logic."""
    if UPSCALER is None:
        raise Exception("Model is not initialized.")
        
    # Upscaling (yahan time lagta hai)
    output_np, _ = UPSCALER.enhance(img_np, outscale=4)
    
    return output_np

# --- 2. FastAPI Setup ---

# FastAPI application ko initialize karein
app = FastAPI(title="Real-ESRGAN Custom Upscaler API")

# --- 3. Custom API Endpoint ---

# Image file upload ke zariye upscaling
@app.post("/api/upscale/file")
async def upscale_image_api(image: UploadFile = File(...)):
    """
    Image file ko upload karein aur 4x upscaled image wapas hasil karein.
    """
    try:
        # File ko PIL Image mein padhna
        image_bytes = await image.read()
        input_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        
        # PIL image ko numpy array mein convert karna
        img_np = np.array(input_image)
        
        # Upscaling
        output_np = run_upscaler(img_np)
        
        # NumPy array ko wapas PIL Image mein convert karna
        output_image = Image.fromarray(output_np)
        
        # Image ko BytesIO mein save karna
        img_io = io.BytesIO()
        output_image.save(img_io, format='PNG')
        img_io.seek(0)
        
        # StreamingResponse se image ko wapas bhejna
        return StreamingResponse(img_io, media_type="image/png")

    except Exception as e:
        return JSONResponse(status_code=500, content={"message": f"Processing error: {str(e)}"})

# --- 4. Gradio UI Integration ---

def upscale_for_gradio(input_image: Image.Image):
    """Gradio UI ke liye wrapper function."""
    try:
        img_np = np.array(input_image.convert("RGB"))
        output_np = run_upscaler(img_np)
        return Image.fromarray(output_np)
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio Interface define karna
gr_interface = gr.Interface(
    fn=upscale_for_gradio,
    inputs=gr.Image(type="pil", label="Low-Resolution Image Upload Karein"),
    outputs=gr.Image(type="pil", label="4x Upscaled (High-Quality) Image"),
    title="⭐ Real-ESRGAN: AI Image Upscaler (UI & Custom API)",
    description="Apni images ko 4x size mein badhayein. Yeh app Custom REST API aur Gradio UI dono offer karta hai.",
    allow_flagging="never"
)

# Gradio ko FastAPI app mein mount karna
# '/gradio' path par UI available hoga
app = gr.mount_gradio_app(app, gr_interface, path="/") 

# --- 5. Uvicorn Server Setup ---

# Yeh tabhi run hoga jab aap file ko directly chalayenge (lekin Docker mein yeh entry point hoga)
if __name__ == "__main__":
    # Hugging Face Spaces Docker mein port 7860 par chalne ki umeed rakhta hai.
    # Hamara server isi port par run hoga.
    uvicorn.run(app, host="0.0.0.0", port=7860)