import io
import cv2
import numpy as np
import requests
import os
import gc
import torch
from fastapi import FastAPI, UploadFile, File
from fastapi.responses import Response
from realesrgan import RealESRGANer
from basicsr.archs.rrdbnet_arch import RRDBNet

app = FastAPI()

# 🧠 Setup Real-ESRGAN Model
model_url = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'
model_path = 'RealESRGAN_x4plus.pth'

if not os.path.exists(model_path):
    print("Downloading AI model... please wait.")
    response = requests.get(model_url)
    with open(model_path, 'wb') as f:
        f.write(response.content)

# Initialize the AI Engine once
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=400, # Crucial for free CPU tier
    tile_pad=10,
    pre_pad=0,
    half=False
)

@app.get("/")
def home():
    return {"status": "Silent Neural HD Engine Online"}

@app.post("/upscale")
async def upscale(file: UploadFile = File(...)):
    try:
        data = await file.read()
        nparr = np.frombuffer(data, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        # Process image
        output, _ = upsampler.enhance(img, outscale=4)
        
        # Convert to JPG
        _, encoded_img = cv2.imencode('.jpg', output, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
        
        # Cleanup memory immediately
        del img
        del output
        gc.collect()
        
        return Response(content=encoded_img.tobytes(), media_type="image/jpeg")
    
    except Exception as e:
        print(f"Error: {e}")
        return {"error": "Processing failed. Image might be too large."}

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