import logging
import sys

from .._utils import SocketClient, cuid, setup_logging
from ..cli.auth import ensure_login
from ..config import settings
from ..types import ServiceParams
from ._async_holder import AsyncHolder
from ._http import HttpClient
from .rest_client import RestClient
from .sampling_client import SamplingClient
from .training_client import TrainingClient

logger = logging.getLogger(__name__)


class ServiceClient:
    def __init__(
        self,
        host: str | None = None,
        api_key: str | None = None,
    ):
        """
        ServiceClient 是初始化 EMC API 的入口。它提供以下方法：
        - 为模型训练工作流程生成 TrainingClient 实例
        - 生成用于文本生成和推理的 SamplingClient 实例
        - 其它与模型相关的 RestClient 实例

        Parameters
        ----------
        host: 服务地址, 可以自定义为本地服务用于调试
        api_key: 直接传入 API Key
        """
        # 如果传递域名或者 ip 则修改为最新的域名或者 ip
        if host is not None:
            settings.host = host
        if api_key is not None:
            settings.api_key = api_key

        # 1. 获取用户信息
        try:
            user_info = ensure_login(
                api_key=api_key,
                relogin=False,
                save=False if api_key else True,
            )
        except Exception as e:
            logger.error(f"🙃 登录失败: {e}")
            sys.exit(1)
        self.user_info = user_info
        self.user_id: str = user_info.get("user_id", "")
        self.username: str = user_info.get("username", "")

        if not self.user_id and not self.username:
            logger.error("🙃 无法获取用户信息")
            sys.exit(1)

        # 初始化日志记录器
        setup_logging()

        logger.info(f"用户 {user_info['username']}, 连接服务到 {settings.base_url}")

        # 2. 随机生成会话 id，12 位字符串
        self.run_id = cuid(12)

        # 3. 创建 socket 连接会话
        socket = SocketClient(
            run_id=self.run_id,
            user_id=self.user_id,
            username=self.username,
        )
        socket.connect()
        self.socket = socket

        # 4. 获取可用模型列表
        http_client = HttpClient(api_key=settings.api_key)
        self.http_client = http_client
        self.supported_models = http_client.list_models()

        # 5. 创建一个异步任务执行器
        self.holder = AsyncHolder()

    def get_supported_models(self) -> list[str]:
        """
        获取可用的模型列表，例如 ['Qwen/Qwen2.5-3B']
        """
        return self.supported_models

    def create_lora_training_client(
        self,
        base_model: str,
        rank: int = 32,
        seed: int | None = None,
        train_mlp: bool = True,
        train_attn: bool = True,
        train_unembed: bool = False,
        trainable_token_indices: list[int] | dict[str, list[int]] | None = None,
        lora_path: str | None = None,
    ) -> TrainingClient:
        """
        创建一个用于 LoRA 微调的任务

        Parameters
        ----------
        base_model: 基础模型，例如 'Qwen/Qwen2.5-3B'
        rank: LoRA rank，默认值为 32，范围在 4-64
        seed: 用于初始化的随机种子
        train_mlp: 是否训练 MLP 层（默认为 True）
        train_attn: 是否训练注意力层（默认为 True）
        train_unembed: 是否训练 lm_head 层（默认为 False）
        trainable_token_indices: 启用指定 token 的 Embedding 训练，例如 [1,2,3] 或者 {"embed_tokens": [1,2,3]}，传入 None 表示不启用

        Returns
        -------
        TrainingClient
            包含训练状态的 TrainingClient 实例

        """
        return self.holder.run_coroutine_threadsafe(
            self.create_lora_training_client_async(
                base_model=base_model,
                rank=rank,
                seed=seed,
                train_mlp=train_mlp,
                train_attn=train_attn,
                train_unembed=train_unembed,
                trainable_token_indices=trainable_token_indices,
                lora_path=lora_path,
            )
        ).result()

    async def create_lora_training_client_async(
        self,
        base_model: str,
        rank: int = 32,
        seed: int | None = None,
        train_mlp: bool = True,
        train_attn: bool = True,
        train_unembed: bool = False,
        trainable_token_indices: list[int] | dict[str, list[int]] | None = None,
        lora_path: str | None = None,
    ) -> TrainingClient:
        """
        创建一个用于 LoRA 微调的任务

        Parameters
        ----------
        base_model: 基础模型，例如 'Qwen/Qwen2.5-3B'
        rank: LoRA rank，默认值为 32，范围在 4-64
        seed: 用于初始化的随机种子
        train_mlp: 是否训练 MLP 层（默认为 True）
        train_attn: 是否训练注意力层（默认为 True）
        train_unembed: 是否训练 lm_head 层（默认为 False）
        trainable_token_indices: 启用指定 token 的 Embedding 训练，例如 [1,2,3] 或者 {"embed_tokens": [1,2,3]}，传入 None 表示不启用

        Returns
        -------
        TrainingClient
            包含训练状态的 TrainingClient 实例

        """
        # 参数校验
        if base_model not in self.supported_models:
            logger.error(
                f"🙃 模型 {base_model} 并不被支持，当前支持的模型列表为: {self.supported_models}"
            )
            sys.exit(1)
        if rank not in range(4, 65):
            logger.error(f"🙃 rank 参数范围为 4-64，传入的值 {rank} 不在允许范围内")
            sys.exit(1)
        if trainable_token_indices is not None:
            if train_unembed:
                logger.error("🙃 trainable_token_indices 与 train_unembed 不能同时启用")
                sys.exit(1)
            if not trainable_token_indices:
                logger.error("🙃 trainable_token_indices 不能为空列表或字典")
                sys.exit(1)

        # 创建一个训练任务
        task_id = cuid(12)  # 随机生成任务 id
        lora = {
            "base_model": base_model,
            "rank": rank,
            "seed": seed,
            "train_mlp": train_mlp,
            "train_attn": train_attn,
            "train_unembed": train_unembed,
            "trainable_token_indices": trainable_token_indices,
            "lora_path": lora_path,
        }

        task_info = self.http_client.create_task(
            "TRAIN", self.run_id, base_model, self.user_id, task_id, lora=lora
        )

        socket_task = self.socket.for_task("TRAIN", task_id, task_info["runnerId"])
        # 首先初始化一个 Lora 模型

        future = await socket_task.send_event_async("init-lora", lora)
        await future

        return TrainingClient(
            task_id,
            base_model,
            lora,
            socket_task,
            ServiceParams(
                http_client=self.http_client,
                run_id=self.run_id,
                user_id=self.user_id,
                socket=self.socket,
            ),
            holder=self.holder,
        )

    def create_sampling_client(
        self,
        base_model: str = "",
        model_path: str | None = None,
    ) -> SamplingClient:
        """
        创建一个模型采样客户端

        Parameters
        ----------
        model_path: Lora 模型路径
        base_model: 基础模型，例如 Qwen/Qwen2.5-3B

        Returns
        -------
        SamplingClient
            模型采样客户端实例

        """
        return self.holder.run_coroutine_threadsafe(
            self.create_sampling_client_async(
                base_model=base_model,
                model_path=model_path,
            )
        ).result()

    async def create_sampling_client_async(
        self,
        base_model: str = "",
        model_path: str | None = None,
    ) -> SamplingClient:
        """
        创建一个模型采样客户端

        Parameters
        ----------
        model_path: Lora 模型路径
        base_model: 基础模型，例如 Qwen/Qwen2.5-3B

        Returns
        -------
        SamplingClient
            模型采样客户端实例

        """
        # 参数校验
        if base_model is None:
            logger.error("🙃 必须传入 base_model 参数指定基础模型")
            sys.exit(1)
        if base_model not in self.supported_models:
            logger.error(
                f"🙃 模型 {base_model} 并不被支持，当前支持的模型列表为: {self.supported_models}"
            )
            sys.exit(1)

        # 创建推理采样任务
        task_id = cuid(12)  # 随机生成任务 id
        task_info = self.http_client.create_task(
            "SAMPLE", self.run_id, base_model, self.user_id, task_id
        )

        socket_task = self.socket.for_task("SAMPLE", task_id, task_info["runnerId"])
        # 如果传递有 Lora 模型路径则可以尝试初始化
        if model_path:
            param = {"lora_path": model_path, "lora_name": task_id}
            await socket_task.send_event_async("load-lora", param)

        return SamplingClient(
            task_id=task_id,
            base_model=base_model,
            socket=socket_task,
            holder=self.holder,
        )

    def create_rest_client(self) -> RestClient:
        """创建 Rest 客户端"""
        return RestClient(self.http_client, self.user_info)

    def _create_rest_client_from_state_prepare(self, path: str) -> TrainingClient:
        rest_client = self.create_rest_client()
        checkpoint_info: dict = rest_client.get_checkpoint_info(path)
        lora_info: dict = checkpoint_info.get("lora", {})

        training_client = self.create_lora_training_client(
            base_model=checkpoint_info["model"],
            rank=lora_info.get("rank", 32),
            seed=lora_info.get("seed", None),
            train_mlp=lora_info.get("train_mlp", True),
            train_attn=lora_info.get("train_attn", True),
            train_unembed=lora_info.get("train_unembed", False),
            trainable_token_indices=lora_info.get("trainable_token_indices", None),
            lora_path=path,
        )

        return training_client

    def create_training_client_from_state(self, path: str) -> TrainingClient:
        """从已有的训练状态创建 TrainingClient 实例

        Parameters
        ----------
        path : str
            训练状态文件的路径, 即之前训练过程中保存的 checkpoint 的路径

        Returns
        -------
        TrainingClient
            包含训练状态的 TrainingClient 实例
        """
        return self.holder.run_coroutine_threadsafe(
            self.create_training_client_from_state_async(path=path)
        ).result()

    async def create_training_client_from_state_async(
        self, path: str
    ) -> TrainingClient:
        """从已有的训练状态创建 TrainingClient 实例

        Parameters
        ----------
        path : str
            训练状态文件的路径, 即之前训练过程中保存的 checkpoint 的路径

        Returns
        -------
        TrainingClient
            包含训练状态的 TrainingClient 实例
        """
        training_client = self._create_rest_client_from_state_prepare(path)

        return training_client

    def create_training_client_from_state_with_optimizer(
        self, path: str
    ) -> TrainingClient:
        """从已有的训练状态创建 TrainingClient 实例

        Parameters
        ----------
        path : str
            训练状态文件的路径, 即之前训练过程中保存的 checkpoint 的路径

        Returns
        -------
        TrainingClient
            包含训练状态的 TrainingClient 实例
        """
        return self.holder.run_coroutine_threadsafe(
            self.create_training_client_from_state_with_optimizer_async(path=path)
        ).result()

    async def create_training_client_from_state_with_optimizer_async(
        self, path: str
    ) -> TrainingClient:
        """从已有的训练状态创建 TrainingClient 实例

        Parameters
        ----------
        path : str
            训练状态文件的路径, 即之前训练过程中保存的 checkpoint 的路径

        Returns
        -------
        TrainingClient
            包含训练状态的 TrainingClient 实例
        """
        training_client = self._create_rest_client_from_state_prepare(path)

        load_future = await training_client._load_optimizer_state_async(path=path)
        await load_future

        return training_client
