import time
from abc import abstractmethod, ABC
from typing import Optional, TypeVar, Union, List
import aiohttp

from huibiao_framework.backend.metrics import PrometheusTimeCostMetricSetting
from huibiao_framework.config import ModelConfig
from loguru import logger
from pydantic import BaseModel


T = TypeVar("T", bound=BaseModel)
V = TypeVar("V", bound=BaseModel)


class AbstractLlmModelClient(ABC):
    def __init__(
        self,
        client_name: str,
        *,
        session: Optional[aiohttp.ClientSession] = None,
        url: str = None,
    ):
        self.__session = session
        self.__name = client_name
        assert self.__session is not None, "会话不能为空"
        self.__url = url if url else ModelConfig.REQUEST_URL

        self.headers = {"Content-Type": "application/json"}

        if ModelConfig.VLLM_BEAR_TOKEN:
            self.headers["Authorization"] = f"Bearer {ModelConfig.VLLM_BEAR_TOKEN}"

    @property
    def url(self):
        return self.__url

    @PrometheusTimeCostMetricSetting.client_metrics_handler("LLmClient")
    async def post(self, dto: T, session_id: str = "") -> dict:
        session_tag = self.session_tag(session_id)
        logger.debug(f"{session_tag},contentLien={self.cal_content_len(dto)}")
        start_time = time.perf_counter()
        try:
            async with self.session.post(
                self.url, json=dto.model_dump(exclude_none=True), headers=self.headers
            ) as resp:
                sp_time = time.perf_counter() - start_time
                logger.debug(f"{session_tag},resp-{resp.status},cost {sp_time:.2f}s")
                resp.raise_for_status()
                response_data = await resp.json()
                return response_data
        except Exception as e:
            logger.error(f"{session_tag} failed, {str(e)}")
            raise e

    @abstractmethod
    def cal_content_len(self, dto: T) -> Union[int, List[int]]:
        return [len(i.content) for i in dto.messages]

    @abstractmethod
    async def query(self, dto: T, session_id: str = "") -> V:
        pass

    @abstractmethod
    async def simple_query(self, content, session_id: str = "", **kwargs) -> str:
        """
        简单的一次性问答
        LLMException
        """
        pass

    @property
    def session(self):
        return self.__session

    @property
    def client_name(self) -> Optional[str]:
        return self.__name

    def session_tag(self, session_id: str) -> str:
        session_id_suffix = f"[{session_id}]" if session_id else ""
        return f"[{self.client_name}]{session_id_suffix}"
