"""
orchestrated workspace - true turn-based multi-agent conversation.

this module provides turn-based orchestration where agents take turns
in a controlled conversation flow with full history visibility.
"""

import asyncio
import json
import logging
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Optional

from synqed.orchestrator import Orchestrator, OrchestrationResult
from synqed.workspace import Workspace, WorkspaceState, MessageType

logger = logging.getLogger(__name__)


class CollaborationMode(Enum):
    """defines how agents interact during collaboration."""
    COLLABORATIVE = "collaborative"  # supportive, synergy-seeking
    ADVERSARIAL = "adversarial"      # force conflict, critique, then synthesis
    INDEPENDENT = "independent"       # no interaction between agents


@dataclass
class AgentRole:
    """an agent role in the conversation."""
    role_id: str
    description: str
    agent_id: str
    agent_name: str
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class ConversationPlan:
    """a plan for a turn-based conversation."""
    plan_id: str
    original_task: str
    agent_roles: list[AgentRole]
    num_turns: int
    selected_agents: list[str]  # agent ids needed
    reasoning: str
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class ExecutionResult:
    """result of executing a turn-based conversation."""
    plan: ConversationPlan
    workspace_id: str
    final_result: str
    success: bool
    workspace_messages: list[dict[str, Any]] = field(default_factory=list)
    error: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)


class OrchestratedWorkspace:
    """
    turn-based orchestration for multi-agent conversations.
    
    this orchestrator:
    1. analyzes the user's task
    2. selects which agents to use
    3. creates a temporary workspace
    4. runs a turn-based conversation between agents
    5. composes final result from conversation history
    
    example:
        ```python
        # create orchestrated workspace
        orchestrated = OrchestratedWorkspace(
            orchestrator=orchestrator,
            enable_agent_discussion=True,
            collaboration_mode=CollaborationMode.COLLABORATIVE,
            display_messages=True,
        )
        
        # register agents
        orchestrated.register_agent(writer_agent)
        orchestrated.register_agent(editor_agent)
        
        # execute task with turn-based conversation
        result = await orchestrated.execute_task(
            "write a short story about a robot learning to paint"
        )
        
        print(f"success: {result.success}")
        print(f"result: {result.final_result}")
        ```
    """
    
    def __init__(
        self,
        orchestrator: Orchestrator,
        enable_agent_discussion: bool = True,
        collaboration_mode: CollaborationMode = CollaborationMode.ADVERSARIAL,
        workspace_persistence: bool = False,
        num_turns: int = 6,
        message_callback: Optional[Any] = None,
        display_messages: bool = False,
        display_verbose: bool = False,
    ):
        """
        initialize the orchestrated workspace.
        
        args:
            orchestrator: orchestrator instance for agent selection
            enable_agent_discussion: whether to enable turn-based conversation
            collaboration_mode: how agents interact (adversarial, collaborative, or independent)
            workspace_persistence: whether to persist workspace state
            num_turns: number of conversation turns (default: 6)
            message_callback: optional custom callback function(msg_dict) for real-time messages
            display_messages: enable built-in message display (default: false)
            display_verbose: show full message details (default: false)
        """
        self.orchestrator = orchestrator
        self.enable_agent_discussion = enable_agent_discussion
        self.collaboration_mode = collaboration_mode
        self.workspace_persistence = workspace_persistence
        self.num_turns = num_turns
        self.display_verbose = display_verbose
        
        # setup message callback
        if display_messages:
            # create a simple display callback
            def display_callback(msg_dict):
                # print with simple output
                sender = msg_dict.get('sender_name', 'system')
                content = msg_dict.get('content', '')
                print(f"[{sender}] {content}")
                
                # call custom callback if provided
                if message_callback:
                    message_callback(msg_dict)
            
            self.message_callback = display_callback
        else:
            self.message_callback = message_callback
        
        # track agents
        self._agents: dict[str, Any] = {}  # agent_id -> agent
        
        logger.info("orchestratedworkspace initialized")
    
    def register_agent(self, agent: Any) -> None:
        """
        Register an agent for orchestration.
        
        Args:
            agent: Agent instance to register
        """
        agent_id = agent.name
        self._agents[agent_id] = agent
        
        # Register with orchestrator
        self.orchestrator.register_agent(
            agent_card=agent.card,
            agent_url=agent.url,
            agent_id=agent_id,
        )
        
        logger.info(f"Registered agent: {agent_id}")
    
    async def execute_task(
        self,
        task: str,
        context: dict[str, Any] | None = None,
        workspace_name: str | None = None,
    ) -> ExecutionResult:
        """
        execute a task with turn-based conversation between agents.
        
        args:
            task: the task to execute
            context: optional additional context
            workspace_name: optional custom workspace name
            
        returns:
            executionresult with conversation plan, results, and status
        """
        logger.info(f"starting turn-based execution: {task[:100]}...")
        
        try:
            # step 1: create conversation plan (select agents)
            plan = await self._create_conversation_plan(task, context)
            logger.info(f"created plan with {len(plan.agent_roles)} agents for {plan.num_turns} turns")
            
            # step 2: create temporary workspace
            workspace = await self._create_workspace(
                plan=plan,
                workspace_name=workspace_name or f"conversation: {task[:50]}",
            )
            logger.info(f"created workspace: {workspace.workspace_id}")
            
            # step 3: execute turn-based conversation
            final_result = await self._execute_conversation(plan, workspace)
            
            # step 4: complete workspace
            await workspace.complete()
            
            # capture workspace messages before closing
            workspace_messages = [msg.to_dict() for msg in workspace.get_messages()]
            
            # export if persistence enabled
            if self.workspace_persistence:
                export_path = await workspace.export_workspace()
                logger.info(f"workspace exported: {export_path}")
            
            # close workspace
            await workspace.close()
            
            # return result
            return ExecutionResult(
                plan=plan,
                workspace_id=workspace.workspace_id,
                final_result=final_result,
                success=True,
                workspace_messages=workspace_messages,
                error=None,
                metadata={
                    "num_agents": len(plan.selected_agents),
                    "num_turns": plan.num_turns,
                },
            )
        
        except Exception as e:
            logger.error(f"turn-based execution failed: {e}", exc_info=True)
            
            # try to capture messages if workspace exists
            workspace_messages = []
            if 'workspace' in locals():
                try:
                    workspace_messages = [msg.to_dict() for msg in workspace.get_messages()]
                except:
                    pass
            
            return ExecutionResult(
                plan=plan if 'plan' in locals() else None,
                workspace_id=workspace.workspace_id if 'workspace' in locals() else None,
                final_result="",
                success=False,
                workspace_messages=workspace_messages,
                error=str(e),
            )
    
    async def _create_conversation_plan(
        self,
        task: str,
        context: dict[str, Any] | None,
    ) -> ConversationPlan:
        """
        create a conversation plan by selecting which agents to use.
        
        args:
            task: the task to execute
            context: optional context
            
        returns:
            conversationplan with agent roles and turn count
        """
        # for simple demos, just use all registered agents
        agent_roles = []
        for agent_id, agent in self._agents.items():
            agent_roles.append(AgentRole(
                role_id=f"role_{agent_id}",
                description=agent.card.description if hasattr(agent, 'card') else agent.description,
                agent_id=agent_id,
                agent_name=agent.name,
            ))
        
        # create simple conversation plan
        plan = ConversationPlan(
            plan_id=f"conversation_{task[:20]}",
            original_task=task,
            agent_roles=agent_roles,
            num_turns=self.num_turns,
            selected_agents=[r.agent_id for r in agent_roles],
            reasoning="turn-based conversation between registered agents",
            metadata={},
        )
        
        return plan
    
    
    async def _create_workspace(
        self,
        plan: ConversationPlan,
        workspace_name: str,
    ) -> Workspace:
        """
        create a workspace and add the required agents.
        
        args:
            plan: conversation plan
            workspace_name: name for the workspace
            
        returns:
            configured workspace instance
        """
        # create workspace
        workspace = Workspace(
            name=workspace_name,
            description=f"turn-based conversation: {plan.original_task[:100]}",
            enable_persistence=self.workspace_persistence,
            auto_cleanup=not self.workspace_persistence,
            display_messages=False,  # disable workspace's display - orchestrated workspace handles it
        )
        
        # set shared state with plan information
        workspace.set_shared_state("conversation_plan", {
            "original_task": plan.original_task,
            "agent_roles": [
                {
                    "id": role.role_id,
                    "description": role.description,
                    "agent_name": role.agent_name,
                }
                for role in plan.agent_roles
            ],
            "num_turns": plan.num_turns,
        })
        
        # add required agents to workspace
        for agent_id in plan.selected_agents:
            if agent_id in self._agents:
                agent = self._agents[agent_id]
                workspace.add_agent(
                    agent=agent,
                    metadata={},
                )
                logger.info(f"added agent to workspace: {agent_id}")
        
        # register message callback if provided
        if self.message_callback is not None:
            async def message_handler(msg):
                """handle workspace messages in real-time."""
                try:
                    self.message_callback(msg.to_dict())
                except Exception as e:
                    logger.error(f"message callback error: {e}")
            
            workspace.on_message(message_handler)
        
        # start workspace
        await workspace.start()
        
        return workspace
    
    async def _execute_conversation(
        self,
        plan: ConversationPlan,
        workspace: Workspace,
    ) -> str:
        """
        execute turn-based conversation between agents.
        
        args:
            plan: conversation plan
            workspace: workspace for execution
            
        returns:
            final result string composed from conversation
        """
        # get participants
        participants = {p["name"]: p["participant_id"] for p in workspace.list_participants()}
        
        if self.enable_agent_discussion and len(plan.agent_roles) > 1:
            # execute turn-based dialogue
            final_result = await self._execute_turn_taking_dialogue(
                agent_roles=plan.agent_roles,
                workspace=workspace,
                participants=participants,
                num_turns=plan.num_turns,
            )
        else:
            # single agent or no discussion - just ask once
            if len(plan.agent_roles) > 0:
                role = plan.agent_roles[0]
                participant_id = participants.get(role.agent_name)
                if participant_id:
                    prompt = f"task: {plan.original_task}\n\nyour role: {role.description}\n\nplease complete this task."
                    final_result = await workspace.send_message_to_agent(participant_id, prompt)
                else:
                    final_result = "no agents available for execution."
            else:
                final_result = "no agents registered."
        
        return final_result
    
    
    async def _execute_turn_taking_dialogue(
        self,
        agent_roles: list[AgentRole],
        workspace: Workspace,
        participants: dict[str, str],
        num_turns: int = 6,
    ) -> str:
        """
        execute true turn-based dialogue with deterministic turn order.
        
        this implements the fixed architecture where:
        1. agents take turns in a fixed sequence (no rounds, just turns)
        2. only one agent responds per turn
        3. each agent sees full conversation history from all previous turns
        4. no broadcasts that trigger both agents
        5. clear logging shows who speaks when
        
        for story writing demo, turn sequence is:
        - turn 1: writer proposes character and concept
        - turn 2: editor responds to writer's idea
        - turn 3: writer writes story draft based on editor feedback
        - turn 4: editor critiques and revises
        - turn 5: writer final revision
        - turn 6: editor final polish
        
        args:
            agent_roles: agent roles participating in conversation
            workspace: workspace with conversation history
            participants: name to participant_id mapping
            num_turns: number of turns in the conversation (default: 6)
            
        returns:
            final result string composed from conversation
        """
        if len(agent_roles) < 2:
            # not enough agents for dialogue
            return "insufficient agents for conversation."
        
        # determine turn order - just alternate between agents
        turn_order = [role for role in agent_roles]
        
        # build turn announcement and store as workspace context (not sent to agents)
        turn_announcement = self._build_turn_order_announcement(turn_order, workspace)
        
        # store announcement as system message in conversation history
        workspace.append_conversation_message(
            sender_id="orchestrator",
            sender_name="orchestrator",
            receiver_id="all",
            receiver_name="all",
            role="system",
            content=turn_announcement,
            message_type=MessageType.SYSTEM,
        )
        
        logger.info(f"[orchestrator] turn order: {' → '.join([role.agent_name for role in turn_order])}")
        logger.info(f"[orchestrator] {turn_order[0].agent_name} will speak first")
        logger.info(f"[orchestrator] {num_turns} total turns planned\n")
        
        # turn-based dialogue loop - fixed number of turns
        for turn_num in range(num_turns):
            logger.info(f"{'='*60}")
            logger.info(f"[orchestrator] turn {turn_num + 1}/{num_turns}")
            logger.info(f"{'='*60}")
            
            # determine which agent speaks this turn (alternate between agents)
            current_role = turn_order[turn_num % len(turn_order)]
            current_agent_name = current_role.agent_name
            participant_id = participants.get(current_agent_name)
            
            if not participant_id:
                logger.warning(f"agent not found: {current_agent_name}")
                continue
            
            logger.info(f"[orchestrator] → sending message to {current_agent_name}")
            
            # build prompt from conversation history
            prompt = self._build_turn_prompt_from_history(
                current_role=current_role,
                turn_num=turn_num,
                max_turns=num_turns,
                workspace=workspace,
            )
            
            # send directed message only to this agent
            try:
                logger.info(f"[{current_agent_name}] processing...")
                response = await workspace.send_message_to_agent(participant_id, prompt)
                
                # add agent's response to conversation history
                workspace.append_conversation_message(
                    sender_id=participant_id,
                    sender_name=current_agent_name,
                    receiver_id="all",
                    receiver_name="all",
                    role="assistant",
                    content=response,
                    message_type=MessageType.AGENT,
                )
                
                logger.info(f"[{current_agent_name}] responded (turn {turn_num + 1} complete)\n")
                
            except Exception as e:
                logger.error(f"turn {turn_num + 1} failed for {current_agent_name}: {e}")
        
        logger.info(f"{'='*60}")
        logger.info(f"[orchestrator] all {num_turns} turns completed")
        logger.info(f"{'='*60}\n")
        
        # compose final result from conversation history
        logger.info("[orchestrator] composing final result from conversation history\n")
        final_result = await self._compose_final_result_from_history(
            agent_roles=agent_roles,
            workspace=workspace,
        )
        
        return final_result
    
    def _build_turn_order_announcement(
        self,
        turn_order: list[AgentRole],
        workspace: Workspace,
    ) -> str:
        """build turn order announcement (stored as context, not sent to agents)."""
        plan_info = workspace.get_shared_state("conversation_plan", {})
        
        if self.collaboration_mode == CollaborationMode.ADVERSARIAL:
            parts = [
                "⚔️ adversarial dialogue: turn-taking mode",
                "",
                f"complex task: {plan_info.get('original_task', '')}",
                "",
                "turn order (agents will alternate):",
            ]
            
            for i, role in enumerate(turn_order, 1):
                parts.append(f"  {i}. {role.agent_name}: {role.description}")
            
            parts.extend([
                "",
                "protocol:",
                f"• {turn_order[0].agent_name} speaks first",
                "• each agent responds directly to the previous agent's output",
                "• agents iterate back and forth until satisfied",
                "• agents must challenge weak thinking and maximize quality",
            ])
        else:
            parts = [
                "🤝 collaborative dialogue: turn-taking mode",
                "",
                f"working together on: {plan_info.get('original_task', '')}",
                "",
                "turn order:",
            ]
            
            for i, role in enumerate(turn_order, 1):
                parts.append(f"  {i}. {role.agent_name}: {role.description}")
            
            parts.extend([
                "",
                f"• {turn_order[0].agent_name} will share their thinking first",
                "• then agents will take turns building on each other's ideas",
            ])
        
        return "\n".join(parts)
    
    def _build_turn_prompt_from_history(
        self,
        current_role: AgentRole,
        turn_num: int,
        max_turns: int,
        workspace: Workspace,
    ) -> str:
        """
        build prompt for current turn based on conversation history.
        
        this builds a prompt that includes:
        1. the overall task and agent's role
        2. full conversation history with sender labels so agents can see what others said
        3. the current turn instruction
        
        args:
            current_role: the agent role for current agent
            turn_num: current turn number (0-indexed)
            max_turns: maximum turns allowed
            workspace: workspace with conversation history
            
        returns:
            prompt string
        """
        plan_info = workspace.get_shared_state("conversation_plan", {})
        current_agent_name = current_role.agent_name
        
        # get recent conversation history (last 15 messages to provide good context)
        recent_history = workspace.get_recent_conversation_history(15)
        
        # build conversation context from history
        conversation_context = []
        if recent_history:
            conversation_context.append("conversation so far:")
            conversation_context.append("")
            for msg in recent_history:
                # show all agent messages (skip system messages except important ones)
                if msg["role"] in ["assistant"]:
                    sender = msg["sender_name"]
                    content = msg["content"]
                    conversation_context.append(f"[{sender}]")
                    conversation_context.append(content)
                    conversation_context.append("")
        
        is_first_turn = turn_num == 0
        is_late_turn = turn_num >= max_turns - 2
        
        # define what this turn should accomplish based on turn number
        turn_instructions = {
            0: "propose a character concept and story idea",
            1: "respond to the writer's proposal with feedback and suggestions",
            2: "write a story draft incorporating the editor's feedback",
            3: "critique the draft and provide revision suggestions",
            4: "revise the story based on the editor's critique",
            5: "provide final polish and refinement",
        }
        
        current_instruction = turn_instructions.get(turn_num, "continue the conversation")
        
        if self.collaboration_mode == CollaborationMode.ADVERSARIAL:
            if is_first_turn:
                # first agent's opening move
                parts = [
                    f"you are {current_agent_name}",
                    "",
                    f"task: {plan_info.get('original_task', '')}",
                    f"your role: {current_role.description}",
                    "",
                    f"turn {turn_num + 1}/{max_turns}",
                    "",
                    "you are going first.",
                    f"your job: {current_instruction}",
                    "",
                    "requirements:",
                    "• be clear and specific",
                    "• take a position based on your expertise",
                    "• the editor will respond to you next",
                ]
            else:
                # subsequent turns - respond to previous agent
                parts = [
                    f"you are {current_agent_name}",
                    "",
                    f"task: {plan_info.get('original_task', '')}",
                    f"your role: {current_role.description}",
                    "",
                ]
                
                if conversation_context:
                    parts.extend(conversation_context)
                
                parts.extend([
                    f"turn {turn_num + 1}/{max_turns}",
                    f"your job: {current_instruction}",
                    "",
                ])
                
                if is_late_turn:
                    parts.extend([
                        "this is one of the final turns.",
                        "",
                        "requirements:",
                        "• directly respond to what the other agent just said",
                        "• challenge weak points or suggest improvements",
                        "• produce high-quality work",
                        "• if this is your final output, make it complete",
                    ])
                else:
                    parts.extend([
                        "requirements:",
                        "• directly engage with what the other agent said",
                        "• build on or critique their work",
                        "• be specific and constructive",
                    ])
        else:
            # collaborative mode
            if is_first_turn:
                parts = [
                    f"you are {current_agent_name}",
                    "",
                    f"task: {plan_info.get('original_task', '')}",
                    f"your role: {current_role.description}",
                    "",
                    f"turn {turn_num + 1}/{max_turns}",
                    "",
                    "you're going first!",
                    f"your job: {current_instruction}",
                    "",
                    "requirements:",
                    "• be clear and specific",
                    "• share your creative ideas",
                    "• the editor will respond next",
                ]
            else:
                parts = [
                    f"you are {current_agent_name}",
                    "",
                    f"task: {plan_info.get('original_task', '')}",
                    f"your role: {current_role.description}",
                    "",
                ]
                
                if conversation_context:
                    parts.extend(conversation_context)
                
                parts.extend([
                    f"turn {turn_num + 1}/{max_turns}",
                    f"your job: {current_instruction}",
                    "",
                ])
                
                if is_late_turn:
                    parts.extend([
                        "we're approaching the end.",
                        "",
                        "requirements:",
                        "• respond to what the other agent just said",
                        "• incorporate their suggestions",
                        "• produce polished, complete work",
                    ])
                else:
                    parts.extend([
                        "requirements:",
                        "• respond to what the other agent shared",
                        "• build on their work constructively",
                        "• be helpful and specific",
                    ])
        
        return "\n".join(parts)
    
    async def _compose_final_result_from_history(
        self,
        agent_roles: list[AgentRole],
        workspace: Workspace,
    ) -> str:
        """
        compose final result from conversation history without additional agent calls.
        
        this extracts the final story and character details from the conversation
        and formats them into a complete result.
        
        args:
            agent_roles: agent roles participating in conversation
            workspace: workspace with conversation history
            
        returns:
            composed final result string
        """
        # get full conversation history from workspace
        conversation_history = workspace.get_all_conversation_history()
        
        # extract agent messages (skip system messages)
        agent_messages = [
            msg for msg in conversation_history
            if msg.get("role") in ["assistant", "user"] and msg.get("sender_name") != "orchestrator"
        ]
        
        if not agent_messages:
            return "no conversation history found."
        
        # get the last few messages which should contain the final story
        last_messages = agent_messages[-4:] if len(agent_messages) >= 4 else agent_messages
        
        # build final result
        result_parts = [
            "final story result",
            "=" * 60,
            "",
        ]
        
        # find story content from the last writer message
        writer_name = agent_roles[0].agent_name if agent_roles else "creative writer"
        editor_name = agent_roles[1].agent_name if len(agent_roles) > 1 else "professional editor"
        
        story_content = None
        character_info = None
        
        # look for story in recent messages
        for msg in reversed(agent_messages):
            sender = msg.get("sender_name", "").lower()
            content = msg.get("content", "")
            
            if "writer" in sender and not story_content:
                story_content = content
            
            if not character_info and len(content) < 300:
                # early messages are likely character concepts
                character_info = content
        
        # add character info if found
        if character_info:
            result_parts.extend([
                "character concept:",
                character_info,
                "",
            ])
        
        # add story
        if story_content:
            result_parts.extend([
                "final story:",
                story_content,
                "",
            ])
        else:
            # fallback to last message
            result_parts.extend([
                "final output:",
                agent_messages[-1].get("content", ""),
                "",
            ])
        
        result_parts.extend([
            "=" * 60,
            f"completed through {len(agent_messages)} turns of dialogue between {writer_name} and {editor_name}",
        ])
        
        return "\n".join(result_parts)
    
    def __repr__(self) -> str:
        """string representation."""
        return (
            f"OrchestratedWorkspace(agents={len(self._agents)}, "
            f"mode={self.collaboration_mode.value}, "
            f"discussion_enabled={self.enable_agent_discussion})"
        )
