import uuid

import cv2
import numpy as np
from fastapi import APIRouter, File, Form, HTTPException, UploadFile

from engine.loader import get_engine, get_inference_lock
from engine.annotator import annotate, to_base64
from schema import DefectBoxSchema, DetectResponse

router = APIRouter()


@router.post("/detect", response_model=DetectResponse)
async def detect(
    image: UploadFile = File(...),
    conf: float = Form(0.45),
    return_annotated: bool = Form(False),
):
    raw = await image.read()
    arr = np.frombuffer(raw, dtype=np.uint8)
    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if img is None:
        raise HTTPException(status_code=400, detail="无法解码图像，请检查文件格式")

    engine = get_engine()
    async with get_inference_lock():
        result = engine.detect(img, conf=conf)

    annotated_b64 = None
    if return_annotated:
        annotated_b64 = to_base64(annotate(img, result))

    return DetectResponse(
        task_id=str(uuid.uuid4()),
        verdict=result.verdict,
        defect_count=result.defect_count,
        max_severity=result.max_severity,
        avg_confidence=result.avg_confidence,
        duration_ms=result.duration_ms,
        image_width=result.image_width,
        image_height=result.image_height,
        model_version=result.model_version,
        defects=[
            DefectBoxSchema(
                class_id=d.class_id,
                class_name=d.class_name,
                class_name_zh=d.class_name_zh,
                confidence=d.confidence,
                severity=d.severity,
                box_xyxy=d.box_xyxy,
            )
            for d in result.defects
        ],
        annotated_image_base64=annotated_b64,
    )