import asyncio
import json
from pathlib import Path

import cv2
import numpy as np
import torch
from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI(title="Reading Trace L2CS")
pipeline = None


def get_pipeline():
    global pipeline
    if pipeline is None:
        from l2cs import Pipeline
        pipeline = Pipeline(
            weights=Path("models/L2CSNet_gaze360.pkl"),
            arch="ResNet50",
            device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
        )
    return pipeline


@app.get("/health")
def health():
    get_pipeline()
    return {"status": "ok", "device": "cuda" if torch.cuda.is_available() else "cpu"}


def infer(data: bytes):
    frame = cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_COLOR)
    if frame is None:
        return {"error": "invalid frame"}
    height, width = frame.shape[:2]
    try:
        result = get_pipeline().step(frame)
    except ValueError:
        return {"faces": 0}
    yaw = np.asarray(result.yaw).reshape(-1)
    pitch = np.asarray(result.pitch).reshape(-1)
    if not yaw.size:
        return {"faces": 0}
    boxes = np.asarray(getattr(result, "bboxes", []))
    index, fx, fy = 0, .5, .5
    if boxes.ndim == 2 and boxes.shape[0] == yaw.size and boxes.shape[1] >= 4:
        areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
        index = int(np.argmax(areas))
        x0, y0, x1, y1 = boxes[index, :4]
        fx, fy = float((x0 + x1) / 2 / width), float((y0 + y1) / 2 / height)
    return {"yaw": float(yaw[index]), "pitch": float(pitch[index]), "fx": fx, "fy": fy, "faces": int(yaw.size)}


@app.websocket("/ws")
async def websocket_gaze(socket: WebSocket):
    await socket.accept()
    try:
        while True:
            data = await socket.receive_bytes()
            if len(data) > 2_000_000:
                await socket.close(code=1009)
                return
            result = await asyncio.to_thread(infer, data)
            await socket.send_text(json.dumps(result))
    except WebSocketDisconnect:
        pass