"""Module for managing a conversation between AI agents."""

import json
from collections.abc import Iterator
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, TypedDict, cast

from partial_json_parser import ensure_json  # type: ignore[import-untyped] # pyright: ignore[reportMissingTypeStubs]
from pydantic import BaseModel, Field, create_model

from .ai_agent import AIAgent


@dataclass
class ConversationManager:
    """Manager for a conversation between AI agents."""

    class _ConversationLogItem(TypedDict):
        agent: str
        content: str

    agents: list[AIAgent]
    initial_message: str | None
    use_markdown: bool = False
    allow_termination: bool = False
    _conversation_log: list[_ConversationLogItem] = field(default_factory=list, init=False)
    _original_system_prompts: list[str] = field(init=False)
    _output_format: type[BaseModel] = field(init=False)

    def __post_init__(self) -> None:  # noqa: D105
        self._original_system_prompts = [agent.system_prompt for agent in self.agents]

        output_format_kwargs: dict[str, Any] = {
            "message": (str, Field(description="Message content")),
        }

        if self.allow_termination:
            output_format_kwargs["terminate"] = (
                bool,
                Field(
                    title="Terminate",
                    description="Terminate conversation if you believe it has reached a natural conclusion. "
                    + "Do not set this field to `true` unless you are certain the conversation should end.",
                ),
            )

        self._output_format = create_model("OutputFormat", **output_format_kwargs)

        # Modify system prompt to include termination instructions if allowed
        additional_instructions: str = ""

        if self.use_markdown:
            additional_instructions += (
                "You may use Markdown for text formatting. "
                "Examples: *italic*, **bold**, `code`, [link](https://example.com), etc.\n\n"
            )

        # Updated system prompts for each agent. This gives the agents more context about the conversation and their
        # role. It also makes the agents use dialogue markers in their responses, which helps other AIs better interpret
        # the conversation.
        #
        # This dialogue marker is stripped from the actual output stream generated by run_conversation(), but other AI
        # agents can still see it in the conversation history, and it helps them generate more coherent responses.
        for agent in self.agents:
            other_agents = ", ".join([a.name for a in self.agents if a != agent])
            agent.system_prompt = (
                "This is a conversation between AI agents.\n\n"
                + f"You are named {agent.name}. The other agents are {other_agents}. "
                + "Your task is to play the role you're given and continue the conversation.\n\n"
                + f"Always start your message with your name followed by a colon (e.g., '{agent.name}: Hello')\n\n"
                + f"This is the prompt for your role: {agent.system_prompt}\n\n"
                + additional_instructions
            ).rstrip()

    def save_conversation(self, filename: Path) -> None:
        """Save the conversation log to a file.

        Args:
            filename (Path): Path to save the conversation log to
        """
        with open(filename, "w", encoding="utf-8") as f:
            if filename.suffix == ".json":

                def agent_to_dict(agent_idx: int) -> dict[str, Any]:
                    agent = self.agents[agent_idx]

                    return {
                        "name": agent.name,
                        "model": agent.model,
                        "temperature": agent.temperature,
                        "ctx_size": agent.ctx_size,
                        "system_prompt": self._original_system_prompts[agent_idx],
                    }

                # Save conversation log as JSON
                json.dump(
                    {
                        "agents": [agent_to_dict(i) for i in range(len(self.agents))],
                        "conversation": self._conversation_log,
                    },
                    f,
                    ensure_ascii=False,
                    indent=4,
                )
                return

            # Save conversation log as plain text
            for i, agent in enumerate(self.agents, start=1):
                _ = f.write(f"=== Agent {i} ===\n\n")
                _ = f.write(f"Name: {agent.name}\n")
                _ = f.write(f"Model: {agent.model}\n")
                _ = f.write(f"Temperature: {agent.temperature}\n")
                _ = f.write(f"Context Size: {agent.ctx_size}\n")
                _ = f.write(f"System Prompt: {self._original_system_prompts[i - 1]}\n\n")

            _ = f.write("=== Conversation ===\n\n")

            for i, msg in enumerate(self._conversation_log):
                if i > 0:
                    _ = f.write("\n" + "\u2500" * 80 + "\n\n")

                _ = f.write(f"{msg['agent']}: {msg['content']}\n")

    def run_conversation(self) -> Iterator[tuple[str, Iterator[str]]]:
        """Generate an iterator of conversation responses.

        Yields:
            (str, Iterator[str]): A tuple of the agent name and an iterator of the accumulated response.

                                  Note: The iterator returns the entire message up until the newest chunk received,
                                  not just the new chunk.

                                  For example, if the first iteration yields "Hello, ", the second iteration will yield
                                  "Hello, world!" instead of just "world!".
        """
        turn_count = 0

        def add_agent_response(agent_idx: int, response: dict[str, Any]) -> None:
            """Add a message from an agent to the conversation log and the agents' message history.

            Args:
                agent_idx (int): Index of the agent the message is from
                message (str): Message content
            """
            # The agents should get the full JSON response as context, to reinforce the response format and help them
            # generate more coherent responses.
            for i, agent in enumerate(self.agents):
                agent.add_message(
                    self.agents[agent_idx].name,
                    # Use "assistant" instead of "user" for the agent's own messages.
                    "assistant" if i == agent_idx else "user",
                    str(response),
                )

            # For the conversation log, only the message content is needed.
            # Strip the dialogue marker from the message before adding it to the conversation log.
            message: str = response["message"][len(self.agents[agent_idx].name) + 2 :]
            self._conversation_log.append({"agent": self.agents[agent_idx].name, "content": message})

        # If a non-empty initial message is provided, start with it.
        if self.initial_message is not None:
            # Make the first agent the one to say the initial message, and the second agent the one to respond.
            add_agent_response(0, {"message": self.initial_message})
            yield (self.agents[0].name, iter([self.initial_message]))
            turn_count += 1

        while True:
            agent_idx = turn_count % len(self.agents)
            current_agent = self.agents[agent_idx]
            response_stream = current_agent.get_response(self._output_format)

            # Will be populated with the full JSON response once the response stream is exhausted.
            response_json: dict[str, Any] = {}

            def parse_partial_json(json_string: str) -> dict[str, Any]:
                """Parse a partial JSON response using the partial JSON parser, and return the JSON object."""
                # Don't use `partial_json_parser.loads()` directly because it doesn't have good type hints.
                return cast(dict[str, Any], json.loads(ensure_json(json_string)))

            def stream_chunks() -> Iterator[str]:
                nonlocal response_json

                response: str = ""

                # Accumulate chunks until the message field is found in the JSON response.
                for response_chunk in response_stream:
                    response += response_chunk
                    response_json = parse_partial_json(response)

                    if "message" in response_json:
                        break

                # Message field is found, yield the entire message gradually as new chunks arrive.
                for response_chunk in response_stream:
                    response += response_chunk
                    response_json = parse_partial_json(response)
                    message: str = response_json["message"]

                    # Yield the message with the dialogue marker stripped.
                    yield message[len(current_agent.name) + 2 :]

            yield (current_agent.name, stream_chunks())

            add_agent_response(agent_idx, response_json)

            # Check if the conversation should be terminated.
            if response_json.get("terminate", False):
                break

            turn_count += 1
