"""CAMEL Society - Multi-agent collaboration framework."""

from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Callable, Dict, List, Optional, Tuple

from .agent import RolePlayingAgent
from .conversation import Conversation, Message, MessageRole
from .types import Role, Task, TaskSpec, AgentType, DEFAULT_ROLES


@dataclass
class SocietyResult:
    """Result of a CAMEL society collaboration."""
    success: bool
    final_output: str
    conversation: Conversation
    total_turns: int
    execution_time: float = 0.0
    agent_contributions: Dict[str, int] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "success": self.success,
            "final_output": self.final_output,
            "conversation": self.conversation.to_dict(),
            "total_turns": self.total_turns,
            "execution_time": self.execution_time,
            "agent_contributions": self.agent_contributions,
        }

    def to_transcript(self) -> str:
        """Get the conversation transcript."""
        return self.conversation.to_transcript()


class CAMELSociety:
    """Orchestrates multi-agent collaboration through role-playing.

    CAMELSociety implements the core concept from the CAMEL paper:
    autonomous communicative agents that collaborate through
    structured role-playing and inception prompting.

    Example:
        >>> society = CAMELSociety()
        >>> society.add_agent(Role(name="Programmer", description="..."))
        >>> society.add_agent(Role(name="Reviewer", description="..."))
        >>> result = society.collaborate("Write a sorting algorithm")
    """

    def __init__(
        self,
        respond_fn: Optional[Callable[[str, List[Dict], str], str]] = None,
        max_turns: int = 10,
        verbose: bool = False,
    ):
        """Initialize the CAMEL society.

        Args:
            respond_fn: Function to generate responses (context, messages, role_name)
            max_turns: Maximum collaboration turns
            verbose: Whether to print debug info
        """
        self._agents: List[RolePlayingAgent] = []
        self._respond_fn = respond_fn
        self._max_turns = max_turns
        self._verbose = verbose
        self._conversation = Conversation(max_turns=max_turns * 2)

    def add_agent(self, role: Role) -> RolePlayingAgent:
        """Add an agent to the society.

        Args:
            role: The role for the new agent

        Returns:
            The created agent
        """
        agent = RolePlayingAgent(
            role=role,
            respond_fn=self._create_respond_fn(role),
            conversation=self._conversation,
            verbose=self._verbose,
        )
        self._agents.append(agent)
        return agent

    def _create_respond_fn(self, role: Role) -> Callable:
        """Create a respond function for an agent."""
        if self._respond_fn:
            return lambda ctx, msgs: self._respond_fn(ctx, msgs, role.name)
        return None

    def collaborate(
        self,
        task: str,
        agents: Optional[List[Role]] = None,
        termination_fn: Optional[Callable[[str], bool]] = None,
    ) -> SocietyResult:
        """Run a collaborative task between agents.

        Args:
            task: The task description
            agents: Optional list of roles (uses existing if not provided)
            termination_fn: Function to check if task is complete

        Returns:
            SocietyResult with collaboration outcome
        """
        import time
        start_time = time.time()

        # Add provided agents
        if agents:
            for role in agents:
                self.add_agent(role)

        if len(self._agents) < 2:
            raise ValueError("Need at least 2 agents for collaboration")

        # Clear conversation
        self._conversation.clear()

        # Add task as initial context
        task_msg = Message(
            content=f"Task: {task}",
            role=MessageRole.SYSTEM,
            sender_name="System",
        )
        self._conversation.messages.append(task_msg)

        if self._verbose:
            print(f"\n=== Starting Collaboration ===")
            print(f"Task: {task}")
            print(f"Agents: {[a.role.name for a in self._agents]}")

        # Track contributions
        contributions = {agent.role.name: 0 for agent in self._agents}

        # Run collaboration loop
        assistant_idx = 0
        user_idx = 1

        for turn in range(self._max_turns):
            # Assistant speaks
            assistant = self._agents[assistant_idx]
            assistant_msg = assistant.respond(
                f"Turn {turn + 1}: Please contribute to the task."
            )
            contributions[assistant.role.name] += 1

            # Check for termination
            if termination_fn and termination_fn(assistant_msg.content):
                if self._verbose:
                    print(f"\nTask completed by {assistant.role.name}")
                break

            # User responds
            user = self._agents[user_idx]
            user_msg = user.respond(assistant_msg.content)
            contributions[user.role.name] += 1

            # Check for termination
            if termination_fn and termination_fn(user_msg.content):
                if self._verbose:
                    print(f"\nTask completed by {user.role.name}")
                break

        # Get final output
        last_msg = self._conversation.messages[-1] if self._conversation.messages else None
        final_output = last_msg.content if last_msg else "No output generated"

        execution_time = time.time() - start_time

        if self._verbose:
            print(f"\n=== Collaboration Complete ===")
            print(f"Turns: {self._conversation.current_turn}")
            print(f"Time: {execution_time:.2f}s")

        return SocietyResult(
            success=True,
            final_output=final_output,
            conversation=self._conversation,
            total_turns=self._conversation.current_turn,
            execution_time=execution_time,
            agent_contributions=contributions,
        )

    def quick_collaborate(
        self,
        task: str,
        role1: str = "programmer",
        role2: str = "reviewer",
    ) -> SocietyResult:
        """Quick collaboration with predefined roles.

        Args:
            task: The task description
            role1: First role name (from DEFAULT_ROLES)
            role2: Second role name (from DEFAULT_ROLES)

        Returns:
            SocietyResult
        """
        r1 = DEFAULT_ROLES.get(role1, Role(name=role1.title(), description=f"A {role1}"))
        r2 = DEFAULT_ROLES.get(role2, Role(name=role2.title(), description=f"A {role2}"))

        # Clear existing agents
        self._agents.clear()

        return self.collaborate(task, agents=[r1, r2])

    def get_conversation(self) -> Conversation:
        """Get the current conversation.

        Returns:
            The conversation object
        """
        return self._conversation

    def get_agents(self) -> List[RolePlayingAgent]:
        """Get all agents in the society.

        Returns:
            List of agents
        """
        return self._agents.copy()

    def to_dict(self) -> Dict[str, Any]:
        return {
            "agents": [a.to_dict() for a in self._agents],
            "max_turns": self._max_turns,
            "conversation": self._conversation.to_dict(),
        }