# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0

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

## Key Types

| Group           | Type                      | Description                                                           |
|-----------------|---------------------------|-----------------------------------------------------------------------|
| Core Prompt     | `ParsedPrompt`            | Prompt after parsing metadata and template body.                      |
|                 | `PromptInputConfig`       | Configuration settings related to the input variables of a prompt.    |
|                 | `PromptMetadata`          | Metadata extracted from prompt frontmatter (config, tools, etc.).     |
|                 | `PromptOutputConfig`      | Configuration settings related to the expected output of a prompt.    |
|                 | `RenderedPrompt`          | Final output after rendering a prompt template.                       |
|-----------------|---------------------------|-----------------------------------------------------------------------|
| Message/Content | `DataPart`                | Content part containing structured data.                              |
|                 | `Document`                | Represents an external document used for context.                     |
|                 | `MediaContent`            | Describes the content details within a `MediaPart`.                   |
|                 | `MediaPart`               | Content part representing media (image, audio, video).                |
|                 | `Message`                 | Represents a single message in a conversation history.                |
|                 | `Part`                    | Union of parts (Text, Data, Media, Tool, Pending).                    |
|                 | `PendingMetadata`         | Defines the required metadata structure for a `PendingPart`.          |
|                 | `PendingPart`             | Content part indicating pending or awaited content.                   |
|                 | `Role`                    | Enum defining roles in a conversation (USER, MODEL, TOOL, SYSTEM).    |
|                 | `TextPart`                | Content part containing plain text.                                   |
|                 | `ToolRequestContent`      | Describes the details of a tool request within a `ToolRequestPart`.   |
|                 | `ToolRequestPart`         | Content part representing a request to invoke a tool.                 |
|                 | `ToolResponseContent`     | Describes the details of a tool response within a `ToolResponsePart`. |
|                 | `ToolResponsePart`        | Content part representing the result from a tool execution.           |
|-----------------|---------------------------|-----------------------------------------------------------------------|
| Tooling         | `ToolArgument`            | Type alias representing either a tool name or a full ToolDefinition.  |
|                 | `ToolDefinition`          | Defines a tool that can be called by a model.                         |
|                 | `ToolResolver`            | Type alias for a function resolving a tool name to a ToolDefinition.  |
|-----------------|---------------------------|-----------------------------------------------------------------------|
| Runtime         | `DataArgument`            | Runtime data (input variables, history, context) for rendering.       |
|                 | `PromptFunction`          | Protocol defining the interface for a callable prompt function.       |
|-----------------|---------------------------|-----------------------------------------------------------------------|
| Utility/Schema  | `HasMetadata`             | Base model for types that can include arbitrary metadata.             |
|                 | `JsonSchema`              | JSON schema definition. 'Any' allows flexibility.                     |
|                 | `PartialResolver`         | function resolving a partial name to a template string.               |
|                 | `Schema`                  | generic schema, represented as a dictionary.                          |
|                 | `SchemaResolver`          | function resolving a schema name to a JSON schema.                    |

## Type Relationships

### Prompt Structure

```
+-------------------+
|     BaseModel     |
+-------------------+
        |
+-------------------+
|    HasMetadata    |
+-------------------+
        |
        +------------------- PromptMetadata (+ model, tools, config, input, output, ...)
                                |
                                +-- ParsedPrompt (+ template)
                                |
                                +-- RenderedPrompt (+ messages: list[Message])
```

### Message/Content Structure

```
+-------------------+
|    HasMetadata    |
+-------------------+
        |
        +------------------- Message (role, content)
        |
        +------------------- Document (content)
                                |
                                +---> content: list[Part]

Part = Union[
    TextPart, DataPart, MediaPart,
    ToolRequestPart, ToolResponsePart, PendingPart
]

ToolRequestPart ---contains---> ToolRequestContent (name, input, ref)
ToolResponsePart --contains---> ToolResponseContent (name, output, ref)
MediaPart --------contains---> MediaContent (url, content_type)
PendingPart ------contains---> PendingMetadata (pending=True)
```
"""

from __future__ import annotations

from collections.abc import Callable
from enum import Enum
from typing import (
    Any,
    Generic,
    Literal,
    Protocol,
    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.
    """

    data: dict[str, Any] | list[dict[str, Any]] | str | list[str] | 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.
        save_path: File path where generated image output should be saved.
                   Required when ``format`` is ``'image'``.
    """

    format: Literal['json', 'text', 'image'] | None = None
    schema_: Schema | None = Field(
        default=None,
        validation_alias=AliasChoices('schema', 'schema_'),
        serialization_alias='schema',
    )
    save_path: str | None = Field(default=None, alias='save_path')
    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_image_save_path(self) -> PromptOutputConfig:
        """Validate that save_path is provided when format is 'image'.

        Returns:
            The validated instance.

        Raises:
            ValueError: If format is 'image' but save_path is not set.
        """
        if self.format == 'image' and not self.save_path:
            raise ValueError("output.save_path is required when output.format is 'image'")
        return self


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

    Can be specified in the prompt frontmatter as::

        adapter: openai                  # simple string form

        adapter:                         # full config form
          name: openai
          base_url: https://api.deepseek.com
          api_key: sk-...

    Attributes:
        name: The adapter name (e.g. 'openai', 'anthropic', 'google').
        base_url: Optional custom API base URL for third-party endpoints.
        api_key: Optional API key (prefer environment variables instead).
    """

    name: str
    base_url: str | None = None
    api_key: str | None = None


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

    Attributes:
        max_workers: The maximum number of concurrent workers for batch processing.
        output_dir: Optional directory for batch output files.
        jsonl: Whether to output batch results in JSONL format.
    """

    max_workers: int = 5
    output_dir: str | None = None
    jsonl: bool = False


class PromptMetadata(HasMetadata, Generic[ModelConfigT]):
    """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:
        name: Optional name override within the metadata itself.
        variant: Optional variant override within the metadata.
        version: Optional version override within the metadata.
        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 nested dictionary holding extension fields from the
             frontmatter, organized by namespace.
    """

    name: str | None = None
    variant: str | None = None
    version: str | None = None
    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, dict[str, Any]] | None = None
    runtime: RuntimeConfig | None = None
    model_config = ConfigDict(populate_by_name=True)


class ParsedPrompt(PromptMetadata[ModelConfigT], Generic[ModelConfigT]):
    """Represents a prompt after parsing its metadata and template.

    Attributes:
        template: The core template string, with frontmatter removed.
    """

    template: str


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(BaseModel, Generic[InputT]):
    """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(HasMetadata, Generic[InputT]):
    """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(BaseModel, Generic[OutputT]):
    """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(HasMetadata, Generic[OutputT]):
    """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) -> PendingMetadata:
        """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(str, Enum):
    """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(BaseModel, Generic[VariablesT]):
    """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(PromptMetadata[ModelConfigT], Generic[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(Protocol[ModelConfigT]):
    """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.
        """
        ...



