import cv2
import numpy as np
import base64
import logging
from fastapi import APIRouter, UploadFile, File, Form, Query
from models.schemas import ImageDetectionResult

router = APIRouter()
logger = logging.getLogger(__name__)

@router.post("/detect/image", response_model=ImageDetectionResult)
async def detect_image(
    file: UploadFile = File(...),
    model_id: str = Query("fire_detection"),
    confidence: float = Query(0.5),
    iou: float = Query(0.45)
):
    from main import model_service
    from services.detection_service import DetectionService
    
    detection_service = DetectionService(model_service)
    
    try:
        contents = await file.read()
        nparr = np.frombuffer(contents, np.uint8)
        frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        if frame is None:
            return ImageDetectionResult(
                success=False,
                message="无法读取图片",
                data={}
            )
        
        result = await detection_service.detect_image(frame, model_id, confidence, iou)
        
        if result['success']:
            annotated_frame = detection_service.draw_detections(frame, result['detections'])

            # 将标注后的图片转换为 base64
            _, buffer = cv2.imencode('.jpg', annotated_frame)
            img_base64 = base64.b64encode(buffer).decode('utf-8')

            return ImageDetectionResult(
                success=True,
                message="检测完成",
                data={
                    "detections": result['detections'],
                    "image_base64": img_base64,
                    "stats": result['stats']
                }
            )
        else:
            return ImageDetectionResult(
                success=False,
                message=result['message'],
                data={}
            )
            
    except Exception as e:
        logger.error(f"图片检测失败: {e}")
        return ImageDetectionResult(
            success=False,
            message=f"检测失败: {str(e)}",
            data={}
        )