from typing import Any, Dict, List, Optional
from pygeai_orchestration.core.base import BasePattern, PatternConfig, PatternResult, PatternType
from pygeai_orchestration.core.common import Message, MessageRole, Conversation, State, StateStatus
import logging

logger = logging.getLogger("pygeai_orchestration")


class AgentRole:
    """
    Represents an agent with a specific role in multi-agent collaboration.

    :param name: str - The unique name/identifier for this agent role.
    :param agent: BaseAgent - The agent instance performing this role.
    :param role_description: str - Description of this agent's responsibilities and expertise.
    """

    def __init__(self, name: str, agent, role_description: str):
        self.name = name
        self.agent = agent
        self.role_description = role_description
        self.contributions: List[str] = []

    def __repr__(self) -> str:
        return f"AgentRole({self.name}: {self.role_description})"


class MultiAgentPattern(BasePattern):
    """
    Multi-agent pattern for collaborative problem-solving with specialized agents.

    This pattern orchestrates multiple agents with different roles and expertise
    to collaboratively solve complex tasks. A coordinator agent manages the workflow,
    distributes work, and synthesizes contributions into a final answer.

    The multi-agent pattern is ideal for:
    - Tasks requiring diverse expertise (e.g., research, analysis, writing)
    - Complex problems benefiting from different perspectives
    - Scenarios where specialized agents outperform generalists
    - Workflow requiring coordination and synthesis of multiple outputs

    The pattern follows this workflow:
    1. **Coordination**: Coordinator agent analyzes task and creates distribution plan
    2. **Parallel/Sequential Execution**: Each specialized agent contributes based on role
    3. **Synthesis**: Coordinator integrates all contributions into coherent final answer

    This approach enables division of labor, specialization, and richer solutions
    than single-agent approaches.
    """

    def __init__(
        self, agents: List[AgentRole], coordinator_agent, config: Optional[PatternConfig] = None
    ):
        """
        Initialize the Multi-Agent pattern.

        :param agents: List[AgentRole] - List of agent roles participating in collaboration.
        :param coordinator_agent: BaseAgent - The coordinator agent managing workflow and synthesis.
        :param config: Optional[PatternConfig] - Pattern configuration. If None, uses default with max_iterations=15.
        """
        if config is None:
            config = PatternConfig(
                name="multi_agent", pattern_type=PatternType.MULTI_AGENT, max_iterations=15
            )
        super().__init__(config)
        self.agent_roles = {role.name: role for role in agents}
        self.coordinator = coordinator_agent
        self._conversation = Conversation(id=f"multi-agent-{id(self)}")
        self._state = State()

    async def execute(self, task: str, context: Optional[Dict[str, Any]] = None) -> PatternResult:
        """
        Execute the multi-agent pattern on the given task.

        The method coordinates multiple specialized agents to collaboratively solve
        the task. The coordinator first creates a distribution plan, then each agent
        contributes based on its role, and finally the coordinator synthesizes all
        contributions into a unified answer.

        :param task: str - The task requiring collaborative multi-agent effort.
        :param context: Optional[Dict[str, Any]] - Additional context for execution. Defaults to None.
        :return: PatternResult - Contains success status, synthesized answer, and agent contributions.
        :raises PatternExecutionError: If the multi-agent pattern execution fails.

        Example:
            >>> roles = [
            ...     AgentRole("researcher", research_agent, "Research and gather facts"),
            ...     AgentRole("analyst", analysis_agent, "Analyze data and identify patterns"),
            ...     AgentRole("writer", writing_agent, "Create clear written summaries")
            ... ]
            >>> pattern = MultiAgentPattern(agents=roles, coordinator_agent=coordinator)
            >>> result = await pattern.execute("Analyze the impact of AI on healthcare")
            >>> print(f"Final answer: {result.result}")
            >>> print(f"Contributions: {result.metadata['contributions']}")
        """
        self.reset()
        for role in self.agent_roles.values():
            role.contributions.clear()
        self._state.update_status(StateStatus.RUNNING)

        logger.info(
            f"Starting multi-agent pattern with {len(self.agent_roles)} agents for task: {task[:50]}..."
        )

        try:
            coordination_plan = await self._coordinate_task(task)
            logger.info(f"Coordination plan: {coordination_plan[:100]}...")

            all_contributions = {}

            for agent_name, role in self.agent_roles.items():
                if self.current_iteration >= self.config.max_iterations:
                    break

                self.increment_iteration()
                logger.debug(f"Agent '{agent_name}' processing task")

                state_data = {
                    "iteration": self.current_iteration,
                    "agent_name": agent_name,
                    "role": role.role_description,
                    "task": task,
                    "coordination_plan": coordination_plan,
                    "previous_contributions": all_contributions,
                }

                step_result = await self.step(state_data)
                contribution = step_result.get("contribution")
                role.contributions.append(contribution)
                all_contributions[agent_name] = contribution

            final_result = await self._synthesize_contributions(task, all_contributions)

            self._state.update_status(StateStatus.COMPLETED)

            return PatternResult(
                success=True,
                result=final_result,
                iterations=self.current_iteration,
                metadata={
                    "agents": list(self.agent_roles.keys()),
                    "contributions": all_contributions,
                    "coordination_plan": coordination_plan,
                },
            )

        except Exception as e:
            logger.error(f"Multi-agent pattern failed: {str(e)}")
            self._state.update_status(StateStatus.FAILED)
            return PatternResult(
                success=False, result=None, iterations=self.current_iteration, error=str(e)
            )

    async def step(self, state: Dict[str, Any]) -> Dict[str, Any]:
        agent_name = state.get("agent_name")
        role_desc = state.get("role")
        task = state.get("task")
        coordination_plan = state.get("coordination_plan", "")
        previous_contributions = state.get("previous_contributions", {})

        role = self.agent_roles[agent_name]

        context_str = ""
        if previous_contributions:
            context_str = "\n\nContributions from other agents:\n"
            for name, contrib in previous_contributions.items():
                context_str += f"- {name}: {contrib[:200]}...\n"

        agent_prompt = (
            f"You are acting as: {role_desc}\n\n"
            f"Task: {task}\n\n"
            f"Coordination plan: {coordination_plan}\n"
            f"{context_str}\n"
            f"Provide your contribution based on your role:"
        )

        contribution = await role.agent.generate(agent_prompt)

        self._conversation.add_message(
            Message(
                role=MessageRole.ASSISTANT,
                content=contribution,
                metadata={"agent": agent_name, "role": role_desc},
            )
        )

        return {"contribution": contribution, "agent": agent_name}

    async def _coordinate_task(self, task: str) -> str:
        agents_desc = "\n".join(
            [f"- {name}: {role.role_description}" for name, role in self.agent_roles.items()]
        )

        coordination_prompt = (
            f"Task: {task}\n\n"
            f"Available agents:\n{agents_desc}\n\n"
            "Create a coordination plan describing how these agents should work together "
            "to accomplish the task. Be specific about what each agent should focus on."
        )

        plan = await self.coordinator.generate(coordination_prompt)
        self._conversation.add_message(
            Message(role=MessageRole.SYSTEM, content=plan, metadata={"type": "coordination_plan"})
        )

        return plan

    async def _synthesize_contributions(self, task: str, contributions: Dict[str, str]) -> str:
        synthesis_prompt = f"Original task: {task}\n\nContributions from each agent:\n"

        for agent_name, contribution in contributions.items():
            role_desc = self.agent_roles[agent_name].role_description
            synthesis_prompt += f"\n{agent_name} ({role_desc}):\n{contribution}\n"

        synthesis_prompt += (
            "\nSynthesize these contributions into a comprehensive final answer. "
            "Integrate insights from all agents and resolve any conflicts or overlaps."
        )

        final_answer = await self.coordinator.generate(synthesis_prompt)

        self._conversation.add_message(
            Message(
                role=MessageRole.ASSISTANT,
                content=final_answer,
                metadata={"type": "final_synthesis"},
            )
        )

        return final_answer

    def get_agent_contributions(self, agent_name: str) -> List[str]:
        role = self.agent_roles.get(agent_name)
        return role.contributions.copy() if role else []

    def get_conversation(self) -> Conversation:
        return self._conversation
