# 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

"""Parses dotprompt templates and extracts metadata.

This module is responsible for parsing dotprompt templates. It handles extracting
metadata from YAML frontmatter and converting the template content into a
structured format, specifically a list of messages and their parts (text,
media, or metadata).

Key functionalities include:

- Extracting YAML frontmatter and the main template body.
- Parsing the YAML frontmatter into a structured metadata object, handling
  reserved keywords and namespaced entries.
- Splitting the template body into message sources based on role and history
  markers.
- Converting message sources into structured messages, processing media and
  section markers within the content.
- Handling the insertion of historical messages into the conversation flow.
"""

import logging
import re
from dataclasses import dataclass, field
from typing import Any, TypeVar

import yaml
from pydantic import ValidationError

from dotpromptz.typing import (
    DataArgument,
    MediaContent,
    MediaPart,
    Message,
    ParsedPrompt,
    Part,
    PendingMetadata,
    PendingPart,
    Role,
    TextPart,
)

logger = logging.getLogger(__name__)

T = TypeVar('T')


@dataclass
class MessageSource:
    """A message with a source string and optional content and metadata."""

    role: Role
    source: str | None = None
    content: list[Part] | None = None
    metadata: dict[str, Any] | None = field(default_factory=dict)


# Prefixes for the role markers in the template.
ROLE_MARKER_PREFIX = '<<<dotprompt:role:'

# Prefixes for the history markers in the template.
HISTORY_MARKER_PREFIX = '<<<dotprompt:history'

# Prefixes for the media markers in the template.
MEDIA_MARKER_PREFIX = '<<<dotprompt:media:'

# Prefixes for the section markers in the template.
SECTION_MARKER_PREFIX = '<<<dotprompt:section'

# Regular expression to match YAML frontmatter delineated by `---` markers.
# Allows blank lines and license headers (lines starting with #) before the first ---.
FRONTMATTER_AND_BODY_REGEX = re.compile(
    r'^(?:(?:#[^\n]*|[ \t]*)\n)*---\s*(?:\r\n|\r|\n)([\s\S]*?)(?:\r\n|\r|\n)---\s*(?:\r\n|\r|\n)([\s\S]*)$'
)

# Regular expression to match <<<dotprompt:role:xxx>>> and
# <<<dotprompt:history>>> markers in the template.
#
# Examples of matching patterns:
# - <<<dotprompt:role:user>>>
# - <<<dotprompt:role:system>>>
# - <<<dotprompt:history>>>
#
# Note: Only lowercase letters are allowed after 'role:'.
ROLE_AND_HISTORY_MARKER_REGEX = re.compile(r'(<<<dotprompt:(?:role:[a-z]+|history))>>>')

# Regular expression to match <<<dotprompt:media:url>>> and
# <<<dotprompt:section>>> markers in the template.
#
# Examples of matching patterns:
# - <<<dotprompt:media:url>>>
# - <<<dotprompt:section>>>
MEDIA_AND_SECTION_MARKER_REGEX = re.compile(r'(<<<dotprompt:(?:media:url|section).*?)>>>')

# List of reserved keywords that are handled specially in the metadata of a
# .prompt file. These keys are processed differently from extension metadata.
RESERVED_METADATA_KEYWORDS = [
    # NOTE: KEEP SORTED
    'adapter',
    'config',
    'description',
    'ext',
    'input',

    'name',
    'output',
    'raw',
    'runtime',
    'toolDefs',
    'tools',
    'variant',
    'version',
]


def split_by_regex(source: str, regex: re.Pattern[str]) -> list[str]:
    """Splits string by regexp while filtering out empty/whitespace-only pieces.

    Args:
        source: The source string to split into parts.
        regex: The regular expression to use for splitting.

    Returns:
        An array of non-empty string pieces.
    """

    def filter_empty(s: str) -> bool:
        return bool(s.strip())

    return list(filter(filter_empty, regex.split(source)))


def split_by_role_and_history_markers(rendered_string: str) -> list[str]:
    """Splits a rendered string into pieces based on role and history markers.

    Empty/whitespace-only pieces are filtered out.

    Args:
        rendered_string: The template string to split.

    Returns:
        Array of non-empty string pieces.
    """
    return split_by_regex(rendered_string, ROLE_AND_HISTORY_MARKER_REGEX)


def split_by_media_and_section_markers(source: str) -> list[str]:
    """Split the source into pieces based on media and section markers.

    Empty/whitespace-only pieces are filtered out.

    Args:
        source: The source string to split into parts

    Returns:
        An array of string parts
    """
    return split_by_regex(source, MEDIA_AND_SECTION_MARKER_REGEX)


def convert_namespaced_entry_to_nested_object(
    key: str,
    value: Any,
    obj: dict[str, dict[str, Any]] | None = None,
) -> dict[str, dict[str, Any]]:
    """Processes a namespaced key-value pair into a nested object structure.

    For example, 'foo.bar': 'value' becomes { foo: { bar: 'value' } }.

    Args:
        key: The dotted namespace key (e.g., 'foo.bar')
        value: The value to assign
        obj: The object to add the namespaced value to

    Returns:
        The updated target object
    """
    # NOTE: Goes only a single level deep.
    if obj is None:
        obj = {}

    last_dot_index = key.rindex('.')
    ns = key[:last_dot_index]
    field = key[last_dot_index + 1 :]

    # Ensure the namespace exists.
    obj.setdefault(ns, {})
    obj[ns][field] = value

    return obj


def extract_frontmatter_and_body(source: str) -> tuple[str, str]:
    """Extracts the YAML frontmatter and body from a document.

    Args:
        source: The source document containing frontmatter and template

    Returns:
        A tuple containing the frontmatter and body If the pattern does not
        match, both the values returned will be empty.
    """
    match = FRONTMATTER_AND_BODY_REGEX.match(source)
    if match:
        frontmatter, body = match.groups()
        return frontmatter, body
    return '', ''


def parse_document(source: str, *, strict: bool = False) -> ParsedPrompt[T]:
    """Parses document containing YAML frontmatter and template content.

    The frontmatter contains metadata and configuration for the prompt.

    Args:
        source: The source document containing frontmatter and template.
        strict: If ``True``, raise exceptions instead of returning a
            degraded ``ParsedPrompt``.  Default ``False`` keeps the
            legacy lenient behaviour.

    Returns:
        Parsed prompt with metadata and template content.

    Raises:
        yaml.YAMLError: If *strict* and the frontmatter YAML is invalid.
        ValidationError: If Pydantic validation fails (always raised for
            value-level validators such as ``save_path`` checks).
        Exception: If *strict* and ``ParsedPrompt`` construction fails.
    """
    frontmatter, body = extract_frontmatter_and_body(source)
    if not frontmatter:
        # No frontmatter, return a basic ParsedPrompt with just the template
        return ParsedPrompt(ext={}, config=None, metadata={}, tool_defs=None, template=source)

    try:
        parsed_metadata = yaml.safe_load(frontmatter)
        if parsed_metadata is None:
            parsed_metadata = {}

        raw = dict(parsed_metadata)
        pruned: dict[str, Any] = {'ext': {}, 'config': {}, 'metadata': {}}
        ext: dict[str, dict[str, Any]] = {}

        # Process each key in the raw metadata
        for key, value in raw.items():
            if key in RESERVED_METADATA_KEYWORDS:
                pruned[key] = value
            elif '.' in key:
                convert_namespaced_entry_to_nested_object(key, value, ext)


        try:
            # Wrap raw input value in PromptInputConfig format
            raw_input = raw.get('input')
            input_config = {'data': raw_input} if raw_input is not None else None

            return ParsedPrompt(
                name=raw.get('name'),
                description=raw.get('description'),
                variant=raw.get('variant'),
                version=raw.get('version'),
                adapter=raw.get('adapter'),
                input=input_config,
                output=raw.get('output'),
                tool_defs=raw.get('toolDefs'),
                tools=raw.get('tools'),
                runtime=raw.get('runtime'),
                ext=ext,
                config=pruned.get('config'),
                metadata=pruned.get('metadata', {}),
                raw=raw,
                template=body.strip(),
            )
        except ValidationError as ve:
            # Re-raise validation errors from model validators (e.g.
            # image format requires save_path) so callers see them
            # immediately.  Plain coercion failures (wrong type for a
            # field) are still handled gracefully below.
            if any(e['type'] == 'value_error' for e in ve.errors()):
                raise
            if strict:
                raise
            logger.warning('Dotprompt: Error building a parsed prompt object, error=%s', ve)
            # Return a basic ParsedPrompt with just the template
            return ParsedPrompt(
                ext={},
                config=None,
                metadata={},
                tool_defs=None,
                template=body.strip(),
            )
        except Exception as e:
            if strict:
                raise
            logger.warning('Dotprompt: Error building a parsed prompt object, error=%s', e)
            # Return a basic ParsedPrompt with just the template
            return ParsedPrompt(
                ext={},
                config=None,
                metadata={},
                tool_defs=None,
                template=body.strip(),
            )
    except ValidationError:
        raise
    except Exception as e:
        if strict:
            raise
        logger.warning('Dotprompt: Error parsing YAML frontmatter, error=%s', e)
        # Return a basic ParsedPrompt with just the template
        return ParsedPrompt(
            ext={},
            config=None,
            metadata={},
            tool_defs=None,
            template=source.strip(),
        )


def to_messages(
    rendered_string: str,
    data: DataArgument[Any] | None = None,
) -> list[Message]:
    """Converts a rendered template string into an array of messages.

    Processes role markers and history placeholders to structure the
    conversation.

    Args:
        rendered_string: The rendered template string to convert
        data: Optional data containing message history

    Returns:
        List of structured messages
    """
    current_message = MessageSource(role=Role.USER, source='')
    message_sources = [current_message]

    for piece in split_by_role_and_history_markers(rendered_string):
        if piece.startswith(ROLE_MARKER_PREFIX):
            role = piece[len(ROLE_MARKER_PREFIX) :]

            if current_message.source and current_message.source.strip():
                # If the current message has content, create a new message
                current_message = MessageSource(role=Role(role), source='')
                message_sources.append(current_message)
            else:
                # Otherwise, update the role of the current message
                current_message.role = Role(role)

        elif piece.startswith(HISTORY_MARKER_PREFIX):
            # Add the history messages to the message sources
            msgs: list[Message] = []
            if data and data.messages:
                msgs = data.messages
            history_messages = transform_messages_to_history(msgs)
            if history_messages:
                message_sources.extend([
                    MessageSource(
                        role=msg.role,
                        content=msg.content,
                        metadata=msg.metadata,
                    )
                    for msg in history_messages
                ])

            # Add a new message source for the model
            current_message = MessageSource(role=Role.MODEL, source='')
            message_sources.append(current_message)

        else:
            # Otherwise, add the piece to the current message source
            current_message.source = (current_message.source or '') + piece

    messages = message_sources_to_messages(message_sources)
    return insert_history(messages, data.messages if data else None)


def message_sources_to_messages(
    message_sources: list[MessageSource],
) -> list[Message]:
    """Processes an array of message sources into an array of messages.

    Args:
        message_sources: List of message sources

    Returns:
        List of structured messages
    """
    messages: list[Message] = []
    for m in message_sources:
        if m.content or m.source:
            message = Message(
                role=m.role,
                content=(m.content if m.content is not None else to_parts(m.source or '')),
            )

            if m.metadata:
                message.metadata = m.metadata

            messages.append(message)

    return messages


def transform_messages_to_history(
    messages: list[Message],
) -> list[Message]:
    """Adds history metadata to an array of messages.

    Args:
        messages: Array of messages to transform

    Returns:
        Array of messages with history metadata added
    """
    return [
        Message(
            role=message.role,
            content=message.content,
            metadata={**(message.metadata or {}), 'purpose': 'history'},
        )
        for message in messages
    ]


def messages_have_history(messages: list[Message]) -> bool:
    """Checks if the messages have history metadata.

    Args:
        messages: The messages to check

    Returns:
        True if the messages have history metadata, False otherwise
    """
    return any(msg.metadata and msg.metadata.get('purpose') == 'history' for msg in messages)


def insert_history(
    messages: list[Message],
    history: list[Message] | None = None,
) -> list[Message]:
    """Inserts historical messages into the conversation.

    The history is inserted at:
    - The end of the conversation if there is no history or no user message.
    - Before the last user message if there is a user message.

    The history is not inserted:
    - If it already exists in the messages.
    - If there is no user message.

    Args:
        messages: Current array of messages
        history: Historical messages to insert

    Returns:
        Messages with history inserted
    """
    # If we have no history or find an existing instance of history, return the
    # original messages unmodified.
    if not history or messages_have_history(messages):
        return messages

    if len(messages) == 0:
        return history

    last_message = messages[-1]
    if last_message.role == Role.USER:
        # If the last message is a user message, insert the history before it.
        messages = messages[:-1]
        messages.extend(history)
        messages.append(last_message)
    else:
        # Otherwise, append the history to the end of the messages.
        messages.extend(history)
    return messages


def to_parts(source: str) -> list[Part]:
    """Converts a source string into an array of parts.

    Also processes media and section markers.

    Args:
        source: The source string to convert into parts

    Returns:
        Array of structured parts (text, media, or metadata)
    """
    return [parse_part(piece) for piece in split_by_media_and_section_markers(source)]


def parse_part(piece: str) -> Part:
    """Parses a part from a piece of rendered template.

    Args:
        piece: The piece to parse

    Returns:
        Part, PendingPart, TextPart, or MediaPart
    """
    if piece.startswith(MEDIA_MARKER_PREFIX):
        return parse_media_part(piece)
    elif piece.startswith(SECTION_MARKER_PREFIX):
        return parse_section_part(piece)
    else:
        return parse_text_part(piece)


def parse_media_part(piece: str) -> MediaPart:
    """Parses a media part from a piece of rendered template.

    Args:
        piece: The piece to parse

    Returns:
        Media part

    Raises:
        ValueError: If the media piece is invalid
    """
    if not piece.startswith(MEDIA_MARKER_PREFIX):
        raise ValueError(f'Invalid media piece: {piece}; expected prefix {MEDIA_MARKER_PREFIX}')

    fields = piece.split(' ')
    n = len(fields)
    if n == 3:
        _, url, content_type = fields
    elif n == 2:
        _, url = fields
        content_type = None
    else:
        raise ValueError(f'Invalid media piece: {piece}; expected 2 or 3 fields, found {n}')

    media_content = MediaContent(
        url=url,
        content_type=(content_type if content_type and content_type.strip() else None),
    )
    return MediaPart(media=media_content)


def parse_section_part(piece: str) -> PendingPart:
    """Parses a section part from a piece of rendered template.

    Args:
        piece: The piece to parse

    Returns:
        Section part

    Raises:
        ValueError: If the section piece is invalid
    """
    if not piece.startswith(SECTION_MARKER_PREFIX):
        raise ValueError(f'Invalid section piece: {piece}; expected prefix {SECTION_MARKER_PREFIX}')

    fields = piece.split(' ')
    if len(fields) == 2:
        section_type = fields[1]
    else:
        raise ValueError(f'Invalid section piece: {piece}; expected 2 fields, found {len(fields)}')

    # Use the helper method to set purpose
    pending_metadata = PendingMetadata.with_purpose(section_type)
    return PendingPart(metadata=pending_metadata)


def parse_text_part(piece: str) -> TextPart:
    """Parses a text part from a piece of rendered template.

    Args:
        piece: The piece to parse

    Returns:
        Text part
    """
    return TextPart(text=piece)
