from typing import Any

from pydantic import model_validator

from ..base import BaseModel
from .model_input import ModelInput
from .tensor_data import TensorData


class Datum(BaseModel):
    model_input: ModelInput
    """输入 token ids，为数字列表"""

    loss_fn_inputs: dict[str, TensorData | Any]
    """损失函数，包含的参数有 weights, target_tokens 等"""

    @model_validator(mode="before")
    @classmethod
    def convert_tensordata(cls, data: Any) -> Any:
        if not isinstance(data, dict):
            return data

        loss_fn_inputs = data.get("loss_fn_inputs")
        if not isinstance(loss_fn_inputs, dict):
            return data

        converted = {}
        for key, value in loss_fn_inputs.items():
            converted[key] = TensorData.model_validate(value)

        return {
            **data,
            "loss_fn_inputs": converted,
        }

    def get_data(self, loss_fn: str = "cross_entropy") -> dict:
        """校验参数，并设置为传入给训练服务的参数模型"""
        if loss_fn == "cross_entropy":
            return self._get_data_cross_entropy()
        elif loss_fn == "importance_sampling":
            return self._get_data_importance_sampling()
        elif loss_fn == "ppo":
            return self._get_data_ppo()
        else:
            raise ValueError(f"Unsupported loss function type: {loss_fn}")

    def _get_data_cross_entropy(self):
        if "target_tokens" not in self.loss_fn_inputs:
            raise ValueError("target_tokens is required")
        if len(self.model_input) != len(self.loss_fn_inputs["target_tokens"]):
            raise ValueError(
                f"input tokens ({len(self.model_input)}) and target tokens ({len(self.loss_fn_inputs['target_tokens'])}) must have the same length"
            )
        if "weights" in self.loss_fn_inputs:
            weights = self.loss_fn_inputs["weights"]
        else:
            weights = TensorData(
                data=[1.0] * len(self.model_input),
                dtype="float32",
                shape=[len(self.model_input)],
            )
        return {
            "input_ids": self.model_input.tolist(),
            "target_ids": self.loss_fn_inputs["target_tokens"].tolist(),
            "weights": weights.tolist(),
        }

    def _get_data_importance_sampling(self):
        if "target_tokens" not in self.loss_fn_inputs:
            raise ValueError("target_tokens is required")
        if "logprobs" not in self.loss_fn_inputs:
            raise ValueError("logprobs is required")
        if "advantages" not in self.loss_fn_inputs:
            raise ValueError("advantages is required")
        if (
            len(self.model_input) != len(self.loss_fn_inputs["target_tokens"])
            or len(self.model_input) != len(self.loss_fn_inputs["logprobs"])
            or len(self.model_input) != len(self.loss_fn_inputs["advantages"])
        ):
            raise ValueError(
                "input tokens, target tokens, logprobs, and advantages must have the same length"
            )

        return {
            "input_ids": self.model_input.tolist(),
            "target_ids": self.loss_fn_inputs["target_tokens"].tolist(),
            "logprobs": self.loss_fn_inputs["logprobs"].tolist(),
            "advantages": self.loss_fn_inputs["advantages"].tolist(),
        }

    def _get_data_ppo(self):
        return self._get_data_importance_sampling()
