"""
llm_agents.multi_agent
=======================
Multi-agent conversation framework (AIMA §2.4.2 multi-agent environments).

Provides:
  AgentRole       – named agent with a specific persona / system prompt
  MultiAgentChat  – round-robin or orchestrated conversation between agents
  Orchestrator    – LLM supervisor that routes tasks to specialist agents

Usage – debate between two agents
----------------------------------
    from ml_ai_library.llm_agents.multi_agent import AgentRole, MultiAgentChat

    affirmative = AgentRole("Pro", "You argue FOR the motion.", model="gpt-4o-mini")
    negative    = AgentRole("Con", "You argue AGAINST the motion.", model="gpt-4o-mini")

    chat = MultiAgentChat([affirmative, negative], rounds=3)
    transcript = chat.run("Motion: AI will surpass human intelligence by 2030.")
    for speaker, text in transcript:
        print(f"[{speaker}] {text}\n")

Usage – orchestrated specialist routing
----------------------------------------
    from ml_ai_library.llm_agents.multi_agent import AgentRole, Orchestrator

    math_agent   = AgentRole("Mathematician", "Solve maths problems.", ...)
    code_agent   = AgentRole("Programmer",    "Write Python code.", ...)
    orchestrator = Orchestrator(agents=[math_agent, code_agent], model="gpt-4o")
    answer = orchestrator.run("Write a Python function to compute Fibonacci numbers.")
"""

from __future__ import annotations
from typing import List, Tuple

from ml_ai_library.llm_agents.base_agent import BaseAgent, Tool


# ---------------------------------------------------------------------------
# AgentRole
# ---------------------------------------------------------------------------

class AgentRole:
    """
    A named LLM agent with a specific persona.

    Parameters
    ----------
    name    : display name (used as speaker label)
    persona : system prompt defining the agent's role and behaviour
    model   : litellm model string
    tools   : optional list of Tool objects
    """

    def __init__(
        self,
        name: str,
        persona: str,
        model: str = "gpt-4o-mini",
        tools: List[Tool] = None,
        temperature: float = 0.7,
    ) -> None:
        self.name  = name
        self._agent = BaseAgent(
            model=model,
            system=persona,
            tools=tools or [],
            temperature=temperature,
        )

    def respond(self, message: str) -> str:
        return self._agent.chat(message)

    def reset(self):
        self._agent.reset()


# ---------------------------------------------------------------------------
# MultiAgentChat – round-robin conversation
# ---------------------------------------------------------------------------

class MultiAgentChat:
    """
    Round-robin multi-agent conversation.

    Each agent sees the previous agent's message and responds in turn.

    Parameters
    ----------
    agents  : ordered list of AgentRole objects
    rounds  : number of full rounds (each agent speaks once per round)
    """

    def __init__(self, agents: List[AgentRole], rounds: int = 3):
        self.agents = agents
        self.rounds = rounds

    def run(self, topic: str) -> List[Tuple[str, str]]:
        """
        Run the multi-agent conversation.

        Returns
        -------
        transcript : list of (agent_name, utterance) tuples
        """
        transcript: List[Tuple[str, str]] = []
        for agent in self.agents:
            agent.reset()

        # Opening: first agent responds to topic
        message = topic
        for _ in range(self.rounds):
            for agent in self.agents:
                reply = agent.respond(message)
                transcript.append((agent.name, reply))
                message = reply                     # next agent responds to this

        return transcript


# ---------------------------------------------------------------------------
# Orchestrator – LLM supervisor that routes tasks to specialists
# ---------------------------------------------------------------------------

_ORCHESTRATOR_SYSTEM = """\
You are an orchestrator that routes tasks to specialist agents.

Available agents:
{agent_list}

For each user request, decide which agent should handle it.
Return ONLY JSON: {{"agent": "<name>", "task": "<refined task description>"}}
"""


class Orchestrator:
    """
    LLM orchestrator that delegates tasks to specialist agents.

    Parameters
    ----------
    agents  : list of AgentRole objects (the specialist pool)
    model   : litellm model for the orchestrator itself
    """

    def __init__(
        self,
        agents: List[AgentRole],
        model: str = "gpt-4o",
        temperature: float = 0.0,
    ) -> None:
        self.agents  = {a.name: a for a in agents}
        agent_list   = "\n".join(
            f"- {a.name}: {a._agent.history[0]['content'][:120]}..."
            for a in agents
        )
        self._router = BaseAgent(
            model=model,
            system=_ORCHESTRATOR_SYSTEM.format(agent_list=agent_list),
            temperature=temperature,
        )

    def run(self, task: str) -> str:
        """
        Route `task` to the most appropriate specialist agent and return its reply.
        Falls back to the first agent if routing fails.
        """
        import json, re
        self._router.reset()
        routing = self._router.chat(task)
        match   = re.search(r'\{.*\}', routing, re.DOTALL)
        if match:
            try:
                data   = json.loads(match.group())
                name   = data.get("agent", "")
                sub_task = data.get("task", task)
                if name in self.agents:
                    return self.agents[name].respond(sub_task)
            except (json.JSONDecodeError, KeyError):
                pass
        # Fallback: first agent
        first = next(iter(self.agents.values()))
        return first.respond(task)

    def run_parallel(self, tasks: List[str]) -> List[Tuple[str, str]]:
        """
        Route multiple independent tasks in parallel (using threading).

        Returns list of (agent_name, reply) tuples in input order.
        """
        import concurrent.futures
        results = []
        with concurrent.futures.ThreadPoolExecutor() as pool:
            futures = {pool.submit(self.run, t): t for t in tasks}
            for future in concurrent.futures.as_completed(futures):
                results.append((futures[future], future.result()))
        return results
