"""
Workspace Engine - The core intelligence layer.

This is where everything comes together:
- Message passing
- Protocol enforcement
- Scheduling
- Shared state management
- Lifecycle management

This is NOT a wrapper. This is the SUBSTRATE.
"""

from dataclasses import dataclass, field
from enum import Enum
from typing import Optional, Callable, Any
from datetime import datetime
import asyncio
import logging
import uuid

from synqed.message_model import (
    StructuredMessage,
    MessageType,
    MessagePriority,
    MessageMetadata,
)
from synqed.protocols import Protocol, ProtocolPhase, ProtocolConfig, ProtocolType, create_protocol
from synqed.scheduler import Scheduler, SchedulerConfig, SchedulingStrategy
from synqed.shared_state import SharedState

logger = logging.getLogger(__name__)


class WorkspaceLifecycleState(Enum):
    """Lifecycle states of the workspace."""
    UNINITIALIZED = "uninitialized"
    INITIALIZING = "initializing"
    ACTIVE = "active"
    PAUSED = "paused"
    TERMINATING = "terminating"
    TERMINATED = "terminated"
    FAILED = "failed"


@dataclass
class AgentInfo:
    """Information about an agent in the workspace."""
    agent_id: str
    agent_name: str
    role: Optional[str] = None
    skills: list[str] = field(default_factory=list)
    constraints: dict[str, Any] = field(default_factory=dict)
    
    # Execution interface
    executor: Optional[Callable] = None  # Function to call the agent
    
    # Metadata
    joined_at: datetime = field(default_factory=datetime.now)
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class WorkspaceConfig:
    """Configuration for the workspace."""
    workspace_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    name: str = "Workspace"
    description: str = ""
    
    # Protocol
    protocol_type: ProtocolType = ProtocolType.CRITIQUE_AND_REVISE
    protocol_config: Optional[ProtocolConfig] = None
    
    # Scheduler
    scheduler_config: Optional[SchedulerConfig] = None
    
    # Constraints
    max_rounds: int = 5
    max_messages: int = 100
    max_total_tokens: Optional[int] = None
    timeout_seconds: Optional[int] = 300
    
    # Behavior
    auto_start: bool = False
    persist_state: bool = False
    
    # Callbacks
    on_message: Optional[Callable[[StructuredMessage], None]] = None
    on_phase_change: Optional[Callable[[ProtocolPhase], None]] = None
    on_round_end: Optional[Callable[[int, dict], None]] = None


@dataclass
class WorkspaceStatistics:
    """Statistics about workspace execution."""
    workspace_id: str
    start_time: datetime
    end_time: Optional[datetime] = None
    
    # Execution
    total_rounds: int = 0
    total_messages: int = 0
    total_tokens: int = 0
    
    # By agent
    messages_by_agent: dict[str, int] = field(default_factory=dict)
    tokens_by_agent: dict[str, int] = field(default_factory=dict)
    
    # By type
    messages_by_type: dict[str, int] = field(default_factory=dict)
    
    # Termination
    termination_reason: Optional[str] = None
    success: bool = False
    
    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "workspace_id": self.workspace_id,
            "start_time": self.start_time.isoformat(),
            "end_time": self.end_time.isoformat() if self.end_time else None,
            "duration_seconds": (
                (self.end_time - self.start_time).total_seconds()
                if self.end_time else None
            ),
            "total_rounds": self.total_rounds,
            "total_messages": self.total_messages,
            "total_tokens": self.total_tokens,
            "messages_by_agent": self.messages_by_agent,
            "tokens_by_agent": self.tokens_by_agent,
            "messages_by_type": self.messages_by_type,
            "termination_reason": self.termination_reason,
            "success": self.success,
        }


class WorkspaceEngine:
    """
    The Workspace Engine - Core intelligence layer for multi-agent collaboration.
    
    Lifecycle:
    1. INIT
       - Load agent cards
       - Load task
       - Assign roles
       - Set rules
       - Spawn workspace state
    
    2. ROUNDS
       - Each agent receives state
       - Agent outputs structured message
       - Workspace merges results
       - Workspace updates state
       - Workspace decides if more rounds are needed
    
    3. TERMINATION
       - Output combined result
       - Return artifacts
       - Store transcript
       - Destroy workspace
    
    This is the HEART of Synq. If this doesn't work, nothing works.
    """
    
    def __init__(self, config: WorkspaceConfig):
        self.config = config
        self.workspace_id = config.workspace_id
        
        # Lifecycle
        self.state = WorkspaceLifecycleState.UNINITIALIZED
        self.created_at = datetime.now()
        self.started_at: Optional[datetime] = None
        self.terminated_at: Optional[datetime] = None
        
        # Agents
        self.agents: dict[str, AgentInfo] = {}
        
        # Core components (initialized in init())
        self.protocol: Optional[Protocol] = None
        self.scheduler: Optional[Scheduler] = None
        self.shared_state: Optional[SharedState] = None
        
        # Messages
        self.messages: list[StructuredMessage] = []
        self.message_index: dict[str, StructuredMessage] = {}  # message_id -> message
        
        # Task
        self.task: Optional[str] = None
        self.context: dict[str, Any] = {}
        
        # Statistics
        self.stats = WorkspaceStatistics(
            workspace_id=self.workspace_id,
            start_time=self.created_at,
        )
        
        # Result
        self.final_result: Optional[str] = None
        
        self.logger = logging.getLogger(f"{__name__}.{self.workspace_id[:8]}")
    
    # ========================================================================
    # LIFECYCLE: INITIALIZATION
    # ========================================================================
    
    async def initialize(
        self,
        task: str,
        agents: list[AgentInfo],
        context: Optional[dict[str, Any]] = None,
    ) -> None:
        """
        Initialize the workspace.
        
        Args:
            task: The task to accomplish
            agents: List of agents to participate
            context: Optional additional context
        """
        if self.state != WorkspaceLifecycleState.UNINITIALIZED:
            raise RuntimeError(f"Cannot initialize workspace in state: {self.state}")
        
        self.state = WorkspaceLifecycleState.INITIALIZING
        self.logger.info(f"Initializing workspace: {self.config.name}")
        
        # Store task and context
        self.task = task
        self.context = context or {}
        
        # Register agents
        for agent in agents:
            self.agents[agent.agent_id] = agent
            self.stats.messages_by_agent[agent.agent_id] = 0
            self.stats.tokens_by_agent[agent.agent_id] = 0
        
        self.logger.info(f"Registered {len(agents)} agents")
        
        # Create protocol
        protocol_config = self.config.protocol_config or ProtocolConfig(
            max_rounds=self.config.max_rounds,
            max_total_tokens=self.config.max_total_tokens,
        )
        self.protocol = create_protocol(self.config.protocol_type, protocol_config)
        self.protocol.state.turn_order = list(self.agents.keys())
        
        self.logger.info(f"Created protocol: {self.config.protocol_type.value}")
        
        # Create scheduler
        scheduler_config = self.config.scheduler_config or SchedulerConfig(
            strategy=SchedulingStrategy.PROTOCOL_DRIVEN,
        )
        agent_tuples = [(a.agent_id, a.agent_name) for a in agents]
        self.scheduler = Scheduler(scheduler_config, self.protocol, agent_tuples)
        
        self.logger.info(f"Created scheduler: {scheduler_config.strategy.value}")
        
        # Create shared state
        self.shared_state = SharedState()
        
        # Initialize shared state with task
        self.shared_state.set("task", task, "system")
        self.shared_state.set("context", self.context, "system")
        self.shared_state.set_read_only("task")
        self.shared_state.set_read_only("context")
        
        # Initialize agent memories
        for agent_id in self.agents:
            memory = self.shared_state.get_agent_memory(agent_id)
            memory.set("task", task)
            memory.set("role", self.agents[agent_id].role)
        
        self.logger.info("Initialized shared state")
        
        # Send initialization message
        init_msg = self._create_system_message(
            content=self._build_initialization_message(),
            phase=ProtocolPhase.INITIALIZATION,
        )
        self._record_message(init_msg)
        
        # Advance protocol to first real phase
        self.protocol.advance_phase()
        
        # Update state
        self.state = WorkspaceLifecycleState.ACTIVE
        self.started_at = datetime.now()
        self.logger.info("Workspace initialized and active")
    
    def _build_initialization_message(self) -> str:
        """Build the initialization message for agents."""
        lines = [
            f"🚀 WORKSPACE INITIALIZED: {self.config.name}",
            "",
            f"TASK: {self.task}",
            "",
            f"PROTOCOL: {self.config.protocol_type.value}",
            f"PARTICIPANTS: {', '.join(a.agent_name for a in self.agents.values())}",
            "",
            "PHASES:",
        ]
        
        for phase in self.protocol.get_phases():
            lines.append(f"  {phase.value}")
        
        lines.extend([
            "",
            "RULES:",
            f"  • Maximum rounds: {self.config.max_rounds}",
            f"  • Maximum messages: {self.config.max_messages}",
            f"  • Follow the protocol phases",
            f"  • Use structured message types",
            f"  • Collaborate to solve the task",
            "",
            "Let's begin!",
        ])
        
        return "\n".join(lines)
    
    # ========================================================================
    # LIFECYCLE: EXECUTION (ROUNDS)
    # ========================================================================
    
    async def execute_round(self) -> dict[str, Any]:
        """
        Execute one round of the protocol.
        
        Returns:
            Round statistics
        """
        if self.state != WorkspaceLifecycleState.ACTIVE:
            raise RuntimeError(f"Cannot execute round in state: {self.state}")
        
        round_num = self.stats.total_rounds + 1
        self.logger.info(f"=== ROUND {round_num} START ===")
        
        # Start round in scheduler
        self.scheduler.start_round(round_num)
        
        # Start round in protocol
        if round_num > 1:
            self.protocol.advance_round()
        
        # Send round start message
        round_start_msg = self._create_system_message(
            content=f"🔄 ROUND {round_num} - Phase: {self.protocol.state.current_phase.value}",
            phase=self.protocol.state.current_phase,
        )
        self._record_message(round_start_msg)
        
        # Execute turns until round complete
        turn_count = 0
        max_turns_per_round = len(self.agents) * 10  # Safety limit
        
        while turn_count < max_turns_per_round:
            # Check if phase should advance
            if self.protocol.should_advance_phase(self.messages):
                self.logger.info(f"Advancing from phase: {self.protocol.state.current_phase.value}")
                self.protocol.advance_phase()
                
                # Notify
                if self.config.on_phase_change:
                    try:
                        self.config.on_phase_change(self.protocol.state.current_phase)
                    except Exception as e:
                        self.logger.error(f"Phase change callback error: {e}")
                
                # Check if protocol complete
                if self.protocol.state.current_phase == ProtocolPhase.COMPLETED:
                    break
                
                # Send phase change message
                phase_msg = self._create_system_message(
                    content=f"📌 PHASE: {self.protocol.state.current_phase.value}",
                    phase=self.protocol.state.current_phase,
                )
                self._record_message(phase_msg)
            
            # Get next speaker
            next_speaker = self.scheduler.get_next_speaker(self.messages)
            
            if next_speaker is None:
                self.logger.info("No next speaker - round complete")
                break
            
            # Execute agent turn
            try:
                await self._execute_agent_turn(next_speaker)
                turn_count += 1
            except Exception as e:
                self.logger.error(f"Agent turn failed: {e}", exc_info=True)
                # Continue with next agent
        
        # End round
        self.stats.total_rounds = round_num
        round_stats = self.scheduler.end_round()
        
        # Notify
        if self.config.on_round_end:
            try:
                self.config.on_round_end(round_num, round_stats)
            except Exception as e:
                self.logger.error(f"Round end callback error: {e}")
        
        self.logger.info(f"=== ROUND {round_num} END === ({turn_count} turns)")
        
        return round_stats
    
    async def _execute_agent_turn(self, agent_id: str) -> None:
        """
        Execute a single agent turn.
        
        Args:
            agent_id: Agent to execute
        """
        agent = self.agents[agent_id]
        self.logger.debug(f"Executing turn for: {agent.agent_name}")
        
        # Build context for agent
        agent_context = self._build_agent_context(agent_id)
        
        # Call agent executor
        if agent.executor is None:
            self.logger.warning(f"No executor for agent: {agent_id}")
            return
        
        try:
            # Execute agent
            response = await self._call_agent(agent, agent_context)
            
            # Parse response into structured message
            message = self._parse_agent_response(agent_id, response)
            
            # Validate message against protocol
            allowed, reason = self.protocol.is_message_allowed(message, agent_id)
            if not allowed:
                self.logger.warning(f"Message not allowed: {reason}")
                # Send feedback to agent
                feedback_msg = self._create_system_message(
                    content=f"❌ Message rejected: {reason}",
                    recipient_ids=[agent_id],
                )
                self._record_message(feedback_msg)
                return
            
            # Record message
            self._record_message(message)
            self.scheduler.record_turn(agent_id, message)
            self.protocol.record_message(message)
            
            # Update shared state based on message
            self._process_message_effects(message)
            
            # Notify callback
            if self.config.on_message:
                try:
                    self.config.on_message(message)
                except Exception as e:
                    self.logger.error(f"Message callback error: {e}")
        
        except Exception as e:
            self.logger.error(f"Agent execution failed: {e}", exc_info=True)
            raise
    
    def _build_agent_context(self, agent_id: str) -> dict[str, Any]:
        """Build context dictionary for an agent."""
        agent = self.agents[agent_id]
        
        # Get recent messages (last 20)
        recent_messages = self.messages[-20:]
        
        # Get shared state
        state_context = self.shared_state.get_context_for_agent(agent_id)
        
        # Build context
        context = {
            "task": self.task,
            "workspace_id": self.workspace_id,
            "workspace_name": self.config.name,
            "agent_id": agent_id,
            "agent_name": agent.agent_name,
            "agent_role": agent.role,
            "agent_skills": agent.skills,
            "current_phase": self.protocol.state.current_phase.value,
            "current_round": self.stats.total_rounds,
            "protocol": self.config.protocol_type.value,
            "shared_state": state_context["shared_state"],
            "my_memory": state_context["my_memory"],
            "recent_messages": [
                {
                    "sender": msg.sender_name,
                    "type": msg.message_type.value,
                    "content": msg.content,
                    "phase": msg.metadata.phase,
                }
                for msg in recent_messages
            ],
            "other_agents": [
                {"id": a.agent_id, "name": a.agent_name, "role": a.role}
                for a in self.agents.values()
                if a.agent_id != agent_id
            ],
        }
        
        return context
    
    async def _call_agent(self, agent: AgentInfo, context: dict[str, Any]) -> str:
        """Call an agent's executor."""
        # This is where we'd call the actual agent
        # For now, this is a placeholder
        if asyncio.iscoroutinefunction(agent.executor):
            return await agent.executor(context)
        else:
            return agent.executor(context)
    
    def _parse_agent_response(
        self,
        agent_id: str,
        response: str,
    ) -> StructuredMessage:
        """
        Parse agent response into a structured message.
        
        For now, we'll use simple heuristics. In production, agents should
        return structured data directly.
        """
        agent = self.agents[agent_id]
        
        # Determine message type based on phase and content
        message_type = self._infer_message_type(response)
        
        # Create message
        metadata = MessageMetadata(
            round_number=self.stats.total_rounds,
            phase=self.protocol.state.current_phase.value,
            token_count=len(response.split()),  # Rough approximation
        )
        
        message = StructuredMessage(
            message_type=message_type,
            sender_id=agent_id,
            sender_name=agent.agent_name,
            content=response,
            metadata=metadata,
        )
        
        return message
    
    def _infer_message_type(self, content: str) -> MessageType:
        """Infer message type from content (simple heuristic)."""
        content_lower = content.lower()
        
        # Check phase
        phase = self.protocol.state.current_phase
        
        if phase == ProtocolPhase.PROPOSAL:
            return MessageType.SUGGESTION
        elif phase == ProtocolPhase.CRITIQUE:
            return MessageType.CRITIQUE
        elif phase == ProtocolPhase.VOTING:
            return MessageType.VOTE
        elif phase == ProtocolPhase.REFINEMENT:
            return MessageType.SUGGESTION
        
        # Content-based heuristics
        if any(word in content_lower for word in ["i suggest", "i propose", "my idea"]):
            return MessageType.SUGGESTION
        elif any(word in content_lower for word in ["i observe", "i notice", "data shows"]):
            return MessageType.OBSERVATION
        elif any(word in content_lower for word in ["i vote", "my vote"]):
            return MessageType.VOTE
        elif any(word in content_lower for word in ["critique", "concern", "issue"]):
            return MessageType.CRITIQUE
        
        # Default
        return MessageType.RESPONSE
    
    def _process_message_effects(self, message: StructuredMessage) -> None:
        """Process side effects of a message on shared state."""
        # Example: if message contains a plan fragment, store it
        if message.message_type == MessageType.PLAN_FRAGMENT:
            plan_fragments = self.shared_state.get("plan_fragments", [])
            plan_fragments.append({
                "from": message.sender_name,
                "content": message.content,
                "message_id": message.metadata.message_id,
            })
            self.shared_state.set("plan_fragments", plan_fragments, message.sender_id)
        
        # Example: track proposals
        elif message.message_type == MessageType.SUGGESTION:
            self.shared_state.append("proposals", message.sender_id, {
                "from": message.sender_name,
                "content": message.content,
                "message_id": message.metadata.message_id,
            })
        
        # Example: track votes
        elif message.message_type == MessageType.VOTE:
            option = message.structured_data.get("option")
            if option:
                votes = self.shared_state.get("votes", {})
                votes[message.sender_id] = option
                self.shared_state.set("votes", votes, message.sender_id)
    
    # ========================================================================
    # LIFECYCLE: EXECUTION (MAIN LOOP)
    # ========================================================================
    
    async def run(self) -> str:
        """
        Run the workspace until termination.
        
        Returns:
            Final result
        """
        if self.state != WorkspaceLifecycleState.ACTIVE:
            raise RuntimeError(f"Cannot run workspace in state: {self.state}")
        
        self.logger.info("Starting workspace execution")
        
        # Main execution loop
        while True:
            # Check termination conditions
            should_terminate, reason = self._should_terminate()
            if should_terminate:
                self.logger.info(f"Terminating: {reason}")
                await self.terminate(reason, success=True)
                break
            
            # Execute round
            try:
                await self.execute_round()
            except Exception as e:
                self.logger.error(f"Round execution failed: {e}", exc_info=True)
                await self.terminate(f"Round execution failed: {e}", success=False)
                break
        
        return self.final_result or "No result"
    
    def _should_terminate(self) -> tuple[bool, Optional[str]]:
        """Check if workspace should terminate."""
        # Protocol termination
        should_terminate, reason = self.protocol.should_terminate(self.messages)
        if should_terminate:
            return True, reason
        
        # Message limit
        if len(self.messages) >= self.config.max_messages:
            return True, f"Reached message limit ({self.config.max_messages})"
        
        # Token limit
        if self.config.max_total_tokens:
            if self.stats.total_tokens >= self.config.max_total_tokens:
                return True, f"Reached token limit ({self.config.max_total_tokens})"
        
        # Timeout
        if self.config.timeout_seconds and self.started_at:
            elapsed = (datetime.now() - self.started_at).total_seconds()
            if elapsed >= self.config.timeout_seconds:
                return True, f"Timeout after {elapsed:.1f}s"
        
        return False, None
    
    # ========================================================================
    # LIFECYCLE: TERMINATION
    # ========================================================================
    
    async def terminate(self, reason: str, success: bool = True) -> None:
        """
        Terminate the workspace.
        
        Args:
            reason: Reason for termination
            success: Whether termination is successful completion
        """
        if self.state == WorkspaceLifecycleState.TERMINATED:
            return
        
        self.logger.info(f"Terminating workspace: {reason}")
        self.state = WorkspaceLifecycleState.TERMINATING
        
        # Synthesize final result
        self.final_result = await self._synthesize_final_result()
        
        # Update statistics
        self.stats.end_time = datetime.now()
        self.stats.termination_reason = reason
        self.stats.success = success
        self.stats.total_messages = len(self.messages)
        self.stats.total_tokens = sum(
            msg.metadata.token_count for msg in self.messages
        )
        
        # Send termination message
        term_msg = self._create_system_message(
            content=f"🏁 WORKSPACE TERMINATED: {reason}",
        )
        self._record_message(term_msg)
        
        # Update state
        self.state = WorkspaceLifecycleState.TERMINATED
        self.terminated_at = datetime.now()
        
        self.logger.info(f"Workspace terminated. Success: {success}")
    
    async def _synthesize_final_result(self) -> str:
        """Synthesize final result from workspace state."""
        # Get all proposals from final phase
        proposals = self.shared_state.get("proposals", [])
        
        if not proposals:
            return "No proposals generated."
        
        # Simple synthesis: combine all proposals
        lines = [
            "FINAL RESULT:",
            "",
            f"Task: {self.task}",
            "",
            "Agent Contributions:",
        ]
        
        for prop in proposals[-len(self.agents):]:  # Last N proposals
            lines.append(f"\n{prop['from']}:")
            lines.append(prop['content'][:500])
        
        return "\n".join(lines)
    
    # ========================================================================
    # MESSAGE MANAGEMENT
    # ========================================================================
    
    def _create_system_message(
        self,
        content: str,
        phase: Optional[ProtocolPhase] = None,
        recipient_ids: Optional[list[str]] = None,
    ) -> StructuredMessage:
        """Create a system message."""
        metadata = MessageMetadata(
            round_number=self.stats.total_rounds,
            phase=phase.value if phase else self.protocol.state.current_phase.value,
            priority=MessagePriority.HIGH,
        )
        
        return StructuredMessage(
            message_type=MessageType.SYSTEM,
            sender_id="system",
            sender_name="System",
            content=content,
            recipient_ids=recipient_ids or [],
            metadata=metadata,
        )
    
    def _record_message(self, message: StructuredMessage) -> None:
        """Record a message in the workspace."""
        self.messages.append(message)
        self.message_index[message.metadata.message_id] = message
        
        # Update statistics
        sender = message.sender_id
        if sender in self.stats.messages_by_agent:
            self.stats.messages_by_agent[sender] += 1
            self.stats.tokens_by_agent[sender] += message.metadata.token_count
        
        msg_type = message.message_type.value
        self.stats.messages_by_type[msg_type] = self.stats.messages_by_type.get(msg_type, 0) + 1
        
        # Record in agent's memory
        if sender != "system" and sender in self.agents:
            self.shared_state.record_action(
                sender,
                f"Sent {message.message_type.value}: {message.content[:100]}"
            )
    
    # ========================================================================
    # QUERIES
    # ========================================================================
    
    def get_messages(
        self,
        message_type: Optional[MessageType] = None,
        sender_id: Optional[str] = None,
        phase: Optional[ProtocolPhase] = None,
        limit: Optional[int] = None,
    ) -> list[StructuredMessage]:
        """Get messages with optional filters."""
        messages = self.messages
        
        if message_type:
            messages = [m for m in messages if m.message_type == message_type]
        
        if sender_id:
            messages = [m for m in messages if m.sender_id == sender_id]
        
        if phase:
            messages = [m for m in messages if m.metadata.phase == phase.value]
        
        if limit:
            messages = messages[-limit:]
        
        return messages
    
    def get_transcript(self) -> str:
        """Get full transcript of workspace."""
        lines = [
            f"WORKSPACE TRANSCRIPT: {self.config.name}",
            f"ID: {self.workspace_id}",
            f"Task: {self.task}",
            "=" * 80,
            "",
        ]
        
        for msg in self.messages:
            timestamp = msg.metadata.timestamp.strftime("%H:%M:%S")
            lines.append(
                f"[{timestamp}] [{msg.metadata.phase}] "
                f"{msg.sender_name} ({msg.message_type.value}):"
            )
            lines.append(f"  {msg.content}")
            lines.append("")
        
        lines.extend([
            "=" * 80,
            f"Total messages: {len(self.messages)}",
            f"Total rounds: {self.stats.total_rounds}",
            f"Success: {self.stats.success}",
        ])
        
        return "\n".join(lines)
    
    def __repr__(self) -> str:
        """String representation."""
        return (
            f"WorkspaceEngine(id={self.workspace_id[:8]}, "
            f"state={self.state.value}, "
            f"agents={len(self.agents)}, "
            f"messages={len(self.messages)})"
        )

