# app.py

import os
import io
import base64
import numpy as np
from PIL import Image
import torch
import cv2 # Ab yeh safely import hoga

# GFPGAN/ESRGAN Libraries
from gfpgan import GFPGANer 
# GFPGANer class mein hi ESRGAN functionality shamil hoti hai

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

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

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

try:
    # Real-ESRGAN model use kar rahe hain jo GFPGANer mein built-in hai.
    # Upar humne 'realesrgan' package hata diya hai aur sirf isko use kar rahe hain.
    
    # Model name: 'realesrgan-x4plus' ko GFPGANer use karta hai
    UPSCALER = GFPGANer(
        model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.3.0/RealESRGAN_x4plus.pth',
        upscale=4,
        arch='realesrgan', 
        device=DEVICE,
        bg_upsampler=None # Background upsampler off taaki Free Tier par tez chale
    )
    print("Real-ESRGAN (via GFPGANer) 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.")
        
    # GFPGANer.enhance method: input (BGR format), output (BGR format)
    # PIL image RGB mein hoti hai, OpenCV BGR mein expect karta hai.
    
    # 1. RGB (PIL) se BGR (OpenCV) mein convert karein
    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)

    # 2. Upscaling
    _, _, output_bgr = UPSCALER.enhance(img_bgr, has_aligned=False, only_center_face=False, paste_back=True)
    
    # 3. Output BGR se wapas RGB mein convert karein
    output_rgb = cv2.cvtColor(output_bgr, cv2.COLOR_BGR2RGB)
    
    return output_rgb

# --- 2. FastAPI Setup and Endpoints (Wohi Rahenge) ---

app = FastAPI(title="Real-ESRGAN Custom Upscaler API")

# Image file upload ke zariye upscaling
@app.post("/api/upscale/file")
async def upscale_image_api(image: UploadFile = File(...)):
    # ... (API logic wahi rahega) ...
    try:
        image_bytes = await image.read()
        input_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        img_np = np.array(input_image)
        
        output_np = run_upscaler(img_np)
        
        output_image = Image.fromarray(output_np)
        img_io = io.BytesIO()
        output_image.save(img_io, format='PNG')
        img_io.seek(0)
        
        return StreamingResponse(img_io, media_type="image/png")

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


# Gradio Interface
def upscale_for_gradio(input_image: Image.Image):
    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)}"

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="⭐ Free Tier Optimized: Real-ESRGAN Upscaler (UI & Custom API)",
    description="Optimized code structure istemaal karte hue high-quality upscaling.",
    allow_flagging="never"
)

# Gradio ko FastAPI app mein mount karna
app = gr.mount_gradio_app(app, gr_interface, path="/") 

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)