from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import StreamingResponse
from media_enhancer.image_enhancer import EnhancementMethod, Enhancer
from pydantic import BaseModel
from PIL import Image
from io import BytesIO
import base64
import numpy as np

class EnhancementRequest(BaseModel):
    method: EnhancementMethod = EnhancementMethod.gfpgan
    background_enhancement: bool = True
    upscale: int = 2

class _EnhanceBase(BaseModel):
    encoded_base_img: List[str]



app = FastAPI()

@app.get("/")
def greet_json():
    return {"Initializing GlamApp Enhancer"}

@app.post("/enhance")
async def enhance_image(
    file: UploadFile = File(...),
    request: EnhancementRequest = EnhancementRequest()
):
    try:
        if not file.content_type.startswith('image/'):
            raise HTTPException(status_code=400, detail="Invalid file type")
            
        contents = await file.read()
        base64_encoded_image = base64.b64encode(contents).decode('utf-8')
        data = _EnhanceBase(encoded_base_img=[base64_encoded_image])

        enhancer = Enhancer(request.method, request.background_enhancement, request.upscale)

        enhanced_img, original_resolution, enhanced_resolution = await enhancer.enhance(data)

        enhanced_image = Image.fromarray(enhanced_img)
        img_byte_arr = BytesIO()
        enhanced_image.save(img_byte_arr, format='PNG')
        img_byte_arr.seek(0)
        
        return StreamingResponse(img_byte_arr, media_type="image/png")

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))