from abc import ABC, abstractmethod
from typing import Optional, Any
import numpy as np

from models.config import ModelConfig


class ModelManagerBase(ABC):
    @abstractmethod
    async def load_model(self, config: ModelConfig) -> bool:
        pass

    @abstractmethod
    async def unload_model(self, model_name: str) -> bool:
        pass

    @abstractmethod
    async def get_model(self, model_name: str) -> Optional[Any]:
        pass

    @abstractmethod
    async def run_inference(
        self, model_name: str, inputs: dict[str, np.ndarray]
    ) -> dict[str, np.ndarray]:
        pass

    @abstractmethod
    async def warmup_model(
        self, model_name: str, iterations: int = 3
    ) -> bool:
        pass

    @abstractmethod
    async def get_loaded_models(self) -> list[str]:
        pass

    @abstractmethod
    async def get_model_info(self, model_name: str) -> Optional[dict]:
        pass