# -*- coding:utf-8 -*-
"""
@Time : 2023/5/4 4:03 PM
@Author: hxl
@Des: websocket
"""
import asyncio
import json

import cv2
from fastapi import APIRouter, WebSocketDisconnect
from starlette.websockets import WebSocket

from arcsoft.arcface import FaceRecognitionEngine
from helper.websocket_manage import ConnectionManager
from models.user import User

router = APIRouter()

manager = ConnectionManager()


class FaceCameraEngine:
    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
    frame = None

    # 描绘人脸
    async def detect_faces(self, img):
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)
        if len(faces) == 0:
            return img
        for (x, y, w, h) in faces:
            cv2.rectangle(img, (x, y), (x + w, y + h), (255, 0, 0), 2)
            return img

    # 开始推流
    async def face_recoWork(self):
        cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
        while len(manager.active_connections):
            _, frame = cap.read()
            self.frame = frame
            frame = await self.detect_faces(frame)
            _, jpeg = cv2.imencode('.jpg', frame)
            await manager.broadcast_stream(jpeg.tobytes())
            await asyncio.sleep(0.1)
        cap.release()

    async def register(self, user_id: int):
        while len(manager.active_connections):
            with FaceRecognitionEngine() as frc:
                dat = frc.face_registration(self.frame)
            if dat:
                user = await User.get_or_none(id=user_id)
                if user:
                    user.face = dat
                    await user.save()
                    result = {"code": 200, "desc": "success"}
                    await manager.broadcast_json("face", result)
            await asyncio.sleep(2)

    async def login(self):
        dat_list = await User.filter(face__isnull=False, locked=False).values_list('id', 'face')
        while len(manager.active_connections):
            with FaceRecognitionEngine() as frc:
                user_id = frc.face_login(self.frame, dat_list)
            if user_id:
                user = await User.get_or_none(id=user_id)
                result = {"code": 200, "data": user.profile}
                await manager.broadcast_json("face", json.dumps(result))
            await asyncio.sleep(2)


@router.websocket("/login")
async def face_login(websocket: WebSocket):
    await manager.connect(websocket)
    if len(manager.active_connections) == 1:
        fc = FaceCameraEngine()
        task1 = asyncio.create_task(fc.face_recoWork())
        task2 = asyncio.create_task(fc.login())
        await task1, task2
    try:
        while True:
            await websocket.receive_bytes()
    except WebSocketDisconnect:
        manager.disconnect(websocket)


@router.websocket("/register/{id}")
async def face_login(websocket: WebSocket):
    await manager.connect(websocket)
    if len(manager.active_connections) == 1:
        user_id = websocket.path_params.get('id')
        fc = FaceCameraEngine()
        task1 = asyncio.create_task(fc.face_recoWork())
        task2 = asyncio.create_task(fc.register(user_id))
        await task1, task2

    try:
        while True:
            await websocket.receive_bytes()
    except WebSocketDisconnect:
        manager.disconnect(websocket)