from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import JSONResponse, Response, FileResponse
from fastapi.middleware.cors import CORSMiddleware
from ultralytics import YOLO
import torch
import cv2
import numpy as np

app = FastAPI()

# Разрешаем вызовы из фронта того же Space
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Загружаем модель
model = YOLO("best.pt")
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

@app.get("/")
def root():
    return FileResponse("index.html")

def read_image_to_bgr(file_bytes: bytes) -> np.ndarray:
    # Декод JPEG/PNG в BGR
    img_array = np.frombuffer(file_bytes, dtype=np.uint8)
    img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)  # BGR
    return img

def annotate_bgr(results) -> np.ndarray:
    # results[0].plot() возвращает BGR с нарисованными боксами
    return results[0].plot()

def results_to_json(results):
    # Конвертация результатов в чистые боксы/классы/скор
    r = results[0]
    boxes = r.boxes
    out = []
    if boxes is not None and len(boxes) > 0:
        xyxy = boxes.xyxy.cpu().numpy()  # (N,4)
        conf = boxes.conf.cpu().numpy()  # (N,)
        cls = boxes.cls.cpu().numpy().astype(int)  # (N,)
        names = r.names
        for i in range(len(xyxy)):
            x1, y1, x2, y2 = xyxy[i].tolist()
            out.append({
                "bbox": [x1, y1, x2, y2],
                "conf": float(conf[i]),
                "class_id": int(cls[i]),
                "class_name": names[int(cls[i])] if names else str(cls[i])
            })
    return {"detections": out}

@app.post("/predict")
async def predict(
    file: UploadFile = File(...),
    conf: float = Form(0.25),
    iou: float = Form(0.45),
    return_image: int = Form(1)  # 1 = вернуть аннотированное изображение, 0 = вернуть JSON боксов
):
    data = await file.read()
    bgr = read_image_to_bgr(data)
    if bgr is None:
        return JSONResponse({"error": "Invalid image"}, status_code=400)

    # Инференс (без трекинга — кадры независимы; для трекинга можно persist и tracker)
    results = model.predict(
        source=bgr,
        conf=conf,
        iou=iou,
        imgsz=640,
        verbose=False
    )

    if return_image == 1:
        annotated = annotate_bgr(results)  # BGR
        # Кодируем в JPEG для отправки
        ok, buf = cv2.imencode(".jpg", annotated)
        if not ok:
            return JSONResponse({"error": "Encode failed"}, status_code=500)
        return Response(content=buf.tobytes(), media_type="image/jpeg")
    else:
        return JSONResponse(results_to_json(results))