"""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 collections.abc import Callable
from enum import StrEnum
from typing import (
    Any,
    Literal,
    Protocol,
    Self,
    TypeVar,
)

from pydantic import AliasChoices, BaseModel, ConfigDict, Field, model_validator

Schema = Any
"""Type alias for a generic schema."""

JsonSchema = Any
"""Type alias for a JSON schema definition. 'Any' allows flexibility."""

ModelConfigT = TypeVar('ModelConfigT')
"""Generic TypeVar for model configuration within prompts."""

InputT = TypeVar('InputT')
"""Generic TypeVar for input types, typically used in ToolRequestPart."""

OutputT = TypeVar('OutputT')
"""Generic TypeVar for output types, typically used in ToolResponsePart."""

VariablesT = TypeVar('VariablesT', bound=dict[str, Any])
"""Generic TypeVar for prompt input variables within DataArgument."""


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 ToolDefinition(BaseModel):
    """Defines the structure and schemas for a tool callable by a model.

    Attributes:
        name: The unique identifier for the tool.
        description: A human-readable explanation of the tool's purpose
                     and function.
        input_schema: A schema definition for the expected input
                      parameters of the tool.
        output_schema: An optional schema definition for the structure of
                       the tool's output.
    """

    name: str
    description: str | None = None
    input_schema: Schema = Field(..., alias='inputSchema')
    output_schema: Schema | None = Field(default=None, alias='outputSchema')
    model_config = ConfigDict(populate_by_name=True)


ToolArgument = str | ToolDefinition
"""Type alias representing either a tool name or a full ToolDefinition."""


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

    Attributes:
        data: Raw input data, which can be a dict, list of dicts,
               file reference (string), or list of file references.
               Auto-resolved during rendering.
        defaults: Optional dictionary of fixed values merged into every
                  record.  Per-record ``data`` values take priority over
                  ``defaults`` when keys overlap.
    """

    data: dict[str, Any] | list[dict[str, Any]] | str | list[str] | None = None
    defaults: dict[str, Any] | None = None
    model_config = ConfigDict(populate_by_name=True)


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

    Attributes:
        format: Specifies the desired output format.
        schema_: A schema definition constraining the structure of the
                 expected output. Can be passed as `schema` or `schema_` at runtime.
                 Only effective when ``format`` is ``'json'`` or ``'yaml'``.
        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).
        aggregate: Whether to aggregate all batch results into a single file.
    """

    format: Literal['json', 'yaml', 'txt', 'image']
    schema_: Schema | None = Field(
        default=None,
        validation_alias=AliasChoices('schema', 'schema_'),
        serialization_alias='schema',
    )
    output_dir: str | None = None
    file_name: str | None = None
    aggregate: bool = False
    model_config = ConfigDict(populate_by_name=True, protected_namespaces=())

    @property
    def schema(self) -> Schema | None:
        """Backward-compatible alias for ``schema_``."""
        return self.schema_

    @schema.setter
    def schema(self, value: Schema | None) -> None:
        """Backward-compatible alias for ``schema_``."""
        self.schema_ = value

    @model_validator(mode='after')
    def _validate_schema_format(self) -> Self:
        """Validate that schema is only set when format is 'json' or 'yaml'.

        Returns:
            The validated instance.

        Raises:
            ValueError: If schema is set but format is not 'json' or 'yaml'.
        """
        if self.schema_ is not None and self.format not in ('json', 'yaml'):
            raise ValueError("output.schema 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


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 = 3
    retry_delay: float = 1.0


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).
    """

    max_workers: int = 5
    timeout: float | None = None
    retry: RetryConfig | None = None


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

    This is a generic model, allowing the `config` field to hold
    different types of model-specific configurations specified by
    `ModelConfigT`.

    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.
        tools: A list of names referring to tools available to the model.
        tool_defs: A list of inline `ToolDefinition` objects available to
                   the model.
        config: Model-specific configuration parameters.
        input: Configuration specific to the prompt's input variables.
        output: Configuration specific to the prompt's expected output.
        raw: A dictionary holding the raw, unprocessed frontmatter parsed
             from the source.
        ext: A dictionary holding non-reserved extension fields from the
             frontmatter. Namespaced keys (e.g. ``foo.bar``) are nested.
    """

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

    adapter: str | AdapterConfig | None = None
    tools: list[str] | None = None
    tool_defs: list[ToolDefinition] | None = Field(default=None, alias='toolDefs')
    config: ModelConfigT | None = None
    input: PromptInputConfig | None = None
    output: PromptOutputConfig | None = None
    raw: dict[str, Any] | None = None
    ext: dict[str, Any] | None = None
    runtime: RuntimeConfig | None = None
    model_config = ConfigDict(populate_by_name=True)


class ParsedPrompt[ModelConfigT](PromptMetadata[ModelConfigT]):
    """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


class ToolRequestContent[InputT](BaseModel):
    """Describes the details of a tool request within a `ToolRequestPart`.

    Attributes:
        name: The name of the tool being requested.
        input: The input parameters for the tool request.
        ref: An optional reference identifier for tracking this request.
    """

    name: str
    input: InputT | None = None
    ref: str | None = None
    model_config = ConfigDict(arbitrary_types_allowed=True)


class ToolRequestPart[InputT](HasMetadata):
    """A content part representing a request to invoke a tool.

    Attributes:
        tool_request: A `ToolRequestContent` object with request details.
    """

    tool_request: ToolRequestContent[InputT] = Field(..., alias='toolRequest')
    model_config = ConfigDict(
        populate_by_name=True,
        arbitrary_types_allowed=True,
    )


class ToolResponseContent[OutputT](BaseModel):
    """Describes the details of a tool response within a `ToolResponsePart`.

    Attributes:
        name: The name of the tool that produced this response.
        output: The output data returned by the tool.
        ref: An optional reference identifier matching the request.
    """

    name: str
    output: OutputT | None = None
    ref: str | None = None
    model_config = ConfigDict(arbitrary_types_allowed=True)


class ToolResponsePart[OutputT](HasMetadata):
    """A content part representing the result from a tool execution.

    Attributes:
        tool_response: A `ToolResponseContent` object with response details.
    """

    tool_response: ToolResponseContent[OutputT] = Field(..., alias='toolResponse')
    model_config = ConfigDict(
        populate_by_name=True,
        arbitrary_types_allowed=True,
    )


class PendingMetadata(BaseModel):
    """Defines the required metadata structure for a `PendingPart`.

    Attributes:
        pending: A literal boolean True, indicating the pending state.
    """

    pending: Literal[True]
    model_config = ConfigDict(extra='allow')

    @classmethod
    def with_purpose(cls, purpose: str) -> Self:
        """Create a PendingMetadata with a purpose field.

        Args:
            purpose: The purpose of the pending part

        Returns:
            A new PendingMetadata instance with the purpose set
        """
        instance = cls(pending=True)
        # Set purpose as an extra field
        instance.__setattr__('purpose', purpose)
        return instance


class PendingPart(HasMetadata):
    """A content part indicating content is pending or awaited.

    Attributes:
        metadata: Metadata object confirming the pending state.
    """

    metadata: dict[str, Any] | None = None

    model_config = ConfigDict(
        extra='allow',
        populate_by_name=True,
    )

    def __init__(self, **data: Any) -> None:
        """Initialize a PendingPart.

        Args:
            **data: Data for the model, including a PendingMetadata object
                   under the 'metadata' key
        """
        if 'metadata' in data and isinstance(data['metadata'], PendingMetadata):
            # Convert PendingMetadata to dict for HasMetadata compatibility
            metadata_dict = data['metadata'].model_dump()
            data['metadata'] = metadata_dict
        super().__init__(**data)


Part = TextPart | DataPart | MediaPart | ToolRequestPart[Any] | ToolResponsePart[Any] | PendingPart
"""Type alias for any valid content part in a `Message` or `Document`."""


# 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'
    TOOL = 'tool'
    SYSTEM = 'system'


# Role constants with ROLE_ prefix for explicit imports
ROLE_USER = Role.USER
ROLE_MODEL = Role.MODEL
ROLE_TOOL = Role.TOOL
ROLE_SYSTEM = Role.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 Document(HasMetadata):
    """Represents an external document, often used for context.

    Attributes:
        content: A list of `Part` objects making up the document content.
    """

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


class DataArgument[VariablesT: dict[str, Any]](BaseModel):
    """Encapsulates runtime information needed to render a prompt template.

    Attributes:
        input: Values for input variables required by the template.
        docs: List of relevant `Document` objects.
        messages: List of preceding `Message` objects in the history.
        context: Arbitrary dictionary of additional context items.
    """

    input: VariablesT | None = None
    docs: list[Document] | None = None
    messages: list[Message] | None = None
    context: dict[str, Any] | None = None


SchemaResolver = Callable[[str], JsonSchema | None]
"""Type alias for a function resolving a schema name to a JSON schema."""

ToolResolver = Callable[[str], ToolDefinition | None]
"""Type alias for a function resolving a tool name to a ToolDefinition."""

PartialResolver = Callable[[str], str | None]
"""Type alias for a function resolving a partial name to a template string."""


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

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

    messages: list[Message]


class PromptFunction[ModelConfigT](Protocol):
    """Protocol defining the interface for a callable prompt function.

    Implementations are callables taking runtime data and optional
    metadata overrides, returning a ``RenderedPrompt``. They must also
    expose the parsed prompt structure via the ``prompt`` attribute.
    """

    prompt: ParsedPrompt[ModelConfigT]
    """The parsed prompt structure associated with this function."""

    def __call__(
        self,
        data: DataArgument[Any],
        options: PromptMetadata[ModelConfigT] | None = None,
    ) -> RenderedPrompt[ModelConfigT]:
        """Renders the prompt.

        Args:
            data: The runtime `DataArgument`.
            options: Optional `PromptMetadata` to merge/override.

        Returns:
            A `RenderedPrompt` object.
        """
        ...
