from io import BytesIO
from typing import Any

from fastapi import APIRouter, File, HTTPException, UploadFile

router = APIRouter(tags=["Machine Learning"])


# Lazy-loaded model state, same shape as DL_CNN_NumberRecognition: the pretrained
# MobileNetV2 weights (~14 MB) download to the torch hub cache on first request.
MODEL_STATE: dict[str, Any] = {
    "model": None,
    "labels": None,
    "transforms": None,
    "error": None,
}

MAX_IMAGE_BYTES = 10 * 1024 * 1024  # 10 MB
TOP_K = 5


def _ensure_model_loaded() -> None:
    if MODEL_STATE["model"] is not None:
        return

    try:
        from torchvision.models import MobileNet_V2_Weights, mobilenet_v2

        weights = MobileNet_V2_Weights.IMAGENET1K_V1
        model = mobilenet_v2(weights=weights)
        model.eval()

        MODEL_STATE["model"] = model
        MODEL_STATE["labels"] = list(weights.meta["categories"])
        MODEL_STATE["transforms"] = weights.transforms()
        MODEL_STATE["error"] = None
    except Exception as e:
        MODEL_STATE["error"] = str(e)
        raise


@router.post("/models/image-classify", summary="Classify an image with MobileNetV2 (ImageNet)")
async def classify_image(file: UploadFile = File(...)):
    content_type = (file.content_type or "").lower()
    if not content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="Uploaded file must be an image.")

    raw = await file.read()
    if not raw:
        raise HTTPException(status_code=400, detail="Uploaded image is empty.")
    if len(raw) > MAX_IMAGE_BYTES:
        raise HTTPException(status_code=400, detail="Image exceeds the 10 MB size limit.")

    try:
        _ensure_model_loaded()
    except Exception:
        detail = "Model not loaded."
        if MODEL_STATE["error"]:
            detail = f"Model not loaded: {MODEL_STATE['error']}"
        return {"error": detail, "status": 500}

    import torch
    from PIL import Image, UnidentifiedImageError

    try:
        image = Image.open(BytesIO(raw)).convert("RGB")
    except UnidentifiedImageError:
        raise HTTPException(status_code=400, detail="Could not decode the uploaded image.")

    model = MODEL_STATE["model"]
    labels = MODEL_STATE["labels"]
    transforms = MODEL_STATE["transforms"]
    if model is None or labels is None or transforms is None:
        return {"error": "Model not loaded.", "status": 500}

    input_tensor = transforms(image).unsqueeze(0)
    with torch.no_grad():
        logits = model(input_tensor)
        probs = torch.softmax(logits, dim=1).squeeze(0)
        top_probs, top_indices = torch.topk(probs, TOP_K)

    predictions = [
        {"label": labels[int(idx)], "probability": float(prob)}
        for prob, idx in zip(top_probs.tolist(), top_indices.tolist())
    ]

    return {"predictions": predictions}