"""Parses dotprompt templates.

Extracts YAML frontmatter and converts the template body into
structured messages with text, media, and metadata parts.
"""

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

import yaml
from pydantic import ValidationError

from dotpromptz.typing import (
    GenerateConfig,
    MediaContent,
    MediaPart,
    Message,
    ParsedPrompt,
    Part,
    Role,
    TextPart,
)

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 media markers in the template.
MEDIA_MARKER_PREFIX = '<<<dotprompt:media:'


# 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>>> markers in the template.
#
# Examples of matching patterns:
# - <<<dotprompt:role:user>>>
# - <<<dotprompt:role:system>>>
#
# Note: Only lowercase letters are allowed after 'role:'.
ROLE_MARKER_REGEX = re.compile(r'(<<<dotprompt:role:[a-z]+)>>>')

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

# 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',
    'input',
    'output',
    'runtime',
    'toolDefs',
    'tools',
    '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_markers(rendered_string: str) -> list[str]:
    """Splits a rendered string into pieces based on role 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_MARKER_REGEX)


def split_by_media_markers(source: str) -> list[str]:
    """Split the source into pieces based on media 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_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 _build_input_config(raw_input: Any) -> dict[str, Any] | None:
    """Build a ``PromptInputConfig``-compatible dict from raw frontmatter ``input``.

    Supports three formats:
      - **Legacy dict/list**: ``input`` is a dict or list of dicts.
        Wrapped as ``{'data': raw_input}``.
      - **Legacy string/list-of-strings**: ``input`` is a file path string
        or list of file path strings.  Wrapped as ``{'files': raw_input}``
        (file-based input is handled via the ``files`` field).
      - **New (structured)**: ``input`` is a dict containing ``defaults``,
        ``data``, or ``files`` keys.  Forwarded as-is.

    The heuristic: if ``raw_input`` is a dict and contains any of the keys
    ``"defaults"``, ``"data"``, or ``"files"``, it is treated as the
    new structured format.  Otherwise legacy.

    Args:
        raw_input: The value of the ``input`` key from YAML frontmatter.

    Returns:
        A dict suitable for constructing ``PromptInputConfig``, or ``None``
        if ``raw_input`` is ``None``.
    """
    if raw_input is None:
        return None

    # New structured format: dict with recognised config keys
    _structured_keys = {'defaults', 'data', 'files'}
    if isinstance(raw_input, dict) and _structured_keys & raw_input.keys():
        result: dict[str, Any] = {}
        if 'defaults' in raw_input:
            result['defaults'] = raw_input['defaults']
        if 'data' in raw_input:
            result['data'] = raw_input['data']
        if 'files' in raw_input:
            result['files'] = raw_input['files']
        return result

    # Legacy string / list-of-strings → redirect to files
    if isinstance(raw_input, str):
        return {'files': raw_input}
    if isinstance(raw_input, list) and raw_input and all(isinstance(item, str) for item in raw_input):
        return {'files': raw_input}

    # Legacy dict or list-of-dicts → wrap as data
    return {'data': raw_input}

def parse_document(source: str) -> 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.

    Returns:
        Parsed prompt with metadata and template content.

    Raises:
        yaml.YAMLError: If the frontmatter YAML is invalid.
        ValidationError: If Pydantic validation fails (always raised for
            value-level validators such as ``save_path`` checks).
        Exception: If ``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(config=None, metadata={}, tool_defs=None, template=source)

    parsed_metadata = yaml.safe_load(frontmatter)
    if parsed_metadata is None:
        parsed_metadata = {}
    if not isinstance(parsed_metadata, dict):
        raise ValueError(
            'Frontmatter must be a YAML mapping (object/dict), '
            f'got {type(parsed_metadata).__name__}.'
        )
    if any(not isinstance(k, str) for k in parsed_metadata):
        raise ValueError('Frontmatter keys must be strings.')

    # Strict validation: reject any unrecognized top-level keys.
    unrecognized: list[str] = []
    for key in parsed_metadata:
        # Dotted keys (e.g. 'foo.bar') are always unrecognized.
        if '.' in key:
            unrecognized.append(key)
            continue
        if key not in RESERVED_METADATA_KEYWORDS:
            unrecognized.append(key)
    if unrecognized:
        raise ValueError(
            f'Unrecognized frontmatter field(s): {sorted(unrecognized)}. '
            f'Valid fields are: {sorted(RESERVED_METADATA_KEYWORDS)}'
        )

    # version is required whenever frontmatter exists.
    if 'version' not in parsed_metadata:
        raise ValueError('Missing required frontmatter field: version')

    # Validate config dict against GenerateConfig to catch typos.
    raw_config = parsed_metadata.get('config')
    if raw_config is not None and isinstance(raw_config, dict):
        try:
            GenerateConfig(**raw_config)
        except ValidationError as config_err:
            raise ValueError(f'Invalid config field(s): {config_err}') from config_err

    # Wrap raw input value in PromptInputConfig format.
    raw_input = parsed_metadata.get('input')
    input_config = _build_input_config(raw_input)

    return ParsedPrompt(
        description=parsed_metadata.get('description'),
        version=parsed_metadata['version'],
        adapter=parsed_metadata.get('adapter'),
        input=input_config,
        output=parsed_metadata.get('output'),
        tool_defs=parsed_metadata.get('toolDefs'),
        tools=parsed_metadata.get('tools'),
        runtime=parsed_metadata.get('runtime'),
        config=parsed_metadata.get('config'),
        metadata={},
        template=body.strip(),
    )


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

    Processes role markers to structure the conversation.

    Args:
        rendered_string: The rendered template string to convert

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

    for piece in split_by_role_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)

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

    return message_sources_to_messages(message_sources)


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 to_parts(source: str) -> list[Part]:
    """Converts a source string into an array of parts.

    Also processes media 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_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)
    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_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)
