import torch
from transformers import PreTrainedModel
from transformers.modeling_outputs import ModelOutput
from dataclasses import dataclass, field
from typing import Optional

from .configuration_domsense import SystemOneConfig

# 复用原始计算内核（作为独立子模块，权重名与原始 checkpoint 保持一致）
from .src_models.system_one_model import SystemOneModel as _SystemOneCore


@dataclass
class DomSenseOutput(ModelOutput):
    """DomSense 决策模型的前向输出。

    同时兼顾「选择」与「MDP 前瞻」产物，便于上层 / 用户直接消费。
    """

    loss: Optional[torch.FloatTensor] = None
    loss_calib: torch.FloatTensor | None = None
    loss_q: torch.FloatTensor | None = None
    loss_decision: torch.FloatTensor | None = None
    choices: torch.LongTensor | None = None
    choice_probs: torch.FloatTensor | None = None
    transition_probs: torch.FloatTensor | None = None
    expected_returns: torch.FloatTensor | None = None
    q_values: torch.FloatTensor | None = None
    confidences: torch.FloatTensor | None = None
    q_history: Optional[list] = None
    rewards: torch.FloatTensor | None = None
    embeddings: torch.FloatTensor | None = None
    action_embeddings: torch.FloatTensor | None = None


class SystemOneModelForDecision(PreTrainedModel):
    """Hugging Face 标准接口包装：把 DomSense 决策内核暴露成 PreTrainedModel。

    - ``from_pretrained`` 可加载：仓库内 ``config.json`` + ``model.safetensors`` 或
      原始 ``model_state_dict``。
    - ``save_pretrained`` 可导出：结构与权重统一为 HF 标准格式。
    - 内核权重名与原始训练 checkpoint 完全一致（顶层前缀 ``core.``），因此
      转换/反向转换都无损。
    """

    config_class = SystemOneConfig
    base_model_prefix = "core"
    main_input_name = "input_texts"

    def __init__(self, config: SystemOneConfig):
        super().__init__(config)
        self.config = config

        # 关键：构建原始内核（权重名系统与原训练 checkpoint 完全一致）。
        core_cfg = config.to_dict()
        self.core = _SystemOneCore(core_cfg)

        # transformers 5.x 的 PreTrainedModel.__init__ 不自动调用 post_init()，
        # 而 post_init() 负责设置 all_tied_weights_keys 等 from_pretrained 所需的内部属性。
        self.post_init()

    # ------------------------------------------------------------------ #
    # 编码器后端控制：默认 lightweight，保证自包含与离线可复现
    # ------------------------------------------------------------------ #
    def set_encoder_backend(self, backend: str):
        """切换编码器后端：'lightweight'（内置字符嵌入）或 'transformers'（distilbert）。

        注意：切换后端意味着编码器结构与权重不同，需用对应 checkpoint 加载；
        本仓库附带的权重为 lightweight 后端。切换仅在你有对应 distilbert 权重时使用。
        """
        raise NotImplementedError(
            "编码器后端切换需要对应权重；本仓库自带 lightweight 后端权重，"
            "若需 transformers(distilbert) 后端请自行加载对应 checkpoint。"
        )

    # ------------------------------------------------------------------ #
    # 前向
    # ------------------------------------------------------------------ #
    def forward(
        self,
        input_texts=None,
        action_texts=None,
        return_mdp_details: bool = False,
        num_valid_actions=None,
        labels=None,
        **kwargs,
    ):
        """前向传播。

        Args:
            input_texts: List[str] 场景文本。
            action_texts: list of list[str]，每个样本的选项文本。
            return_mdp_details: 是否返回 MDP 内部细节。
            num_valid_actions: (batch,) 有效动作数。
            labels: 可选。结构化 label dict，用于计算 joint loss（尚未接入）。
        """
        if input_texts is None and "texts" in kwargs:
            input_texts = kwargs.pop("texts")
        if input_texts is None:
            raise ValueError("input_texts 不能为空")

        out = self.core.forward(
            input_texts,
            action_texts=action_texts,
            num_valid_actions=num_valid_actions,
            return_mdp_details=return_mdp_details,
        )

        action_emb = None
        if action_texts is not None:
            action_emb, _ = self.core.encode_options(action_texts)

        return DomSenseOutput(
            choices=out["choices"],
            choice_probs=out["choice_probs"],
            transition_probs=out["transition_probs"],
            expected_returns=out["expected_returns"],
            q_values=out["q_values"],
            confidences=out["confidences"],
            q_history=out.get("q_history", []),
            rewards=out.get("rewards"),
            embeddings=out["embeddings"],
            action_embeddings=action_emb,
        )

    def generate_compact_output(self, texts, action_texts=None, num_valid_actions=None):
        """生成紧凑决策输出（与训练/服务层输出契约一致）。"""
        return self.core.generate_compact_output(
            texts, action_texts=action_texts, num_valid_actions=num_valid_actions
        )

    def count_parameters(self):
        return self.core.count_parameters()

    # ------------------------------------------------------------------ #
    # 权重初始化：不随机覆盖已有内核权重；新初始化仅用于兼容 HF 约定
    # ------------------------------------------------------------------ #
    def _init_weights(self, module):
        pass