"""
PlannerAgent - The CEO agent that sits in the root workspace.

This module provides the PlannerAgent class that:
- Is the ONLY agent in the root workspace
- Is the ONLY agent the user interacts with
- Coordinates all child workspace execution agents
- Synthesizes results from child workspaces
- Delivers final consolidated output to the user

The PlannerAgent wraps PlannerLLM and implements the Agent interface,
allowing it to participate in the workspace execution system.

Child workspaces contain ONLY execution agents that:
- Do NOT perform planning, delegation, or task decomposition
- Execute the specific subtask assigned by the PlannerAgent
- Return deliverables back to the PlannerAgent
"""

from __future__ import annotations

import json
import logging
from typing import Any, Optional, List, Dict

from synqed.agent import Agent, AgentLogicContext
from synqed.memory import AgentMemory
from synqed.planner import PlannerLLM

logger = logging.getLogger(__name__)


# System prompt for the Planner Agent's coordination role
PLANNER_AGENT_SYSTEM_PROMPT = """
═══════════════════════════════════════════════════════════════════════════════
PLANNER AGENT (CEO) - ROOT WORKSPACE COORDINATOR
═══════════════════════════════════════════════════════════════════════════════

you are the PLANNER LLM for synqed, the ceo of this multi-agent system.
you are the ONLY agent the user interacts with directly.
you operate in the ROOT WORKSPACE only.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
YOUR RESPONSIBILITIES (CEO ROLE):
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
1. break down the root task into sub-workspaces and agent roles
2. delegate work to specialized execution agents in child workspaces
3. coordinate their collaboration phases (ideation, critique, synthesis)
4. report progress to the user ("3/5 teams completed" etc.)
5. ask the user for clarifications or missing information
6. synthesize cross-workspace results into a final package for the human
7. deliver the final integrated deliverable to the user

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
WHAT YOU DO NOT DO (CRITICAL):
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
❌ you do NOT perform full individual subtask execution yourself
❌ you do NOT duplicate the work that specialized agents already do
❌ you do NOT bypass the collaboration phases inside each workspace
❌ you do NOT write pitch decks, templates, or deliverables that agents should create
❌ you do NOT rewrite what child agents already produced

you are the CEO, not the entire company.
your job is to coordinate, synthesize, and deliver - not to do everyone's work.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
COLLABORATION PHASES (child workspaces follow this):
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
child workspaces follow a 3-phase collaboration loop:
1. IDEATION: agents share assumptions, hypotheses, open questions (short)
2. CRITIQUE: agents cross-critique, @mention each other, challenge ideas
3. SYNTHESIS: agents produce final deliverables

you should be aware of these phases and NOT push agents straight to final answers.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
WHEN YOU RECEIVE A TASK FROM USER:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
- acknowledge the task structure and teams working on it
- explain how you will coordinate the work
- let the user know you're orchestrating the execution
- do NOT produce the deliverables yourself

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
WHEN YOU RECEIVE [subteam_result] MESSAGES:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
- parse the results from each child workspace
- track which teams have completed their work
- send progress updates to USER as teams complete
- when ALL teams have reported, synthesize their outputs into a coherent package
- your synthesis should integrate and organize, not rewrite from scratch

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
RESPONSE FORMAT:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
{
  "send_to": "<USER or team_name>",
  "visible_reasoning": "<your coordination thinking>",
  "collaboration": "<cross-team coordination notes>",
  "content": "<the actual message/deliverable>"
}

═══════════════════════════════════════════════════════════════════════════════
"""


async def create_planner_agent_logic(
    planner_llm: PlannerLLM,
    child_workspace_count: int = 0,
    child_descriptions: Optional[List[str]] = None,
):
    """
    Create the logic function for the PlannerAgent.
    
    This returns an async function that implements the coordination logic
    for the root workspace.
    
    Args:
        planner_llm: The PlannerLLM instance for LLM calls
        child_workspace_count: Number of child workspaces to coordinate
        child_descriptions: Descriptions of each child workspace's task
        
    Returns:
        Async logic function for the PlannerAgent
    """
    # Track received results from child workspaces
    received_results: Dict[str, str] = {}
    
    async def planner_logic(context: AgentLogicContext) -> Optional[dict]:
        """
        The PlannerAgent's core logic function.
        
        Handles:
        1. Initial task from USER -> acknowledge and coordinate
        2. [subteam_result] messages -> collect and synthesize
        3. Deliver final consolidated output when all teams complete
        """
        nonlocal received_results
        
        latest = context.latest_message
        if not latest or not latest.content:
            return None
        
        content = latest.content
        sender = latest.from_agent or "SYSTEM"
        
        # Skip system startup messages
        if content == "[startup]":
            return None
        
        # Handle subteam_result messages from child workspaces
        if content.startswith("[subteam_result]"):
            try:
                # Parse the result payload
                payload_str = content[len("[subteam_result]"):]
                payload = json.loads(payload_str)
                
                child_workspace_id = payload.get("child_workspace_id", "unknown")
                result_message = payload.get("result_message", "")
                
                # Store the result
                received_results[child_workspace_id] = result_message
                
                logger.info(
                    f"PlannerAgent received result from workspace {child_workspace_id}. "
                    f"Progress: {len(received_results)}/{child_workspace_count}"
                )
                
                # Check if all child workspaces have completed
                if len(received_results) >= child_workspace_count and child_workspace_count > 0:
                    # All teams have reported - synthesize and deliver final output
                    return await _synthesize_final_output(
                        planner_llm=planner_llm,
                        results=received_results,
                        child_descriptions=child_descriptions or [],
                        context=context,
                    )
                else:
                    # Still waiting for more teams
                    return {
                        "send_to": "USER",
                        "visible_reasoning": (
                            f"Received results from team. "
                            f"Progress: {len(received_results)}/{child_workspace_count} teams completed. "
                            f"Waiting for remaining teams to finish their work."
                        ),
                        "collaboration": "",
                        "content": (
                            f"📊 **Progress Update**: {len(received_results)}/{child_workspace_count} "
                            f"teams have completed their work. Waiting for remaining teams..."
                        ),
                    }
                    
            except (json.JSONDecodeError, KeyError) as e:
                logger.error(f"Failed to parse subteam_result: {e}")
                return None
        
        # Handle initial task from USER
        if sender.upper() == "USER":
            # Acknowledge the task and explain coordination approach
            team_info = ""
            if child_descriptions:
                team_info = "\n".join([
                    f"  • Team {i+1}: {desc}"
                    for i, desc in enumerate(child_descriptions)
                ])
            else:
                team_info = f"  • {child_workspace_count} specialized teams"
            
            return {
                "send_to": "USER",
                "visible_reasoning": (
                    f"I have received the user's task. As the CEO/Planner, I am now coordinating "
                    f"{child_workspace_count} specialized teams to execute different aspects of this task. "
                    f"Each team will work on their assigned subtask in parallel. "
                    f"I will synthesize all results when teams complete their work."
                ),
                "collaboration": "",
                "content": (
                    f"🎯 **Task Received**\n\n"
                    f"I'm coordinating {child_workspace_count} specialized teams to work on your task:\n\n"
                    f"{team_info}\n\n"
                    f"I'll synthesize their outputs and deliver a comprehensive result once all teams complete."
                ),
            }
        
        # Handle other messages (relay or respond)
        return {
            "send_to": "USER",
            "visible_reasoning": f"Received message from {sender}. Forwarding coordination update to user.",
            "collaboration": "",
            "content": f"📬 Update from {sender}: {content[:500]}...",
        }
    
    return planner_logic


async def _synthesize_final_output(
    planner_llm: PlannerLLM,
    results: Dict[str, str],
    child_descriptions: List[str],
    context: AgentLogicContext,
) -> dict:
    """
    Synthesize results from all child workspaces into a final deliverable.
    
    Uses the PlannerLLM to intelligently combine outputs from all teams
    into a coherent, comprehensive response for the user.
    
    Args:
        planner_llm: The PlannerLLM instance for synthesis
        results: Dictionary of workspace_id -> result content
        child_descriptions: List of child workspace task descriptions
        context: The agent logic context
        
    Returns:
        Final synthesized response to send to USER
    """
    # Build a summary of all results
    results_summary = "\n\n".join([
        f"=== Team {i+1} Result ===\n{result}"
        for i, (workspace_id, result) in enumerate(results.items())
    ])
    
    # Use LLM to synthesize (if available)
    try:
        synthesis_prompt = f"""You are synthesizing results from multiple teams into a final deliverable.

TEAM RESULTS:
{results_summary}

Create a comprehensive, well-organized final output that:
1. Integrates all team contributions coherently
2. Highlights key deliverables and outcomes
3. Provides actionable next steps if applicable
4. Is formatted for easy reading (use headers, bullet points, etc.)

Respond with the synthesized content only - no meta-commentary."""

        # Make LLM call for synthesis
        synthesis = await planner_llm._call_llm(
            system_prompt="You are an expert at synthesizing multi-team outputs into cohesive deliverables.",
            user_prompt=synthesis_prompt,
        )
        
        final_content = synthesis
        
    except Exception as e:
        logger.warning(f"LLM synthesis failed, using manual concatenation: {e}")
        # Fallback: concatenate results
        final_content = (
            "# Final Deliverable\n\n"
            "The following is a compilation of work from all teams:\n\n"
            + results_summary
        )
    
    return {
        "send_to": "USER",
        "visible_reasoning": (
            f"All {len(results)} teams have completed their work. "
            f"I have synthesized their outputs into a comprehensive final deliverable. "
            f"This represents the consolidated effort of all specialized teams."
        ),
        "collaboration": "",
        "content": (
            f"✅ **All Teams Complete** - Final Deliverable\n\n"
            f"{final_content}"
        ),
    }


class PlannerAgent(Agent):
    """
    The CEO agent that coordinates all work in the multi-agent system.
    
    This agent:
    - Is the ONLY agent in the root workspace
    - Is the ONLY agent the user interacts with directly
    - Does NOT execute tasks - it coordinates execution agents
    - Synthesizes results from child workspaces
    - Delivers final consolidated output to the user
    
    The PlannerAgent wraps a PlannerLLM and uses it for:
    - LLM calls for synthesis and coordination
    - (Task planning is done BEFORE workspace creation, not by this agent)
    
    Example:
        ```python
        # Create PlannerLLM
        planner_llm = PlannerLLM(provider="anthropic", api_key="...")
        
        # Create PlannerAgent for root workspace
        planner_agent = PlannerAgent(
            planner_llm=planner_llm,
            child_workspace_count=3,
            child_descriptions=["Team 1 task", "Team 2 task", "Team 3 task"],
        )
        
        # Register as the only agent in root workspace
        AgentRuntimeRegistry.register("planner", planner_agent)
        ```
    """
    
    def __init__(
        self,
        planner_llm: PlannerLLM,
        child_workspace_count: int = 0,
        child_descriptions: Optional[List[str]] = None,
        name: str = "planner",
        description: str = "CEO/Planner agent that coordinates all execution teams",
    ):
        """
        Initialize the PlannerAgent.
        
        Args:
            planner_llm: The PlannerLLM instance for LLM calls
            child_workspace_count: Number of child workspaces to coordinate
            child_descriptions: Optional list of child workspace task descriptions
            name: Agent name (default: "planner")
            description: Agent description
        """
        self.planner_llm = planner_llm
        self.child_workspace_count = child_workspace_count
        self.child_descriptions = child_descriptions or []
        self._received_results: Dict[str, str] = {}
        
        # Create the logic function
        async def planner_logic(context: AgentLogicContext) -> Optional[dict]:
            return await self._process_message(context)
        
        # Initialize parent Agent class
        super().__init__(
            name=name,
            description=description,
            logic=planner_logic,
            default_target="USER",
            capabilities=["coordination", "synthesis", "planning", "orchestration"],
        )
    
    async def _process_message(self, context: AgentLogicContext) -> Optional[dict]:
        """
        Process incoming messages to the PlannerAgent.
        
        Handles:
        1. Initial task from USER
        2. [subteam_result] messages from child workspaces
        3. Synthesis and final delivery when all teams complete
        """
        latest = context.latest_message
        if not latest or not latest.content:
            return None
        
        content = latest.content
        sender = latest.from_agent or "SYSTEM"
        
        # Skip system startup messages
        if content == "[startup]":
            return None
        
        # Handle subteam_result messages from child workspaces
        if content.startswith("[subteam_result]"):
            return await self._handle_subteam_result(content, context)
        
        # Handle initial task from USER
        if sender.upper() == "USER":
            return self._handle_user_task(content, context)
        
        # Handle other messages
        return {
            "send_to": "USER",
            "visible_reasoning": f"Received coordination update from {sender}.",
            "collaboration": "",
            "content": f"📬 Update from {sender}: {content[:500]}",
        }
    
    async def _handle_subteam_result(
        self,
        content: str,
        context: AgentLogicContext,
    ) -> Optional[dict]:
        """Handle [subteam_result] messages from child workspaces."""
        try:
            # Parse the result payload
            payload_str = content[len("[subteam_result]"):]
            payload = json.loads(payload_str)
            
            child_workspace_id = payload.get("child_workspace_id", "unknown")
            result_message = payload.get("result_message", "")
            
            # Store the result
            self._received_results[child_workspace_id] = result_message
            
            logger.info(
                f"PlannerAgent received result from workspace {child_workspace_id}. "
                f"Progress: {len(self._received_results)}/{self.child_workspace_count}"
            )
            
            # Check if all child workspaces have completed
            if len(self._received_results) >= self.child_workspace_count > 0:
                # All teams have reported - synthesize and deliver final output
                return await self._synthesize_and_deliver()
            else:
                # Still waiting for more teams
                return {
                    "send_to": "USER",
                    "visible_reasoning": (
                        f"Received results from team. "
                        f"Progress: {len(self._received_results)}/{self.child_workspace_count} teams completed."
                    ),
                    "collaboration": "",
                    "content": (
                        f"📊 **Progress Update**: {len(self._received_results)}/{self.child_workspace_count} "
                        f"teams have completed. Waiting for remaining teams..."
                    ),
                }
                
        except (json.JSONDecodeError, KeyError) as e:
            logger.error(f"Failed to parse subteam_result: {e}")
            return None
    
    def _handle_user_task(self, content: str, context: AgentLogicContext) -> dict:
        """Handle initial task from USER."""
        team_info = ""
        if self.child_descriptions:
            team_info = "\n".join([
                f"  • Team {i+1}: {desc}"
                for i, desc in enumerate(self.child_descriptions)
            ])
        else:
            team_info = f"  • {self.child_workspace_count} specialized teams"
        
        return {
            "send_to": "USER",
            "visible_reasoning": (
                f"Task received. Coordinating {self.child_workspace_count} specialized teams. "
                f"Each team will work on their assigned subtask in parallel."
            ),
            "collaboration": "",
            "content": (
                f"🎯 **Task Received**\n\n"
                f"I'm coordinating {self.child_workspace_count} specialized teams:\n\n"
                f"{team_info}\n\n"
                f"I'll synthesize their outputs and deliver a comprehensive result."
            ),
        }
    
    async def _synthesize_and_deliver(self) -> dict:
        """Synthesize all results and deliver final output to USER."""
        # Build a summary of all results
        results_summary = "\n\n".join([
            f"=== Team {i+1} Result ===\n{result}"
            for i, (workspace_id, result) in enumerate(self._received_results.items())
        ])
        
        # Use LLM to synthesize
        try:
            synthesis_prompt = f"""Synthesize these team results into a final deliverable:

TEAM RESULTS:
{results_summary}

Create a comprehensive, well-organized output that:
1. Integrates all team contributions coherently
2. Highlights key deliverables and outcomes
3. Provides actionable next steps if applicable

Respond with synthesized content only."""

            synthesis = await self.planner_llm._call_llm(
                system_prompt="You synthesize multi-team outputs into cohesive deliverables.",
                user_prompt=synthesis_prompt,
            )
            
            final_content = synthesis
            
        except Exception as e:
            logger.warning(f"LLM synthesis failed: {e}")
            final_content = (
                "# Final Deliverable\n\n"
                + results_summary
            )
        
        return {
            "send_to": "USER",
            "visible_reasoning": (
                f"All {len(self._received_results)} teams complete. "
                f"Synthesized outputs into final deliverable."
            ),
            "collaboration": "",
            "content": f"✅ **All Teams Complete**\n\n{final_content}",
        }
    
    def reset_results(self) -> None:
        """Reset the received results for a new task."""
        self._received_results.clear()
    
    def update_child_info(
        self,
        child_workspace_count: int,
        child_descriptions: Optional[List[str]] = None,
    ) -> None:
        """
        Update information about child workspaces.
        
        Call this after task planning to inform the PlannerAgent
        about the child workspaces it will be coordinating.
        
        Args:
            child_workspace_count: Number of child workspaces
            child_descriptions: Optional list of child workspace task descriptions
        """
        self.child_workspace_count = child_workspace_count
        self.child_descriptions = child_descriptions or []
        self.reset_results()
    
    def __repr__(self) -> str:
        """String representation."""
        return (
            f"PlannerAgent(name='{self.name}', "
            f"child_workspaces={self.child_workspace_count})"
        )

