"""Schema for mlc-chat-config"""

from typing import Any, Dict, List, Literal, Optional, Union  # noqa: UP035

from pydantic import BaseModel, Field

from mlc_llm.support.constants import MLC_CHAT_CONFIG_VERSION

from .conversation_protocol import Conversation

MLC_CHAT_SYSTEM_DEFAULT = {
    "pad_token_id": 0,
    "bos_token_id": 1,
    "eos_token_id": 2,
    "temperature": 1.0,
    "presence_penalty": 0.0,
    "frequency_penalty": 0.0,
    "repetition_penalty": 1.0,
    "top_p": 1.0,
}
"""system default values."""


class MLCChatConfig(BaseModel):
    """Fields in the dumped `mlc-chat-config.json` file."""

    # Version control
    version: str = MLC_CHAT_CONFIG_VERSION

    # use alias to avoid protected namespace conflict with pydantic
    field_model_type: str = Field(alias="model_type")
    quantization: str
    # use alias to avoid protected namespace conflict with pydantic
    field_model_config: Dict[str, Any] = Field(alias="model_config")  # noqa: UP006
    vocab_size: int
    context_window_size: int
    sliding_window_size: int
    prefill_chunk_size: int
    attention_sink_size: int
    tensor_parallel_shards: int
    pipeline_parallel_stages: int = 1
    # Configuration of text generation
    active_vocab_size: int = None
    temperature: Optional[float] = None
    presence_penalty: Optional[float] = None
    frequency_penalty: Optional[float] = None
    repetition_penalty: Optional[float] = None
    top_p: Optional[float] = None
    # Tokenizer configuration
    tokenizer_files: List[str] = Field(default_factory=list)  # noqa: UP006
    # The content of tokenizer.TokenizerInfo
    tokenizer_info: Dict[str, Any] = Field(default_factory=dict)  # noqa: UP006
    # conversation template
    conv_template: Conversation
    # extra fields from generation_config.json
    # NOTE: they are not being used for now in MLCEngine
    # but we keep them for book-keep purposes
    pad_token_id: Optional[int] = None
    bos_token_id: Optional[int] = None
    eos_token_id: Optional[Union[int, List[int]]] = None  # noqa: UP006

    field_model_task: Literal["chat", "embedding"] = Field(default="chat", alias="model_task")
    embedding_metadata: Optional[Dict[str, Any]] = None  # noqa: UP006

    def get_system_defaults_for_missing_fields(self) -> Dict[str, Any]:  # noqa: UP006
        """Apply system default value for fields that are None

        Note
        ----
        We implement default setting in this way so we can lazily create
        MLCChatConfig, override its optional values then
        apply_system_defaults in the end.
        """
        res = {}
        for key, value in MLC_CHAT_SYSTEM_DEFAULT.items():
            if getattr(self, key) is None:
                res[key] = value
        return res