"""Pydantic V2 models mirroring the TypeScript types for Dotprompt.

This module provides Python equivalents of the core TypeScript types used in the
Dotprompt TS reference implementation.
"""

from enum import StrEnum
from typing import (
    Any,
    Literal,
    Self,
)

from pydantic import BaseModel, ConfigDict, Field, model_validator


class HasMetadata(BaseModel):
    """Base model for types that can include arbitrary metadata.

    Attributes:
        metadata: Arbitrary dictionary for tooling or informational
                  purposes.
    """

    metadata: dict[str, Any] | None = None
    model_config = ConfigDict(extra='allow')


class PromptInputConfig(BaseModel):
    """Configuration settings related to the input data of a prompt.

    Attributes:
        data: Raw input data, which can be a dict or list of dicts.
               Auto-resolved during rendering.
        files: File-based input source. Can be a single file/directory path
               (str), or a list of file/directory paths.
        chain: Whether runtime should provide this prompt's input source from
               the parent chain item output.
        defaults: Optional dictionary of fixed values merged into every
                  record.  Per-record ``data`` / ``files`` values take
                  priority over ``defaults`` when keys overlap.
    """

    data: dict[str, Any] | list[dict[str, Any]] | None = None
    files: str | list[str] | None = None
    chain: bool = False
    defaults: dict[str, Any] | None = None
    model_config = ConfigDict(populate_by_name=True, extra='forbid')

    @model_validator(mode='after')
    def _validate_input_sources(self) -> Self:
        """Ensure only one input source mode is configured.

        Returns:
            The validated instance.

        Raises:
            ValueError: If multiple input source modes are provided.
        """
        if self.data is not None and self.files is not None:
            raise ValueError('"data" and "files" are mutually exclusive; specify only one')
        if self.chain and self.data is not None:
            raise ValueError('"chain" and "data" are mutually exclusive; specify only one')
        if self.chain and self.files is not None:
            raise ValueError('"chain" and "files" are mutually exclusive; specify only one')
        return self


class PromptOutputConfig(BaseModel):
    """Configuration settings related to the expected output of a prompt.

    Attributes:
        format: Specifies the desired output format.
        schema: A JSON Schema object constraining the structure of the
                expected output. Only effective when ``format`` is
                ``'json'`` or ``'yaml'``.
        example: YAML sample text used to infer strict JSON Schema.
            Mutually exclusive with ``schema``.
        output_dir: Optional directory for output files. Supports template variables.
        file_name: File name template (without extension). Supports template variables
                   including ``{{NAME}}`` (the .prompt file stem).

    """

    format: Literal['json', 'yaml', 'txt', 'image']
    schema_value: dict[str, Any] | None = Field(
        default=None,
        validation_alias='schema',
        serialization_alias='schema',
    )
    example: str | None = None
    fixed: Any | None = None
    output_dir: str | None = None
    file_name: str | None = None

    model_config = ConfigDict(
        populate_by_name=True,
        serialize_by_alias=True,
        protected_namespaces=(),
        extra='forbid',
    )

    @model_validator(mode='before')
    @classmethod
    def _validate_schema_type(cls, data: Any) -> Any:
        """Validate and normalize ``output.schema`` and ``output.example``.

        Args:
            data: Raw model input.

        Returns:
            Unmodified model input.

        Raises:
            ValueError: If schema/example contract is invalid.
        """
        if not isinstance(data, dict):
            return data

        normalized = dict(data)

        schema_wrapper = normalized.get('schema')
        if isinstance(schema_wrapper, dict) and 'data' in schema_wrapper:
            unknown = sorted(key for key in schema_wrapper if key not in {'data', 'fixed'})
            if unknown:
                raise ValueError(
                    'output.schema wrapper only supports keys ["data", "fixed"] when using wrapped schema syntax'
                )
            normalized['schema'] = schema_wrapper.get('data')
            if normalized.get('fixed') is not None and 'fixed' in schema_wrapper:
                raise ValueError('"output.fixed" and "output.schema.fixed" are mutually exclusive; specify only one')
            if 'fixed' in schema_wrapper:
                normalized['fixed'] = schema_wrapper.get('fixed')
            data = normalized

        if data.get('schema') is not None and data.get('example') is not None:
            raise ValueError('"output.schema" and "output.example" are mutually exclusive; specify only one')

        schema_value = data.get('schema')
        if schema_value is not None and not isinstance(schema_value, dict):
            raise ValueError('output.schema must be a JSON Schema object; string schema definitions are not supported')

        example = data.get('example')
        if example is not None:
            if not isinstance(example, str):
                raise ValueError('output.example must be a YAML string (use "|" block style)')
            from dotpromptz.inputs import infer_schema_from_example

            normalized = dict(data)
            normalized['schema'] = infer_schema_from_example(example)
            return normalized

        return data

    @model_validator(mode='after')
    def _validate_schema_format(self) -> Self:
        """Validate that schema/example are only effective for structured formats.

        Returns:
            The validated instance.

        Raises:
            ValueError: If schema/example are set but format is unsupported.
        """
        if self.schema_value is not None and self.format not in ('json', 'yaml'):
            raise ValueError("output.schema/output.example are only effective when output.format is 'json' or 'yaml'")
        if self.fixed is not None and self.format not in ('json', 'yaml'):
            raise ValueError("output.fixed is only effective when output.format is 'json' or 'yaml'")
        return self


class AdapterConfig(BaseModel):
    """Configuration for an LLM adapter.

    Can be specified in the prompt frontmatter as::

        adapter: openai                  # simple string form

        adapter:                         # with group filter
          name: openai
          group: azure-east              # use specific credential group

        adapter:
          name: openai
          groups: [azure-east, azure-west]  # round-robin subset

    Attributes:
        name: The adapter name (e.g. 'openai', 'anthropic', 'google').
        group: Optional single credential group name to use.
        groups: Optional list of credential group names to rotate among.
    """

    name: str
    group: str | None = None
    groups: list[str] | None = None

    model_config = ConfigDict(extra='forbid')


class RetryConfig(BaseModel):
    """Configuration for retry behaviour.

    Attributes:
        max_retries: Maximum number of retry attempts per item (default 3).
        retry_delay: Base delay in seconds between retries (default 1.0).
            Delay grows linearly: ``retry_delay * (attempt + 1)``.
            For example, with retry_delay=1.0, delays are 1s, 2s, 3s, …
    """

    max_retries: int = Field(default=3, ge=0)
    retry_delay: float = Field(default=1.0, ge=0.0)

    model_config = ConfigDict(extra='forbid')


class RuntimeConfig(BaseModel):
    """Configuration for runtime execution parameters.

    Attributes:
        max_workers: Maximum number of concurrent workers for batch processing.
        timeout: Per-request timeout in seconds.  ``None`` means no timeout.
        retry: Retry configuration.  ``None`` disables retrying (default).
        next: Optional next ``.prompt`` path triggered per successful item.
    """

    max_workers: int = Field(default=5, ge=1)
    timeout: float | None = Field(default=None, gt=0.0)
    retry: RetryConfig | None = None
    next: str | None = None

    model_config = ConfigDict(extra='forbid')


class GenerateConfig(BaseModel):
    """Strict configuration model for LLM generation parameters.

    Used to validate config dicts at parse time. Any unrecognized
    field (e.g. ``models`` instead of ``model``) will raise a
    ``ValidationError``.

    Attributes:
        model: The model name/identifier.
        thinking: Thinking/reasoning effort level.
        temperature: Sampling temperature.
        top_p: Nucleus sampling parameter.
        max_tokens: Maximum number of tokens to generate.
        stop: Stop sequences.
        frequency_penalty: Frequency penalty.
        presence_penalty: Presence penalty.
        seed: Random seed for reproducibility.
        resolution: Requested output resolution/size for image generation.
        aspect_ratio: Requested aspect ratio for image generation.
        transparent: Whether to generate transparent-background images.
    """

    model: str | None = None
    thinking: Literal['low', 'medium', 'high'] | None = None
    temperature: float | None = None
    top_p: float | None = None
    max_tokens: int | None = None
    stop: list[str] | None = None
    frequency_penalty: float | None = None
    presence_penalty: float | None = None
    seed: int | None = None
    resolution: str | None = None
    aspect_ratio: str | None = None
    transparent: bool | None = None
    model_config = ConfigDict(extra='forbid')


class PromptMetadata(HasMetadata):
    """Metadata associated with a prompt, including configuration.

    Attributes:
        version: The dotpromptz major version this prompt targets.
            When absent, defaults to the current major version.
        description: A human-readable description of the prompt's purpose.

        adapter: Configuration for the LLM adapter to use.
        config: Model-specific configuration parameters.
        input: Configuration specific to the prompt's input variables.
        output: Configuration specific to the prompt's expected output.
    """

    version: int = 1
    description: str | None = None

    adapter: str | AdapterConfig | None = None
    config: dict[str, Any] | None = None
    input: PromptInputConfig | None = None
    output: PromptOutputConfig | None = None
    runtime: RuntimeConfig | None = None
    model_config = ConfigDict(populate_by_name=True, extra='forbid')


class ParsedPrompt(PromptMetadata):
    """Represents a prompt after parsing its metadata and template.

    Attributes:
        template: The core template string, with frontmatter removed.
        warnings: List of non-fatal issues encountered during parsing.
    """

    template: str
    warnings: list[str] = Field(default_factory=list)


class TextPart(HasMetadata):
    """A content part containing a plain text string.

    Attributes:
        text: The textual content of this part.
    """

    text: str


class DataPart(HasMetadata):
    """A content part containing arbitrary structured data.

    Attributes:
        data: A dictionary representing the structured data content.
    """

    data: dict[str, Any]


class MediaContent(BaseModel):
    """Describes the content details within a `MediaPart`.

    Attributes:
        url: The URL where the media resource can be accessed.
        content_type: The MIME type of the media.
    """

    url: str
    content_type: str | None = Field(default=None, alias='contentType')
    model_config = ConfigDict(populate_by_name=True)


class MediaPart(HasMetadata):
    """A content part representing media, like an image, audio, or video.

    Attributes:
        media: A `MediaContent` object with URL and type of the media.
    """

    media: MediaContent


Part = TextPart | DataPart | MediaPart
"""Type alias for any valid content part in a `Message`."""


# Define Role as a proper enum to work with Pydantic models
class Role(StrEnum):
    """Defines the role of a participant in a conversation."""

    USER = 'user'
    MODEL = 'model'
    SYSTEM = 'system'


class Message(HasMetadata):
    """Represents a single turn or message in a conversation history.

    Attributes:
        role: The role of the originator of this message.
        content: A list of `Part` objects making up the message content.
    """

    role: Role
    content: list[Part]
    model_config = ConfigDict(arbitrary_types_allowed=True)


class RenderedPrompt(PromptMetadata):
    """The final output after a prompt template is rendered.

    Attributes:
        messages: The list of `Message` objects resulting from rendering.
    """

    messages: list[Message]
