import logging
import sys
from typing import Any, Union, cast

from .._utils import ForTask
from ..types import APIFuture, SampleResponse, SamplingParams
from ._async_holder import AsyncHolder

logger = logging.getLogger(__name__)


class SamplingClient:
    def __init__(
        self, task_id: str, base_model: str, socket: ForTask, holder: AsyncHolder
    ):
        self.task_id = task_id
        # 使用的基础模型
        self.base_model = base_model

        self.socket = socket
        self.holder = holder

    def _sample_prepare(
        self,
        prompt: Union[list[int], list[list[int]]],
        num_samples: int,
        sampling_params: SamplingParams,
        include_prompt_logprobs: bool,
        topk_prompt_logprobs: int,
    ) -> dict[str, Any] | None:
        p = self._compute_logprobs_prepare(
            prompt=prompt,
        )
        if p is None:
            return None

        batched_prompt = p["payload"]["input_ids"]
        input_tokens = p["input_tokens"]

        # 如果设置了 ignore_eos 则必须设置 max_tokens 参数
        if sampling_params.ignore_eos and not sampling_params.max_tokens:
            logger.error("🙃 当 ignore_eos=True 时，必须设置 max_tokens 参数")
            sys.exit(1)

        return {
            "type": "sample",
            "payload": {
                "input_ids": batched_prompt,
                "params": sampling_params.model_dump(),
                "num_samples": num_samples,
                "include_prompt_logprobs": include_prompt_logprobs,
                "topk_prompt_logprobs": topk_prompt_logprobs,
            },
            "input_tokens": input_tokens,
        }

    def sample(
        self,
        prompt: Union[list[int], list[list[int]]],
        num_samples: int,
        sampling_params: SamplingParams,
        include_prompt_logprobs: bool = False,
        topk_prompt_logprobs: int = 0,
    ) -> APIFuture[SampleResponse]:
        """
        根据模型生成文本补全

        Parameters
        ----------
        prompt: 输入的 token id 列表
        num_samples: 生成样本数量
        sampling_params: 采样参数
        include_prompt_logprobs: 是否包含 prompt 部分的对数概率
        topk_prompt_logprobs: 返回的 prompt 部分的 top-k 对数概率数量

        Returns
        -------

        """
        return self.holder.run_coroutine_threadsafe(
            self.sample_async(
                prompt=prompt,
                num_samples=num_samples,
                sampling_params=sampling_params,
                include_prompt_logprobs=include_prompt_logprobs,
                topk_prompt_logprobs=topk_prompt_logprobs,
            )
        ).result()

    async def sample_async(
        self,
        prompt: Union[list[int], list[list[int]]],
        num_samples: int,
        sampling_params: SamplingParams,
        include_prompt_logprobs: bool = False,
        topk_prompt_logprobs: int = 0,
    ) -> APIFuture[SampleResponse]:
        """
        根据模型生成文本补全

        Parameters
        ----------
        prompt: 输入的 token id 列表
        num_samples: 生成样本数量
        sampling_params: 采样参数
        include_prompt_logprobs: 是否包含 prompt 部分的对数概率
        topk_prompt_logprobs: 返回的 prompt 部分的 top-k 对数概率数量

        Returns
        -------

        """
        p = self._compute_logprobs_prepare(
            prompt=prompt,
        )
        if p is None:
            logger.error("🙃 采样输入无效")
            return APIFuture.completed(
                SampleResponse(
                    sequences=[],
                    prompt_logprobs=[],
                    topk_prompt_logprobs=[],
                    output_tokens=0,
                    elapsed_time=0.0,
                )
            )

        batched_prompt = p["payload"]["input_ids"]
        input_tokens = p["input_tokens"]

        # 如果设置了 ignore_eos 则必须设置 max_tokens 参数
        if sampling_params.ignore_eos and not sampling_params.max_tokens:
            logger.error("🙃 当 ignore_eos=True 时，必须设置 max_tokens 参数")
            sys.exit(1)

        response = await self.socket.send_event_async(
            event_type="sample",
            payload={
                "input_ids": batched_prompt,
                "params": sampling_params.model_dump(),
                "num_samples": num_samples,
                "include_prompt_logprobs": include_prompt_logprobs,
                "topk_prompt_logprobs": topk_prompt_logprobs,
            },
            input_tokens=input_tokens,
        )
        return response.map_model(SampleResponse)

    def _compute_logprobs_prepare(
        self, prompt: list[list[int]] | list[int]
    ) -> dict[str, Any] | None:

        batched_prompt: list[list[int]]

        if not prompt:
            batched_prompt = [[]]
        elif isinstance(prompt[0], int):
            batched_prompt = [cast(list[int], prompt)]
        else:
            batched_prompt = cast(list[list[int]], prompt)

        # 空输入
        if len(batched_prompt[0]) == 0:
            logger.error("🙃 输入为空，请检查你的输入")
            return None

        input_tokens = sum(len(sublist) for sublist in batched_prompt)

        return {
            "type": "compute-logprobs",
            "payload": {
                "input_ids": batched_prompt,
            },
            "input_tokens": input_tokens,
        }

    def compute_logprobs(
        self, prompt: list[list[int]] | list[int]
    ) -> APIFuture[dict[str, Any]]:
        """
        计算提示令牌的对数概率

        Parameters
        ----------
        prompt: 输入的 token id 列表

        Returns
        -------

        """
        return self.holder.run_coroutine_threadsafe(
            self.compute_logprobs_async(prompt=prompt)
        ).result()

    async def compute_logprobs_async(
        self, prompt: list[list[int]] | list[int]
    ) -> APIFuture[dict[str, Any]]:
        """
        计算提示令牌的对数概率

        Parameters
        ----------
        prompt: 输入的 token id 列表

        Returns
        -------

        """
        prep = self._compute_logprobs_prepare(prompt=prompt)
        if prep is None:
            return APIFuture.completed({})

        response = await self.socket.send_event_async(
            event_type=prep["type"],
            payload=prep["payload"],
            input_tokens=prep["input_tokens"],
        )
        return response.map_dict()

    def get_tokenizer(self):
        """
        获取与基础模型匹配的 tokenizer，基于 transformers/modelscope 提供的 AutoTokenizer
        """
        try:
            # from transformers import AutoTokenizer
            from modelscope import AutoTokenizer
        except ImportError:
            raise ImportError(
                "检测到未安装 'modelscope' 库。执行此操作需要该库，"
                "请通过 'pip install modelscope transformers' 或 'pip install emc[full]' 安装。"
            )
        return AutoTokenizer.from_pretrained(self.base_model)
