import os
import tempfile
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import Response
from PIL import Image
import io
from super_image import MsrnModel, ImageLoader

app = FastAPI()

models = {}


def get_model(scale: int):
    if scale not in models:
        model = MsrnModel.from_pretrained("eugenesiow/msrn", scale=scale)
        model = model.eval()
        models[scale] = model
    return models[scale]


def _upscale_image(img: Image.Image, scale: int) -> Image.Image:
    model = get_model(scale)
    inputs = ImageLoader.load_image(img)
    pred = model(inputs)
    fd, path = tempfile.mkstemp(suffix=".png")
    try:
        ImageLoader.save_image(pred, path)
        os.close(fd)
        result = Image.open(path).convert("RGB")
        result.load()
    finally:
        os.unlink(path)
    return result


@app.post("/upscale")
async def upscale(image: UploadFile = File(...), scale: int = Form(2)):
    try:
        img = Image.open(io.BytesIO(await image.read())).convert("RGB")
    except Exception as e:
        return Response(
            f"Image load failed: {type(e).__name__}: {e}",
            status_code=500,
            media_type="text/plain",
        )

    try:
        if scale == 8:
            # Chain: 4x → 2x
            img = _upscale_image(img, 4)
            img = _upscale_image(img, 2)
        else:
            img = _upscale_image(img, scale)
    except Exception as e:
        return Response(
            f"Upscale failed: {type(e).__name__}: {e}",
            status_code=500,
            media_type="text/plain",
        )

    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return Response(buf.getvalue(), media_type="image/png")