"""Inference API for local use and Docker Spaces. No training imports."""

import io
import os
from contextlib import asynccontextmanager

from fastapi import FastAPI, File, HTTPException, UploadFile
from PIL import Image, UnidentifiedImageError

MAX_UPLOAD = 8 * 1024 * 1024


def create_app(recognizer=None):
    @asynccontextmanager
    async def lifespan(app):
        if recognizer is None:
            threads = int(os.getenv("MATH2LATEX_THREADS", "2"))
            if os.getenv("MATH2LATEX_BACKEND", "formulanet") == "formulanet":
                from image_to_latex.formulanet import FormulaNetRecognizer

                app.state.recognizer = FormulaNetRecognizer(
                    os.getenv("MATH2LATEX_MODEL_DIR", "artifacts/formulanet"), threads
                )
            else:
                from image_to_latex.inference import Recognizer

                app.state.recognizer = Recognizer(
                    os.getenv("MATH2LATEX_MODEL_DIR", "artifacts/latex-ocr"),
                    threads,
                    max_tokens=int(os.getenv("MATH2LATEX_MAX_TOKENS", "192")),
                    quantized=os.getenv("MATH2LATEX_INT8", "0") == "1",
                )
        else:
            app.state.recognizer = recognizer
        yield

    app = FastAPI(
        title="Math2LaTeX",
        description="Convert a cropped formula into LaTeX.",
        lifespan=lifespan,
    )

    @app.get("/")
    def health():
        return {"status": "ready", "backend": app.state.recognizer.backend}

    @app.post("/predict/")
    def predict(file: UploadFile = File(...)):
        payload = file.file.read(MAX_UPLOAD + 1)
        if len(payload) > MAX_UPLOAD:
            raise HTTPException(413, "Maximum upload size is 8 MiB")
        try:
            with Image.open(io.BytesIO(payload)) as image:
                result = app.state.recognizer.predict(image)
            return {"data": {"pred": result["latex"], **result}}
        except (
            UnidentifiedImageError,
            ValueError,
            OSError,
            Image.DecompressionBombError,
        ) as exc:
            raise HTTPException(400, str(exc)) from exc

    return app


app = create_app()