from dataclasses import dataclass
from typing import Optional

import numpy as np
from ultralytics import YOLO

from entity.model import AlgoModel
from services.model_service import ModelService
from common.global_logger import logger


@dataclass
class AlgoModelExec:
    algo_model_id: int
    algo_model_info: AlgoModel
    algo_model_exec: Optional[object] = None
    input_size: int = 640


class ModelManager:
    def __init__(self, db, model_warm_up=5):
        self.db = db
        self.model_service = ModelService(self.db)
        self.models = {}
        self.model_warm_up = model_warm_up

        self.load_models()

    def query_model_inuse(self):
        algo_model_list = list(self.model_service.get_models_in_use())
        for algo_model in algo_model_list:
            self.models[algo_model.id] = AlgoModelExec(
                algo_model_id=algo_model.id,
                algo_model_info=algo_model
            )

    def load_models(self):
        self.models = {}
        self.query_model_inuse()
        for algo_model_id, algo_model_exec in self.models.items():
            self.load_model(algo_model_exec)

    def load_model(self, algo_model_exec: AlgoModelExec):
        model_name = algo_model_exec.algo_model_info.name
        model_path = algo_model_exec.algo_model_info.path
        logger.info(f'loading model {model_name}: {model_path}')

        algo_model_exec.algo_model_exec = YOLO(model_path, task='detect')
        if self.model_warm_up > 0:
            logger.info(f'warming up model {model_name}')
            imgsz = algo_model_exec.input_size
            if not isinstance(imgsz, list):
                imgsz = [imgsz, imgsz]
            dummy_input = np.zeros((imgsz[0], imgsz[1], 3))
            for i in range(self.model_warm_up):
                algo_model_exec.algo_model_exec.predict(source=dummy_input, imgsz=imgsz, verbose=False)
            logger.info(f'warm up model {model_name} success!')
        logger.info(f'load model {model_name} success!')

    def reload_model(self, model_id):
        algo_model_exec = self.models.get(model_id)
        if algo_model_exec:
            algo_model_exec.algo_model_exec = None
            self.load_model(algo_model_exec)