"""
Helper functions for extracting and converting messages from dataset examples.

This module provides utilities for the "appended messages" feature, which allows
users to specify a path to conversation messages within dataset examples that
should be appended to prompt templates when running experiments.
"""

import json
from typing import Any, Iterable, Mapping, Optional, Sequence, TypedDict

from typing_extensions import Required, assert_never

from phoenix.db.types.prompts import PromptChatTemplate, PromptTemplateFormat
from phoenix.server.api.types.ChatCompletionMessageRole import ChatCompletionMessageRole
from phoenix.utilities.template_formatters import (
    FStringTemplateFormatter,
    MustacheTemplateFormatter,
    NoOpFormatter,
    TemplateFormatter,
)

# These types are based loosely on the openinference.instrumentation.Message type.
# This makes it easier to leverage openinference.instrumentation helpers
# for extracting OpenInference attributes.


class PlaygroundToolCallFunction(TypedDict, total=False):
    name: str
    arguments: dict[str, Any]


class PlaygroundToolCall(TypedDict, total=False):
    id: str
    function: PlaygroundToolCallFunction


class PlaygroundMessage(TypedDict, total=False):
    role: Required[ChatCompletionMessageRole]
    content: Required[str]
    tool_call_id: str
    tool_calls: Sequence[PlaygroundToolCall]


def create_playground_message(
    role: ChatCompletionMessageRole,
    content: str,
    tool_call_id: Optional[str] = None,
    tool_calls: Optional[Sequence[PlaygroundToolCall]] = None,
) -> PlaygroundMessage:
    msg: PlaygroundMessage = {"role": role, "content": content}
    if tool_call_id is not None:
        msg["tool_call_id"] = tool_call_id
    if tool_calls is not None:
        msg["tool_calls"] = tool_calls
    return msg


# Mapping from OpenAI role strings to internal enum values
_ROLE_MAPPING = {
    "user": ChatCompletionMessageRole.USER,
    "assistant": ChatCompletionMessageRole.AI,
    "model": ChatCompletionMessageRole.AI,
    "system": ChatCompletionMessageRole.SYSTEM,
    "developer": ChatCompletionMessageRole.SYSTEM,
    "tool": ChatCompletionMessageRole.TOOL,
    # Also handle our internal names
    "ai": ChatCompletionMessageRole.AI,
}


def extract_value_from_path(data: dict[str, Any], path: str) -> Any:
    """
    Extract a value from a nested dictionary using dot-notation path.

    Args:
        data: The dictionary to extract from
        path: Dot-notation path (e.g., "messages" or "input.messages")

    Returns:
        The value at the specified path

    Raises:
        KeyError: If the path doesn't exist in the data
        TypeError: If the path traverses through a non-dict value
    """
    if not path:
        raise KeyError("Empty path provided")

    keys = path.split(".")
    current = data

    for key in keys:
        if not isinstance(current, dict):
            raise TypeError(f"Cannot traverse path '{path}': intermediate value is not a dict")
        if key not in current:
            raise KeyError(f"Key '{key}' not found in path '{path}'")
        current = current[key]

    return current


def _role_to_enum(role: str) -> ChatCompletionMessageRole:
    """
    Convert an OpenAI-format role string to ChatCompletionMessageRole enum.

    Args:
        role: Role string (e.g., "user", "assistant", "system", "tool")

    Returns:
        Corresponding ChatCompletionMessageRole enum value
    """
    role_lower = role.lower()
    if role_lower not in _ROLE_MAPPING:
        # Default to USER for unknown roles
        return ChatCompletionMessageRole.USER
    return _ROLE_MAPPING[role_lower]


def convert_openai_message_to_internal(message: dict[str, Any]) -> PlaygroundMessage:
    """
    Convert an OpenAI-format message to the internal PlaygroundMessage dict.

    OpenAI format:
        {"role": "user", "content": "Hello"}
        {"role": "assistant", "content": "Hi", "tool_calls": [...]}
        {"role": "tool", "content": "result", "tool_call_id": "call_123"}

    Internal format:
        PlaygroundMessage dict with role, content, tool_call_id (optional), tool_calls (optional)

    Args:
        message: Message dict in OpenAI format

    Returns:
        PlaygroundMessage dict
    """
    role = _role_to_enum(message.get("role", "user"))

    # Content can be a string or null
    content = message.get("content")
    if content is None:
        content = ""
    elif not isinstance(content, str):
        # Handle array content (multimodal) by extracting text parts
        if isinstance(content, list):
            text_parts = []
            for part in content:
                if isinstance(part, dict) and part.get("type") == "text":
                    text_parts.append(part.get("text", ""))
                elif isinstance(part, str):
                    text_parts.append(part)
            content = "\n".join(text_parts)
        else:
            content = str(content)

    # Extract tool_call_id (for tool response messages)
    tool_call_id: Optional[str] = message.get("tool_call_id")

    # Extract tool_calls (for assistant messages with function calls)
    tool_calls = message.get("tool_calls")

    return create_playground_message(role, content, tool_call_id, tool_calls)


def extract_and_convert_example_messages(
    data: dict[str, Any],
    path: str,
) -> list[PlaygroundMessage]:
    """
    Extract messages from a dataset example and convert them to internal format.

    This is the main entry point for the appended messages feature. It extracts
    a list of messages from a dataset example's input using a dot-notation path,
    then converts each message from OpenAI format to the internal PlaygroundMessage
    format.

    Args:
        data: The dataset example input dictionary
        path: Dot-notation path to the messages list

    Returns:
        List of PlaygroundMessage dicts ready to be appended to a prompt

    Raises:
        KeyError: If the path doesn't exist in the data
        TypeError: If the value at the path is not a list
        ValueError: If messages in the list are not valid message dicts
    """
    messages_raw = extract_value_from_path(data, path)

    if not isinstance(messages_raw, list):
        raise TypeError(f"Value at path '{path}' is not a list (got {type(messages_raw).__name__})")

    messages: list[PlaygroundMessage] = []
    for i, msg in enumerate(messages_raw):
        if not isinstance(msg, dict):
            raise ValueError(f"Message at index {i} is not a dict (got {type(msg).__name__})")
        messages.append(convert_openai_message_to_internal(msg))

    return messages


def prompt_chat_template_to_playground_messages(
    template: PromptChatTemplate,
) -> list[PlaygroundMessage]:
    """
    Convert a PromptChatTemplate (DB model) into the list of PlaygroundMessage
    dicts used by LLM streaming clients.

    Content part mapping:
      - text part        → content str (joined with newline for multiple text parts)
      - tool_call part   → tool_calls entry on the message
      - tool_result part → content str + tool_call_id on the message
    """
    messages: list[PlaygroundMessage] = []
    for msg in template.messages:
        role = _role_to_enum(msg.role)
        text_parts: list[str] = []
        tool_calls: list[PlaygroundToolCall] = []
        tool_call_id: Optional[str] = None

        if isinstance(msg.content, str):
            text_parts.append(msg.content)
        else:
            for part in msg.content:
                if part.type == "text":
                    text_parts.append(part.text)
                elif part.type == "tool_call":
                    try:
                        parsed_args: Any = json.loads(part.tool_call.arguments)
                    except (ValueError, TypeError):
                        parsed_args = {}
                    tool_calls.append(
                        PlaygroundToolCall(
                            id=part.tool_call_id,
                            function=PlaygroundToolCallFunction(
                                name=part.tool_call.name,
                                arguments=parsed_args,
                            ),
                        )
                    )
                elif part.type == "tool_result":
                    tool_call_id = part.tool_call_id
                    result = part.tool_result
                    if isinstance(result, str):
                        text_parts.append(result)
                    elif result is not None:
                        text_parts.append(json.dumps(result))

        content = "\n".join(text_parts)
        messages.append(
            create_playground_message(
                role=role,
                content=content,
                tool_call_id=tool_call_id,
                tool_calls=tool_calls if tool_calls else None,
            )
        )
    return messages


def build_template_variables(
    *,
    input_data: dict[str, Any],
    output_data: Any,
    metadata: dict[str, Any],
    template_variables_path: Optional[str],
) -> Any:
    """
    Build template variables for a dataset revision based on the configured path.

    This function constructs the full context dictionary with input, reference (output),
    and metadata, then extracts the appropriate subset based on the template_variables_path
    configuration.

    Args:
        input_data: The dataset example input dictionary
        output_data: The dataset example expected output (reference)
        metadata: The dataset example metadata dictionary
        template_variables_path: Dot-notation path to extract variables from context,
                                or empty string/None to use the full context

    Returns:
        Dictionary of template variables to use for prompt formatting

    Raises:
        KeyError: If the path doesn't exist in the context
        TypeError: If the path traverses through a non-dict value
    """
    # Build the full context with input, reference (expected output), and metadata
    full_context: dict[str, Any] = {
        "input": input_data,
        "reference": output_data,
        "metadata": metadata,
    }

    # Resolve template variables based on the configured path
    if template_variables_path:
        return extract_value_from_path(full_context, template_variables_path)
    else:
        return full_context


def formatted_messages(
    *,
    messages: Iterable[PlaygroundMessage],
    template_format: PromptTemplateFormat,
    template_variables: Mapping[str, Any],
) -> list[PlaygroundMessage]:
    """
    Formats the messages using the given template options.
    """
    messages_list = list(messages)
    if not messages_list:
        return []
    template_formatter = _template_formatter(template_format=template_format)
    result: list[PlaygroundMessage] = []
    for msg in messages_list:
        formatted_content = template_formatter.format(msg["content"], **template_variables)
        result.append(
            create_playground_message(
                msg["role"],
                formatted_content,
                msg.get("tool_call_id"),
                msg.get("tool_calls"),
            )
        )
    return result


def _template_formatter(template_format: PromptTemplateFormat) -> TemplateFormatter:
    """
    Instantiates the appropriate template formatter for the template format
    """
    # Use equality, not identity: ORM / DB round-trips may yield plain strings that match
    # enum values (e.g. "F_STRING") but are not the same object as PromptTemplateFormat.*.
    if template_format == PromptTemplateFormat.MUSTACHE:
        return MustacheTemplateFormatter()
    if template_format == PromptTemplateFormat.F_STRING:
        return FStringTemplateFormatter()
    if template_format == PromptTemplateFormat.NONE:
        return NoOpFormatter()
    assert_never(template_format)
