import cv2
import numpy as np
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from PIL import Image
import io
import base64

app = FastAPI(title="Upscaler API", version="2.0")

# KRITISCH: Erlaube Anfragen von deinem UI-Space
# Ersetze 'YOUR-UI-SPACE-URL' mit der URL deines UI-Spaces (z.B. 'https://your-username-upscaler-ui.hf.space')
app.add_middleware(
    CORSMiddleware,
    allow_origins=["https://astridkraft-upscaling-ui.hf.space"], 
    allow_credentials=True,
    allow_methods=["POST"],
    allow_headers=["*"],
)

def simple_upscale(image: np.ndarray, scale: float):
    height, width = image.shape[:2]
    new_width = int(width * scale)
    new_height = int(height * scale)
    return cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_CUBIC)

def upscale_image_pil(input_image: Image.Image, scale_factor: str) -> Image.Image:
    try:
        scale_float = float(scale_factor.replace('x', '').replace(',', '.'))
        img_array = np.array(input_image)
        if input_image.mode != 'RGB':
            input_image = input_image.convert('RGB')
        img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
        output = simple_upscale(img_array, scale_float)
        output_rgb = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
        return Image.fromarray(output_rgb)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Upscaling error: {str(e)}")

@app.get("/")
def read_root():
    return {"message": "Upscaler API is running", "endpoint": "/upscale"}

@app.post("/upscale")
async def api_upscale(file: UploadFile = File(...), scale: str = Form("4x")):
    try:
        contents = await file.read()
        image = Image.open(io.BytesIO(contents))
        upscaled = upscale_image_pil(image, scale)
        
        output_buffer = io.BytesIO()
        upscaled.save(output_buffer, format="PNG")
        output_buffer.seek(0)
        
        base64_image = base64.b64encode(output_buffer.getvalue()).decode('utf-8')
        
        return JSONResponse({
            "success": True,
            "original_size": {"width": image.width, "height": image.height},
            "upscaled_size": {"width": upscaled.width, "height": upscaled.height},
            "scale_factor": scale,
            "image_base64": f"data:image/png;base64,{base64_image}"
        })
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Server error: {str(e)}")

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