import io
import cv2
import numpy as np
import uvicorn
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import StreamingResponse
from gfpgan import GFPGANer
from realesrgan import RealESRGANer
from basicsr.archs.rrdbnet_arch import RRDBNet

app = FastAPI(title="AI Image Enhancer Server")

# 1. إعداد نموذج زيادة الدقة (Real-ESRGAN)
model_realesrgan = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
upsampler = RealESRGANer(
    scale=2,
    model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth',
    model=model_realesrgan,
    tile=0,
    tile_pad=10,
    pre_pad=0,
    half=False # اجعلها True إذا كنت تستخدم GPU
)

# 2. إعداد نموذج تحسين الوجه (GFPGAN)
face_enhancer = GFPGANer(
    model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
    upscale=2,
    arch='clean',
    channel_multiplier=2,
    bg_upsampler=upsampler # دمج زيادة الدقة مع تحسين الوجه
)

@app.post("/enhance")
async def enhance_image(file: UploadFile = File(...)):
    # قراءة الصورة من الطلب
    data = await file.read()
    nparr = np.frombuffer(data, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

    # تنفيذ السحر: تحسين الوجه + زيادة الدقة
    # face_enhancer سيقوم بتوضيح الملامح و Real-ESRGAN سيهتم بالخلفية والجودة العامة
    _, _, enhanced_img = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)

    # تحويل الصورة الناتجة إلى صيغة يمكن إرسالها عبر المتصفح
    _, buffer = cv2.imencode(".png", enhanced_img)
    io_buf = io.BytesIO(buffer)
    
    return StreamingResponse(io_buf, media_type="image/png")

@app.get("/")
def home():
    return {"message": "سيرفر تحسين الصور يعمل بنجاح!"}

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