"""Utility functions for processing and formatting LLM-related message data."""

import json
import re

from loguru import logger

from ..enumeration import Role
from ..schema import Message, Trajectory, MemoryNode


def format_messages(
    messages: list[Message | dict],
    add_index: bool = True,
    add_time: bool = True,
    use_name: bool = True,
    add_reasoning: bool = True,
    add_tools: bool = True,
    strip_markdown_headers: bool = True,
    enable_system: bool = False,
) -> str:
    """Formats a list of messages into a single string, optionally filtering system roles."""
    formatted_lines = []
    for i, message in enumerate(messages):
        if isinstance(message, dict):
            message = Message(**message)
        if not enable_system and message.role is Role.SYSTEM:
            continue

        formatted_lines.append(
            message.format_message(
                index=i if add_index else None,
                add_time=add_time,
                use_name=use_name,
                add_reasoning=add_reasoning,
                add_tools=add_tools,
                strip_markdown_headers=strip_markdown_headers,
            ),
        )
    return "\n".join(formatted_lines)


def merge_messages_content(messages: list[Message | dict]) -> str:
    """Merge messages content into a formatted string representation.

    This function processes a list of messages (either Message objects or dicts)
    and formats them into a structured string. Different message roles are
    formatted differently:
    - ASSISTANT: Includes reasoning content, main content, and tool calls
    - USER: Includes the user content
    - TOOL: Includes tool call results

    Each message is prefixed with a step number (starting from 0) to indicate
    its position in the conversation sequence.

    Args:
        messages: List of Message objects or dictionaries to merge. If a dict
            is provided, it will be converted to a Message object.

    Returns:
        Formatted string representation of all messages with step numbers.
        Each message is separated by newlines and includes role information.

    Example:
        ```python
        messages = [
            Message(role=Role.USER, content="What's the weather?"),
            Message(role=Role.ASSISTANT, content="Let me check",
                    tool_calls=[ToolCall(name="get_weather", arguments={})])
        ]
        result = merge_messages_content(messages)
        # Returns formatted string with step numbers and role information
        ```
    """
    content_collector = []
    for i, message in enumerate(messages):
        if isinstance(message, dict):
            message = Message(**message)

        if message.role is Role.ASSISTANT:
            line = (
                f"### step.{i} role={message.role.value} content=\n{message.reasoning_content}\n\n{message.content}\n"
            )
            if message.tool_calls:
                for tool_call in message.tool_calls:
                    line += f" - tool call={tool_call.name}\n   params={tool_call.arguments}\n"
            content_collector.append(line)

        elif message.role is Role.USER:
            line = f"### step.{i} role={message.role.value} content=\n{message.content}\n"
            content_collector.append(line)

        elif message.role is Role.TOOL:
            line = f"### step.{i} role={message.role.value} tool call result=\n{message.content}\n"
            content_collector.append(line)

    return "\n".join(content_collector)


def parse_json_experience_response(response: str) -> list[dict]:
    """Parse JSON formatted experience response"""
    try:
        # Extract JSON blocks
        json_pattern = r"```json\s*([\s\S]*?)\s*```"
        json_blocks = re.findall(json_pattern, response)

        if json_blocks:
            parsed = json.loads(json_blocks[0])

            # Handle array format
            if isinstance(parsed, list):
                experiences = []
                for exp_data in parsed:
                    if isinstance(exp_data, dict) and (
                        ("when_to_use" in exp_data and "experience" in exp_data)
                        or ("condition" in exp_data and "experience" in exp_data)
                    ):
                        experiences.append(exp_data)

                return experiences

            # Handle single object
            elif isinstance(parsed, dict) and (
                ("when_to_use" in parsed and "experience" in parsed)
                or ("condition" in parsed and "experience" in parsed)
            ):
                return [parsed]

        # Fallback: try to parse entire response
        parsed = json.loads(response)
        if isinstance(parsed, list):
            return parsed
        elif isinstance(parsed, dict):
            return [parsed]

    except json.JSONDecodeError as e:
        logger.warning(f"Failed to parse JSON experience response: {e}")

    return []


def get_trajectory_context(trajectory: Trajectory, step_sequence: list[Message]) -> str:
    """Get context of step sequence within trajectory"""
    try:
        # Find position of step sequence in trajectory
        start_idx = 0
        for i, step in enumerate(trajectory.messages):
            if step == step_sequence[0]:
                start_idx = i
                break

        # Extract before and after context
        context_before = trajectory.messages[max(0, start_idx - 2) : start_idx]
        context_after = trajectory.messages[start_idx + len(step_sequence) : start_idx + len(step_sequence) + 2]

        context = f"Query: {trajectory.metadata.get('query', 'N/A')}\n"

        if context_before:
            context += (
                "Previous steps:\n"
                + "\n".join(
                    [f"- {step.content[:100]}..." for step in context_before],
                )
                + "\n"
            )

        if context_after:
            context += "Following steps:\n" + "\n".join([f"- {step.content[:100]}..." for step in context_after])

        return context

    except Exception as e:
        logger.error(f"Error getting trajectory context: {e}")
        return f"Query: {trajectory.metadata.get('query', 'N/A')}"


def extract_content(text: str, language_tag: str = "json", greedy: bool = False):
    """Extracts content from Markdown code blocks and parses it if the tag is JSON."""
    quantifier = ".*" if greedy else ".*?"
    pattern = rf"```\s*{re.escape(language_tag)}\s*({quantifier})\s*```"
    match = re.search(pattern, text, re.DOTALL)

    if not match:
        return None

    content = match.group(1).strip()

    if language_tag == "json":
        try:
            return json.loads(content)
        except json.JSONDecodeError:
            return None
    else:
        return content


def deduplicate_memories(memories: list[MemoryNode]) -> list[MemoryNode]:
    """Deduplicates a list of memories by memory ID."""
    seen_memories: dict[str, MemoryNode] = {}
    for memory in memories:
        if memory.memory_id not in seen_memories:
            seen_memories[memory.memory_id] = memory
    return list(seen_memories.values())
