from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field
from sinapsis_chatbots_base.templates.llm_text_completion_base import (
    LLMCompletionArgs,
    LLMInitArgs,
)
from sinapsis_core.utils.env_var_keys import SINAPSIS_CACHE_DIR
from typing_extensions import Self
from vllm.entrypoints.llm import ModelDType, QuantizationMethods, TokenizerMode
from vllm.sampling_params import StructuredOutputsParams


class vLLMInitArgs(LLMInitArgs):
    """Configuration arguments for initializing a vLLM engine.

    Attributes:
        tokenizer_mode (TokenizerMode): The tokenizer mode. ``"auto"`` will use the fast
            tokenizer if available. Defaults to ``"auto"``.
        trust_remote_code (bool): Whether to allow custom code from the model
            repository. Defaults to ``False``.
        download_dir (str): Directory to download and load the weights.
            Defaults to ``SINAPSIS_CACHE_DIR``.
        tensor_parallel_size (int): Number of GPUs to use for distributed
            execution. Defaults to ``1``.
        dtype (ModelDType): Data type for model weights and activations.
            Defaults to ``"auto"``.
        quantization (QuantizationMethods | None): Method used to quantize
            the weights. Defaults to ``None``.
        seed (int): Random seed for reproducibility. Defaults to ``0``.
        gpu_memory_utilization (float): Fraction of GPU memory to be used
            for the model executor. Defaults to ``0.9``.
        max_num_seqs (int): Maximum number of sequences per iteration.
            Defaults to ``256``.
        max_model_len (int | None): Maximum sequence length for the model.
            Defaults to ``None``.
        cpu_offload_gb (float): Amount of CPU memory (in GB) to offload
            weights to. Defaults to ``0``.
        enforce_eager (bool): Whether to enforce eager execution instead
            of CUDA graphs. Defaults to ``False``.
        disable_log_stats (bool): Whether to disable logging of periodic
            runtime statistics. Defaults to ``False``.
    """

    tokenizer_mode: TokenizerMode = "auto"
    trust_remote_code: bool = False
    download_dir: str = SINAPSIS_CACHE_DIR
    tensor_parallel_size: int = 1
    dtype: ModelDType = "auto"
    quantization: QuantizationMethods | None = None
    seed: int = 0
    gpu_memory_utilization: float = 0.9
    max_num_seqs: int = 256
    max_model_len: int | None = None
    cpu_offload_gb: float = 0
    enforce_eager: bool = False
    disable_log_stats: bool = False


class vLLMMultimodalInitArgs(vLLMInitArgs):
    """Configuration arguments for initializing a vLLM engine for multimodal models.

    Provides sensible defaults for vision-language models so they work out of the box.

    Attributes:
        trust_remote_code (bool): Whether to allow custom code from the model repository.
            Defaults to ``True`` since most VLMs require custom code.
        limit_mm_per_prompt (dict[str, int]): Maximum number of multimodal items per prompt.
            Defaults to ``{"image": 1}`` to allow one image per prompt.
    """

    trust_remote_code: bool = True
    limit_mm_per_prompt: dict[str, int] = Field(default_factory=lambda: {"image": 1})


class PropertyDefinition(BaseModel):
    """Rich definition for a single property in a structured output schema.

    Allows specifying type constraints, enumerated allowed values, nested objects,
    arrays with items, and a description to guide the LLM's output for this field.

    Nested properties can use either simple type strings (e.g., ``{"field": "string"}``)
    or full PropertyDefinition objects.

    Attributes:
        type_ (str): The JSON Schema type (e.g. 'string', 'integer', 'number',
            'boolean', 'object', 'array'). Aliased to 'type'. Defaults to `"string"`.
        description (str | None): Human-readable description of the field's purpose.
            Included in the JSON Schema to help guide the LLM. Defaults to `None`.
        enum (list[str] | None): List of allowed values. When set, the structured output
            will constrain the model to only output one of these values. Defaults to `None`.
        properties (dict[str, str | Self] | None): Nested properties when type is 'object'.
            Allows defining nested objects with their own schema. Values can be simple
            type strings or full PropertyDefinition objects. Defaults to `None`.
        required (list[str] | None): Required fields within this nested object.
            Only applicable when type is 'object'. Defaults to `None`.
        items (str | Self | None): Schema for array items when type is 'array'.
            Can be a simple type string (e.g., ``"string"``, ``"integer"``) or a full
            PropertyDefinition for complex nested structures. Defaults to `None`.
        additional_properties (bool | None): Whether additional properties not defined
            in `properties` are allowed when type is 'object'. Set to `False` to
            strictly enforce the schema, or `True` to allow extra properties.
            Defaults to `None` (not included in schema).
    """

    type_: str = Field(default="string", alias="type")
    description: str | None = None
    enum: list[str] | None = None
    properties: dict[str, str | Self] | None = None
    required: list[str] | None = None
    items: str | Self | None = None
    additional_properties: bool | None = Field(default=None, alias="additionalProperties")

    model_config = ConfigDict(populate_by_name=True)

    @staticmethod
    def _convert_property(prop: str | PropertyDefinition) -> dict[str, Any]:
        """Converts a property to JSON Schema format.

        Args:
            prop: Either a type string or a PropertyDefinition.

        Returns:
            dict[str, Any]: The JSON Schema representation of the property.
        """
        if isinstance(prop, str):
            return {"type": prop}
        return prop.to_json_schema()

    def to_json_schema(self) -> dict[str, Any]:
        """Converts to a JSON Schema property definition dict.

        Returns:
            dict[str, Any]: A dict with 'type' and optionally 'description',
                'enum', 'properties', 'required', 'items', or 'additionalProperties'
                for nested structures.
        """
        result: dict[str, Any] = {"type": self.type_}
        if self.description is not None:
            result["description"] = self.description
        if self.enum is not None:
            result["enum"] = self.enum
        if self.type_ == "object" and self.properties is not None:
            result["properties"] = {k: self._convert_property(v) for k, v in self.properties.items()}
            if self.required is not None:
                result["required"] = self.required
            if self.additional_properties is not None:
                result["additionalProperties"] = self.additional_properties
        if self.type_ == "array" and self.items is not None:
            result["items"] = self._convert_property(self.items)
        return result


PropertyDefinition.model_rebuild()


class SchemaDefinition(BaseModel):
    """JSON Schema definition for structured LLM outputs.

    Properties can be defined in multiple formats:
        - **Simple**: ``{field_name: type_string}`` (e.g. ``{"name": "string"}``)
        - **Rich**: ``{field_name: PropertyDefinition}`` with type, enum, description,
          and optionally nested properties for objects or items for arrays
        - **Nested**: Properties can contain other objects (type='object') or arrays
          (type='array') with their own item schemas

    All formats can be mixed within the same schema.

    Attributes:
        properties (dict[str, str | PropertyDefinition]): Mapping of field names to either
            a type string or a full PropertyDefinition (which supports nesting).
            Defaults to empty dict.
        required (list[str]): List of field names that must be present in the output.
            Defaults to empty list.
        additional_properties (bool | None): Whether additional properties not defined
            in `properties` are allowed at the root level. Set to `False` to strictly
            enforce the schema, or `True` to allow extra properties. Defaults to `None`.
    """

    properties: dict[str, str | PropertyDefinition] = Field(default_factory=dict)
    required: list[str] = Field(default_factory=list)
    additional_properties: bool | None = Field(default=None, alias="additionalProperties")

    @staticmethod
    def _convert_property(prop: str | PropertyDefinition) -> dict[str, Any]:
        """Converts a property to JSON Schema format.

        Args:
            prop: Either a type string or a PropertyDefinition.

        Returns:
            dict[str, Any]: The JSON Schema representation of the property.
        """
        if isinstance(prop, str):
            return {"type": prop}
        return prop.to_json_schema()

    def to_json_schema(self) -> dict[str, Any]:
        """Converts to a full JSON Schema dict.

        Simple string values are expanded to ``{"type": value}``. PropertyDefinition
        instances are converted via their own ``to_json_schema()`` method, which
        handles nested objects and arrays.

        Returns:
            dict[str, Any]: A JSON Schema object with 'type', 'properties',
                'required', and optionally 'additionalProperties'.
        """
        json_properties: dict[str, Any] = {}
        for key, prop in self.properties.items():
            json_properties[key] = self._convert_property(prop)
        result: dict[str, Any] = {
            "type": "object",
            "properties": json_properties,
            "required": self.required,
        }
        if self.additional_properties is not None:
            result["additionalProperties"] = self.additional_properties
        return result


class vLLMResponseFormat(BaseModel):
    """Defines the response format for structured vLLM outputs.

    Attributes:
        type_ (Literal["text", "json_object"]): The output format type. Use 'json_object'
            to constrain the model output to valid JSON. Aliased to 'type'. Defaults to `"text"`.
        schema_ (SchemaDefinition): Schema defining the expected JSON structure. Only used
            when type is 'json_object'. Aliased to 'schema'. Defaults to empty schema.
    """

    type_: Literal["text", "json_object"] = Field(default="text", alias="type")
    schema_: SchemaDefinition = Field(default_factory=SchemaDefinition, alias="schema")

    model_config = ConfigDict(populate_by_name=True)

    def to_vllm_format(self) -> StructuredOutputsParams | None:
        """Converts to vLLM's native StructuredOutputsParams.

        When type is 'json_object' and properties are defined, builds a full JSON Schema
        constraint. When type is 'json_object' with no properties, uses json_object mode.
        Returns None when type is 'text'.

        Returns:
            StructuredOutputsParams | None: Structured output params, or None for text mode.
        """
        if self.type_ != "json_object":
            return None
        if self.schema_.properties:
            return StructuredOutputsParams(json=self.schema_.to_json_schema())
        return StructuredOutputsParams(json_object=True)


class vLLMCompletionArgs(LLMCompletionArgs):
    """Arguments for vLLM text completion.

    Attributes:
        presence_penalty (float): Float that penalizes new tokens based on whether they
            appear in the prompt or generated text so far. Values > 0 encourage the
            model to use new tokens, while values < 0 encourage the model to repeat
            tokens. Defaults to 0.0.
        frequency_penalty (float): Float that penalizes new tokens based on their
            frequency in the prompt or generated text so far. Values > 0 encourage
            the model to use new tokens, while values < 0 encourage the model to
            repeat tokens. Defaults to 0.0.
        repetition_penalty (float): Float that penalizes new tokens based on whether
            they appear in the prompt or generated text so far. Values > 1 encourage
            the model to use new tokens, while values < 1 encourage the model to
            repeat tokens. Defaults to 1.0.
        min_p (float): Float that represents the minimum probability for a token to be
            considered, relative to the probability of the most likely token.
            Must be between 0 and 1. Defaults to 0.0.
        seed (int | None): Random seed to use for the generation. Defaults to None.
        stop (str | list[str] | None): List of strings that stop the generation when
            they are generated. The returned output will not contain the stop strings.
            Defaults to None.
        ignore_eos (bool): Whether to ignore the EOS token and continue generating
            tokens after the EOS token is generated. Defaults to False.
        max_tokens (int): Maximum number of tokens to generate per output sequence.
            Defaults to 16.
        min_tokens (int): Minimum number of tokens to generate per output sequence
            before EOS or stop tokens are generated. Defaults to 0.
        bad_words (list[str] | None): List of words that are not allowed to be
            generated. Defaults to None.
        response_format (vLLMResponseFormat): Constrains the model output to a specific format.
            Use with type 'json_object' to enforce valid JSON output, optionally with a JSON Schema.
    """

    presence_penalty: float = 0.0
    frequency_penalty: float = 0.0
    repetition_penalty: float = 1.0
    min_p: float = 0.0
    seed: int | None = None
    stop: str | list[str] | None = None
    ignore_eos: bool = False
    max_tokens: int = 16
    min_tokens: int = 0
    bad_words: list[str] | None = None
    response_format: vLLMResponseFormat = Field(default_factory=vLLMResponseFormat)
