import logging
from typing import Any, Callable

import torch

from .._utils import ForTask, cuid
from ..lib.sampling_client import SamplingClient
from ..types import (
    AdamParams,
    APIFuture,
    Datum,
    ForwardBackwardOutput,
    OptimStepResponse,
    SampleInput,
    SaveWeightsForSamplerResponse,
    SaveWeightsResponse,
    ServiceParams,
    TensorData,
)
from ._async_holder import AsyncHolder

logger = logging.getLogger(__name__)

APIFutureDict = APIFuture[dict[str, Any]]


class TrainingClient:
    def __init__(
        self,
        task_id: str,
        base_model: str,
        lora: dict,
        socket: ForTask,
        service: ServiceParams,
        holder: AsyncHolder,
    ):
        # 任务 id
        self.task_id = task_id
        # 使用的基础模型
        self.base_model = base_model
        self.lora = lora

        self.socket = socket
        self.service = service
        self.holder = holder

    def _forward_prepare(
        self,
        data: list[Datum],
        loss_fn: str,
        loss_fn_config: dict[str, Any] | None,
        auto_shift: bool,
    ) -> dict[str, Any] | None:
        """整理前向请求的公共参数，返回事件体或 None。"""
        input_tokens = sum(len(sample.model_input) for sample in data)
        batch = [sample.get_data(loss_fn=loss_fn) for sample in data]

        if len(batch) == 0:
            logger.error("🙃 输入样本为空，无法执行，请检查你的输入")
            return None

        return {
            "type": "forward",
            "payload": {
                "batch": batch,
                "loss_fn": loss_fn,
                "loss_fn_config": loss_fn_config,
                "auto_shift": auto_shift,
            },
            "input_tokens": input_tokens,
        }

    def forward(
        self,
        data: list[Datum],
        loss_fn: str = "cross_entropy",
        loss_fn_config: dict[str, Any] | None = None,
        auto_shift: bool = False,
    ) -> APIFuture[ForwardBackwardOutput]:
        """
        仅前向传播，传入一个批次的 tokens (同步版本)

        Parameters
        ----------
        data : list[Datum]
            样本列表, 每个样本包含输入 token ids 和损失函数参数

        loss_fn : str
            损失函数类型, 默认为交叉熵损失, 可选值有:
            - "cross_entropy": 交叉熵损失函数
            - "importance_sampling": 重要性采样损失函数
            - "ppo": PPO 损失函数

        loss_fn_config : dict, optional
            损失函数的配置项, 不同损失函数的配置项不同

        auto_shift : bool
            是否自动对齐输入输出位置, 默认为 False
            - 如果为 True, 模型前向时会自动将 labels 偏移一位以对齐预测目标
            - 如果为 False, 则需要用户自行确保输入输出位置对齐
        """
        return self.holder.run_coroutine_threadsafe(
            self.forward_async(
                data=data,
                loss_fn=loss_fn,
                loss_fn_config=loss_fn_config,
                auto_shift=auto_shift,
            )
        ).result()

    async def forward_async(
        self,
        data: list[Datum],
        loss_fn: str = "cross_entropy",
        loss_fn_config: dict[str, Any] | None = None,
        auto_shift: bool = False,
    ) -> APIFuture[ForwardBackwardOutput]:
        """
        仅前向传播，传入一个批次的 tokens (异步版本)

        Parameters
        ----------
        data : list[Datum]
            样本列表, 每个样本包含输入 token ids 和损失函数参数

        loss_fn : str
            损失函数类型, 默认为交叉熵损失, 可选值有:
            - "cross_entropy": 交叉熵损失函数
            - "importance_sampling": 重要性采样损失函数
            - "ppo": PPO 损失函数

        loss_fn_config : dict, optional
            损失函数的配置项, 不同损失函数的配置项不同

        auto_shift : bool
            是否自动对齐输入输出位置, 默认为 False
            - 如果为 True, 模型前向时会自动将 labels 偏移一位以对齐预测目标
            - 如果为 False, 则需要用户自行确保输入输出位置对齐
        """
        prep = self._forward_prepare(
            data=data,
            loss_fn=loss_fn,
            loss_fn_config=loss_fn_config,
            auto_shift=auto_shift,
        )
        if prep is None:
            logger.error("🙃 输入无效, 请检查你的输入")
            return APIFuture.completed(
                ForwardBackwardOutput(loss_fn_outputs=[], metrics={})
            )

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

    def forward_backward_custom(
        self,
        data: list[Datum],
        loss_fn: Callable[
            [list[Datum], list[torch.Tensor]], tuple[torch.Tensor, dict[str, float]]
        ],
    ) -> APIFuture[ForwardBackwardOutput]:
        """
        自定义损失的前向+反向 (同步版本)
        """
        return self.holder.run_coroutine_threadsafe(
            self.forward_backward_custom_async(
                data=data,
                loss_fn=loss_fn,
            )
        ).result()

    async def forward_backward_custom_async(
        self,
        data: list[Datum],
        loss_fn: Callable[
            [list[Datum], list[torch.Tensor]], tuple[torch.Tensor, dict[str, float]]
        ],
    ) -> APIFuture[ForwardBackwardOutput]:
        """
        自定义损失的前向+反向 (异步版本)
        """
        if not data:
            return APIFuture.completed(
                ForwardBackwardOutput(loss_fn_outputs=[], metrics={})
            )

        forward_future = await self.forward_async(
            data=data,
            loss_fn="cross_entropy",
            loss_fn_config=None,
            auto_shift=False,
        )
        forward_result = await forward_future

        logprobs_list: list[torch.Tensor] = []
        for output in forward_result.loss_fn_outputs:
            raw = output.get("logprobs", TensorData(data=[], dtype="float32", shape=[]))
            logprob = torch.tensor(raw.data, dtype=torch.float32).reshape(raw.shape)
            logprobs_list.append(logprob.clone().detach().requires_grad_(True))

        loss, custom_metrics = loss_fn(data, logprobs_list)
        loss.backward()

        linear_loss_data: list[Datum] = []
        for datum, logprob in zip(data, logprobs_list, strict=True):
            assert logprob.grad is not None, (
                "自定义损失函数的输出无梯度, 请确保 loss_fn 返回的损失张量具有梯度"
            )

            grad = logprob.grad.clone().detach()
            linear_loss_data.append(
                SampleInput(
                    model_input=datum.model_input,
                    loss_fn_inputs={
                        "target_tokens": datum.loss_fn_inputs["target_tokens"],
                        "weights": (-grad).cpu().tolist(),
                    },
                )
            )

        backend_future = await self.forward_backward_async(
            data=linear_loss_data,
            loss_fn="cross_entropy",
            loss_fn_config=None,
            auto_shift=False,
        )

        return backend_future.map(
            lambda r: ForwardBackwardOutput(
                loss_fn_outputs=r.loss_fn_outputs,
                metrics={**r.metrics, **custom_metrics},
            )
        )

    def forward_backward(
        self,
        data: list[Datum],
        loss_fn: str = "cross_entropy",
        loss_fn_config: dict[str, Any] | None = None,
        auto_shift: bool = False,
    ) -> APIFuture[ForwardBackwardOutput]:
        """
        前向+反向传播 (同步版本)

        Parameters
        ----------

        data : list[Datum]
            样本列表, 每个样本包含输入 token ids 和损失函数参数

        loss_fn : str
            损失函数类型, 默认为交叉熵损失, 可选值有:
            - "cross_entropy": 交叉熵损失函数
            - "importance_sampling": 重要性采样损失函数
            - "ppo": PPO 损失函数

        loss_fn_config : dict, optional
            损失函数的配置项, 不同损失函数的配置项不同

        auto_shift : bool
            是否自动对齐输入输出位置, 默认为 False
            - 如果为 True, 模型前向时会自动将 labels 偏移一位以对齐预测目标
            - 如果为 False, 则需要用户自行确保输入输出位置对齐
        """
        return self.holder.run_coroutine_threadsafe(
            self.forward_backward_async(
                data=data,
                loss_fn=loss_fn,
                loss_fn_config=loss_fn_config,
                auto_shift=auto_shift,
            )
        ).result()

    async def forward_backward_async(
        self,
        data: list[Datum],
        loss_fn: str = "cross_entropy",
        loss_fn_config: dict[str, Any] | None = None,
        auto_shift: bool = False,
    ) -> APIFuture[ForwardBackwardOutput]:
        """
        前向+反向传播 (异步版本)

        Parameters
        ----------

        data : list[Datum]
            样本列表, 每个样本包含输入 token ids 和损失函数参数

        loss_fn : str
            损失函数类型, 默认为交叉熵损失, 可选值有:
            - "cross_entropy": 交叉熵损失函数
            - "importance_sampling": 重要性采样损失函数
            - "ppo": PPO 损失函数

        loss_fn_config : dict, optional
            损失函数的配置项, 不同损失函数的配置项不同

        auto_shift : bool
            是否自动对齐输入输出位置, 默认为 False
            - 如果为 True, 模型前向时会自动将 labels 偏移一位以对齐预测目标
            - 如果为 False, 则需要用户自行确保输入输出位置对齐
        """
        prep = self._forward_prepare(
            data=data,
            loss_fn=loss_fn,
            loss_fn_config=loss_fn_config,
            auto_shift=auto_shift,
        )
        if prep is None:
            logger.error("🙃 输入无效, 请检查你的输入")
            return APIFuture.completed(
                ForwardBackwardOutput(loss_fn_outputs=[], metrics={})
            )
        else:
            prep["type"] = "forward-backward"

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

    def optim_step(self, adam_params: AdamParams) -> APIFuture[OptimStepResponse]:
        """
        梯度优化 (同步版本)
        """
        return self.holder.run_coroutine_threadsafe(
            self.optim_step_async(adam_params=adam_params)
        ).result()

    async def optim_step_async(
        self, adam_params: AdamParams
    ) -> APIFuture[OptimStepResponse]:
        """
        梯度优化 (异步版本)
        """
        prep = {
            "type": "optimize-step",
            "payload": adam_params.model_dump(),
        }
        response = await self.socket.send_event_async(
            event_type=prep["type"],
            payload=prep["payload"],
        )
        return response.map_model(OptimStepResponse)

    def _save_state_prepare(self, name: str) -> dict[str, Any]:
        """整理保存 checkpoint 的公共参数。"""
        save_embedding_layers = (
            self.lora.get("train_unembed", False)
            or self.lora.get("trainable_token_indices") is not None
        )

        return {
            "type": "save-checkpoint",
            "payload": {
                "name": name,
                "save_embedding_layers": save_embedding_layers,
            },
        }

    def save_state(self, name: str) -> APIFuture[SaveWeightsResponse]:
        """
        保留权重和优化器状态 (同步版本)
        """
        return self.holder.run_coroutine_threadsafe(
            self.save_state_async(name=name)
        ).result()

    async def save_state_async(self, name: str) -> APIFuture[SaveWeightsResponse]:
        """
        保留权重和优化器状态 (异步版本)
        """
        prep = self._save_state_prepare(
            name=name,
        )
        response = await self.socket.send_event_async(
            event_type=prep["type"],
            payload=prep["payload"],
        )
        return response.map_model(SaveWeightsResponse)

    def save_weights_for_sampler(
        self, name: str
    ) -> APIFuture[SaveWeightsForSamplerResponse]:
        """
        只保留模型权重 (同步版本)
        """
        return self.holder.run_coroutine_threadsafe(
            self.save_weights_for_sampler_async(name=name)
        ).result()

    async def save_weights_for_sampler_async(
        self, name: str
    ) -> APIFuture[SaveWeightsForSamplerResponse]:
        """
        只保留模型权重 (异步版本)
        """
        prep = self._save_state_prepare(
            name=name,
        )

        prep["type"] = "save-weights"
        response = await self.socket.send_event_async(
            event_type=prep["type"],
            payload=prep["payload"],
        )
        return response.map_model(SaveWeightsForSamplerResponse)

    # ========== 断点续训: 加载状态 (TODO: 未对齐) ==========

    async def _load_optimizer_state_async(self, path: str) -> APIFutureDict:
        """
        加载优化器状态，用于断点续训 (异步版本)

        Parameters
        ----------
        path : str
            checkpoint 路径（相对于 SAVE_DIR），由 save_state 返回的 path 字段
        """
        response = await self.socket.send_event_async(
            event_type="load-state",
            payload={"path": path},
        )
        return response.map_dict()

    def get_tokenizer(self):
        """
        获取与基础模型匹配的 tokenizer，基于 transformers/modelscope 提供的 AutoTokenizer

        Raises
        ------
        ImportError
            未安装 `modelscope` 依赖
        """
        try:
            # from transformers import AutoTokenizer
            from modelscope import AutoTokenizer
        except ImportError:
            raise ImportError(
                "检测到未安装 'modelscope' 库。执行此操作需要该库，"
                "请通过 'pip install modelscope transformers' 或 'pip install trio[full]' 安装。"
            )
        return AutoTokenizer.from_pretrained(self.base_model)

    def create_sampling_client(self, model_path: str) -> SamplingClient:
        """
        创建推理采样任务客户端（同步方法）

        Parameters
        ----------
        model_path : str
            LoRA 权重路径
        """
        # 创建推理采样任务
        task_id = cuid(12)
        task_info = self.service.http_client.create_task(
            "SAMPLE",
            self.service.run_id,
            self.base_model,
            self.service.user_id,
            task_id,
        )

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