import base64
import io
import os
from collections import Counter

import torch
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.responses import JSONResponse
from PIL import Image, ImageDraw, ImageFont
from transformers import RTDetrForObjectDetection, RTDetrImageProcessor


MODEL_ID = os.getenv("MODEL_ID", "./model")
CONFIDENCE_THRESHOLD = float(os.getenv("CONFIDENCE_THRESHOLD", "0.35"))
MAX_IMAGE_MB = int(os.getenv("MAX_IMAGE_MB", "15"))

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

processor = RTDetrImageProcessor.from_pretrained(MODEL_ID)
model = RTDetrForObjectDetection.from_pretrained(MODEL_ID)
model.to(DEVICE)
model.eval()

app = FastAPI(
    title="Ice Cream Counter API",
    version="1.0.0",
    description="RT-DETR ice cream product detector and counter. No YOLO.",
)


@app.get("/health")
def health():
    return {
        "status": "ok",
        "model": MODEL_ID,
        "device": str(DEVICE),
        "threshold": CONFIDENCE_THRESHOLD,
    }


def annotate(image, detections):
    image = image.copy()
    draw = ImageDraw.Draw(image)

    try:
        font = ImageFont.truetype("DejaVuSans.ttf", 18)
    except Exception:
        font = ImageFont.load_default()

    for d in detections:
        x1, y1, x2, y2 = d["box"]
        label = f'{d["class"]} {d["confidence"]:.2f}'
        draw.rectangle([x1, y1, x2, y2], outline="red", width=3)

        bbox = draw.textbbox((x1, y1), label, font=font)
        draw.rectangle(bbox, fill="red")
        draw.text((x1, y1), label, fill="white", font=font)

    return image


@app.post("/predict")
async def predict(
    file: UploadFile = File(...),
    return_image: bool = True,
):
    if not file.content_type or not file.content_type.startswith("image/"):
        raise HTTPException(400, "Upload a JPG, PNG, WEBP, or other image file.")

    raw = await file.read()

    if len(raw) > MAX_IMAGE_MB * 1024 * 1024:
        raise HTTPException(
            413,
            f"Image is too large. Maximum is {MAX_IMAGE_MB} MB.",
        )

    try:
        image = Image.open(io.BytesIO(raw)).convert("RGB")
    except Exception as exc:
        raise HTTPException(400, f"Could not read image: {exc}")

    inputs = processor(images=image, return_tensors="pt")
    inputs = {
        k: v.to(DEVICE) if torch.is_tensor(v) else v
        for k, v in inputs.items()
    }

    with torch.inference_mode():
        outputs = model(**inputs)

    target_sizes = torch.tensor(
        [[image.height, image.width]],
        device=DEVICE,
    )

    result = processor.post_process_object_detection(
        outputs,
        threshold=CONFIDENCE_THRESHOLD,
        target_sizes=target_sizes,
    )[0]

    detections = []
    counts = Counter()

    for score, label, box in zip(
        result["scores"],
        result["labels"],
        result["boxes"],
    ):
        score_value = float(score.item())
        label_id = int(label.item())
        class_name = model.config.id2label[label_id]

        coords = [round(float(x), 2) for x in box.tolist()]

        detections.append(
            {
                "class": class_name,
                "confidence": round(score_value, 4),
                "box": coords,
            }
        )
        counts[class_name] += 1

    response = {
        "total": len(detections),
        "counts": dict(sorted(counts.items())),
        "detections": detections,
    }

    if return_image:
        annotated = annotate(image, detections)
        buf = io.BytesIO()
        annotated.save(buf, format="JPEG", quality=90)
        response["annotated_image_base64"] = base64.b64encode(
            buf.getvalue()
        ).decode("ascii")

    return JSONResponse(response)