from typing import Optional
from pathlib import Path
import tempfile

from fastapi import FastAPI, UploadFile, File, Query, HTTPException
from fastapi.middleware.cors import CORSMiddleware

from app.schemas import PredictResponse
from app.utils import decode_image_from_bytes
from app.pipeline import BirdCascadePipeline
from app.image_pipeline import save_image_outputs
from app.video_pipeline import process_video_file


app = FastAPI(
    title="Bird Cascade Service",
    description="YOLOv5 detection + target crop + ConvNeXt fine-grained bird classification",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

pipeline: Optional[BirdCascadePipeline] = None
startup_error: Optional[str] = None


@app.on_event("startup")
def load_models():
    global pipeline, startup_error
    try:
        pipeline = BirdCascadePipeline()
        startup_error = None
    except Exception as e:
        pipeline = None
        startup_error = str(e)


@app.get("/health")
def health():
    return {
        "status": "ok" if pipeline is not None else "error",
        "det_model_loaded": pipeline is not None,
        "cls_model_loaded": pipeline is not None,
        "error": startup_error,
    }


@app.post("/predict", response_model=PredictResponse)
async def predict(
    file: UploadFile = File(...),
    return_image: bool = Query(False, description="是否在返回 JSON 中包含 base64 图片"),
    det_conf: Optional[float] = Query(None, ge=0.0, le=1.0, description="YOLO 检测置信度阈值"),
    iou_thres: Optional[float] = Query(None, ge=0.0, le=1.0, description="NMS IoU 阈值"),
):
    """
    图片识别接口。

    无论 return_image 是否为 true，都会自动保存：
    1. 打框图片；
    2. JSON；
    3. XML。

    保存目录：
    runs/images/
    """
    if pipeline is None:
        raise HTTPException(status_code=500, detail=f"模型未加载成功: {startup_error}")

    try:
        data = await file.read()
        img_bgr = decode_image_from_bytes(data)

        output = pipeline.predict(
            img_bgr,
            return_image=return_image,
            det_conf=det_conf,
            iou_thres=iou_thres,
        )

        saved_outputs = save_image_outputs(
            img_bgr=img_bgr,
            original_filename=file.filename or "input.jpg",
            predict_output=output,
        )

        output["saved_outputs"] = saved_outputs
        return output

    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.post("/predict_video")
async def predict_video(
    file: UploadFile = File(...),
    sample_interval: Optional[int] = Query(None, ge=1, description="视频抽帧间隔，例如 2 表示每 2 帧处理 1 帧"),
    save_annotated_video: bool = Query(True, description="是否保存画框后视频"),
):
    """
    视频识别接口。

    自动保存：
    1. 打框视频；
    2. JSON；
    3. XML；
    4. 关键帧截图；
    5. 原始视频副本。

    保存目录：
    runs/video/
    """
    if pipeline is None:
        raise HTTPException(status_code=500, detail=f"模型未加载成功: {startup_error}")

    try:
        suffix = Path(file.filename or "input.mp4").suffix or ".mp4"

        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            tmp.write(await file.read())
            tmp_path = tmp.name

        summary = process_video_file(
            video_path=tmp_path,
            pipeline=pipeline,
            original_filename=file.filename or "input.mp4",
            sample_interval=sample_interval,
            save_annotated_video=save_annotated_video,
        )

        return summary

    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))