from fastapi import FastAPI, UploadFile, File
from fastapi.responses import FileResponse, HTMLResponse
from fastapi import HTTPException
from fastapi.staticfiles import StaticFiles
import subprocess
import uuid
import os
import shutil

app = FastAPI()

app.mount("/static", StaticFiles(directory="static"), name="static")

OUTPUT_DIR = "outputs"
os.makedirs(OUTPUT_DIR, exist_ok=True)

REALESRGAN_PATH = "realesrgan-ncnn-vulkan"  # should be in PATH


@app.get("/", response_class=HTMLResponse)
def home():
    with open("static/index.html") as f:
        return f.read()


@app.post("/enhance")
async def enhance_image(file: UploadFile = File(...)):
    input_name = f"{uuid.uuid4()}_{file.filename}"
    input_path = os.path.join(OUTPUT_DIR, input_name)

    with open(input_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    output_path = os.path.join(OUTPUT_DIR, f"enhanced_{input_name}")

    cmd = [
        "realesrgan-ncnn-vulkan",
        "-i", input_path,
        "-o", output_path,
        "-s", "2",
        "-m", "/opt/realesrgan/models"
    ]

    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        raise HTTPException(status_code=500, detail=result.stderr)

    enhanced = os.path.join(OUTPUT_DIR, f"enhanced_{input_name}")

    return FileResponse(enhanced, media_type="image/png")