from __future__ import annotations

from typing import Any, Dict, Optional

from transformers import AutoConfig, PretrainedConfig


class OpenThaiSystemOneConfig(PretrainedConfig):
    """Config = a text-tower config (Qwen3.5 text) + slot-head settings [+ vision tower + point head for the -Vision variant].

    vision_config is None for the text-only line (v0.x): the model then wraps `Qwen3_5TextModel` and its state-dict
    layout is unchanged. With a vision_config the model wraps the multimodal `Qwen3_5Model` (visual + language_model)
    and gains a PointHead for `point` questions.
    """

    model_type = "openthai_systemone"
    sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig}

    def __init__(
        self,
        text_config: Optional[Dict[str, Any] | PretrainedConfig] = None,
        n_slots: int = 256,
        abstain_slot: int = 255,
        answer_token_id: Optional[int] = None,
        head_bias: bool = True,
        n_temperatures: int = 3,  # per question type: choice / score / noul (/ point)
        vision_config: Optional[Dict[str, Any] | PretrainedConfig] = None,
        image_token_id: Optional[int] = None,
        vision_start_token_id: Optional[int] = None,
        vision_end_token_id: Optional[int] = None,
        point_token_id: Optional[int] = None,
        point_head: Optional[Dict[str, Any]] = None,  # {"dim": 256, "n_sa_layers": 1, "n_heads": 8}
        **kwargs,
    ):
        if isinstance(text_config, dict):
            text_config = AutoConfig.for_model(**text_config) if "model_type" in text_config else AutoConfig.for_model("qwen3_5_text", **text_config)
        if isinstance(vision_config, dict):
            vc = dict(vision_config)
            vision_config = AutoConfig.for_model(**vc) if "model_type" in vc else AutoConfig.for_model("qwen3_5_vision", **vc)
        self.text_config = text_config
        self.vision_config = vision_config
        self.n_slots = n_slots
        self.abstain_slot = abstain_slot
        self.answer_token_id = answer_token_id
        self.head_bias = head_bias
        self.n_temperatures = n_temperatures
        self.image_token_id = image_token_id
        self.vision_start_token_id = vision_start_token_id
        self.vision_end_token_id = vision_end_token_id
        self.point_token_id = point_token_id
        self.point_head = point_head
        super().__init__(**kwargs)

    @property
    def is_vision(self) -> bool:
        return self.vision_config is not None

    @property
    def hidden_size(self) -> int:
        return self.text_config.hidden_size

    @property
    def vocab_size(self) -> int:
        return self.text_config.vocab_size