"""
Workspace - Collaborative environment for multi-agent systems.

The Workspace provides a temporary sandbox/playground where AI agents can:
- Share files and data
- Collaborate on tasks
- Communicate with each other
- Track progress and artifacts
- Maintain shared state

This enables true multi-agent collaboration beyond simple delegation.
"""

import asyncio
import json
import logging
import shutil
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any, Callable, Coroutine

from synqed.agent import Agent
from synqed.client import Client

logger = logging.getLogger(__name__)


class WorkspaceState(str, Enum):
    """State of the workspace."""
    CREATED = "created"
    ACTIVE = "active"
    PAUSED = "paused"
    COMPLETED = "completed"
    CANCELLED = "cancelled"
    ERROR = "error"


class MessageType(str, Enum):
    """Type of message in the workspace."""
    SYSTEM = "system"
    AGENT = "agent"
    USER = "user"
    ARTIFACT = "artifact"
    STATUS = "status"


@dataclass
class WorkspaceMessage:
    """A message in the workspace."""
    message_id: str
    timestamp: datetime
    sender_id: str
    sender_name: str
    message_type: MessageType
    content: str
    metadata: dict[str, Any] = field(default_factory=dict)
    
    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "message_id": self.message_id,
            "timestamp": self.timestamp.isoformat(),
            "sender_id": self.sender_id,
            "sender_name": self.sender_name,
            "message_type": self.message_type.value,
            "content": self.content,
            "metadata": self.metadata,
        }


@dataclass
class WorkspaceArtifact:
    """An artifact created in the workspace."""
    artifact_id: str
    name: str
    artifact_type: str  # file, data, result, etc.
    content: str | bytes | dict[str, Any]
    created_by: str
    created_at: datetime
    metadata: dict[str, Any] = field(default_factory=dict)
    
    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        # Handle different content types
        if isinstance(self.content, bytes):
            content = f"<binary data: {len(self.content)} bytes>"
        elif isinstance(self.content, dict):
            content = json.dumps(self.content, indent=2)
        else:
            content = str(self.content)
        
        return {
            "artifact_id": self.artifact_id,
            "name": self.name,
            "artifact_type": self.artifact_type,
            "content": content,
            "created_by": self.created_by,
            "created_at": self.created_at.isoformat(),
            "metadata": self.metadata,
        }


@dataclass
class WorkspaceParticipant:
    """A participant in the workspace."""
    participant_id: str
    name: str
    role: str  # agent, user, system
    agent: Agent | None = None
    agent_url: str | None = None
    joined_at: datetime = field(default_factory=datetime.now)
    metadata: dict[str, Any] = field(default_factory=dict)


class Workspace:
    """
    Collaborative workspace for multi-agent systems.
    
    The Workspace creates a temporary environment where agents can collaborate,
    share resources, and work together on complex tasks.
    
    Example:
        ```python
        # Create a workspace
        workspace = Workspace(
            name="Recipe Planning",
            description="Plan meals and create shopping lists"
        )
        
        # Add agents to the workspace
        workspace.add_agent(recipe_agent)
        workspace.add_agent(shopping_agent)
        
        # Start collaboration
        await workspace.start()
        
        # Submit a collaborative task
        result = await workspace.collaborate(
            "Plan a week of dinners and create a shopping list"
        )
        
        # Get artifacts
        artifacts = workspace.get_artifacts()
        
        # Clean up
        await workspace.close()
        ```
    """
    
    def __init__(
        self,
        name: str,
        description: str,
        workspace_id: str | None = None,
        workspace_dir: Path | None = None,
        auto_cleanup: bool = True,
        enable_persistence: bool = False,
        max_messages: int = 1000,
        message_retention_hours: int = 24,
    ):
        """
        Initialize a workspace.
        
        Args:
            name: Name of the workspace
            description: Description of the workspace purpose
            workspace_id: Optional custom workspace ID
            workspace_dir: Optional custom directory for workspace files
            auto_cleanup: Whether to automatically clean up on close
            enable_persistence: Whether to persist workspace state to disk
            max_messages: Maximum number of messages to retain
            message_retention_hours: Hours to retain messages
        """
        self.workspace_id = workspace_id or str(uuid.uuid4())
        self.name = name
        self.description = description
        self.auto_cleanup = auto_cleanup
        self.enable_persistence = enable_persistence
        self.max_messages = max_messages
        self.message_retention_hours = message_retention_hours
        
        # Create workspace directory
        if workspace_dir is None:
            workspace_dir = Path.home() / ".synqed" / "workspaces" / self.workspace_id
        self.workspace_dir = Path(workspace_dir)
        self.workspace_dir.mkdir(parents=True, exist_ok=True)
        
        # State
        self.state = WorkspaceState.CREATED
        self.created_at = datetime.now()
        self.started_at: datetime | None = None
        self.completed_at: datetime | None = None
        
        # Participants
        self._participants: dict[str, WorkspaceParticipant] = {}
        self._clients: dict[str, Client] = {}
        
        # Communication
        self._messages: list[WorkspaceMessage] = []
        self._message_callbacks: list[Callable[[WorkspaceMessage], Coroutine[Any, Any, None]]] = []
        
        # Artifacts
        self._artifacts: dict[str, WorkspaceArtifact] = {}
        
        # Shared state
        self._shared_state: dict[str, Any] = {}
        
        # Background task for message processing
        self._message_processor_task: asyncio.Task | None = None
        self._running = False
        
        logger.info(f"Created workspace: {self.name} ({self.workspace_id})")
    
    def add_agent(
        self,
        agent: Agent | None = None,
        agent_url: str | None = None,
        role: str = "agent",
        metadata: dict[str, Any] | None = None,
    ) -> str:
        """
        Add an agent to the workspace.
        
        Args:
            agent: Agent instance (if local)
            agent_url: URL of a remote agent
            role: Role of the agent in the workspace
            metadata: Additional metadata
            
        Returns:
            Participant ID
            
        Raises:
            ValueError: If neither agent nor agent_url is provided
        """
        if agent is None and agent_url is None:
            raise ValueError("Must provide either agent or agent_url")
        
        participant_id = str(uuid.uuid4())
        
        if agent is not None:
            name = agent.name
            agent_url = agent.url
        else:
            name = agent_url or "RemoteAgent"
        
        participant = WorkspaceParticipant(
            participant_id=participant_id,
            name=name,
            role=role,
            agent=agent,
            agent_url=agent_url,
            metadata=metadata or {},
        )
        
        self._participants[participant_id] = participant
        
        # Add system message
        self._add_message(
            sender_id="system",
            sender_name="System",
            message_type=MessageType.SYSTEM,
            content=f"{name} joined the workspace",
            metadata={"participant_id": participant_id},
        )
        
        logger.info(f"Added agent to workspace: {name}")
        return participant_id
    
    def remove_agent(self, participant_id: str) -> None:
        """
        Remove an agent from the workspace.
        
        Args:
            participant_id: ID of the participant to remove
        """
        if participant_id not in self._participants:
            logger.warning(f"Participant not found: {participant_id}")
            return
        
        participant = self._participants[participant_id]
        
        # Close client if exists
        if participant_id in self._clients:
            # Note: We don't await here, client will be closed during workspace shutdown
            del self._clients[participant_id]
        
        del self._participants[participant_id]
        
        # Add system message
        self._add_message(
            sender_id="system",
            sender_name="System",
            message_type=MessageType.SYSTEM,
            content=f"{participant.name} left the workspace",
        )
        
        logger.info(f"Removed agent from workspace: {participant.name}")
    
    def list_participants(self) -> list[dict[str, Any]]:
        """
        List all participants in the workspace.
        
        Returns:
            List of participant information dictionaries
        """
        return [
            {
                "participant_id": p.participant_id,
                "name": p.name,
                "role": p.role,
                "agent_url": p.agent_url,
                "joined_at": p.joined_at.isoformat(),
                "metadata": p.metadata,
            }
            for p in self._participants.values()
        ]
    
    async def start(self) -> None:
        """
        Start the workspace.
        
        This activates the workspace and begins processing messages.
        """
        if self.state == WorkspaceState.ACTIVE:
            logger.warning("Workspace is already active")
            return
        
        self.state = WorkspaceState.ACTIVE
        self.started_at = datetime.now()
        self._running = True
        
        # Start background message processor
        self._message_processor_task = asyncio.create_task(self._process_messages())
        
        # Add system message
        self._add_message(
            sender_id="system",
            sender_name="System",
            message_type=MessageType.SYSTEM,
            content=f"Workspace '{self.name}' started",
        )
        
        logger.info(f"Started workspace: {self.name}")
    
    async def pause(self) -> None:
        """Pause the workspace."""
        if self.state != WorkspaceState.ACTIVE:
            logger.warning("Workspace is not active")
            return
        
        self.state = WorkspaceState.PAUSED
        
        self._add_message(
            sender_id="system",
            sender_name="System",
            message_type=MessageType.SYSTEM,
            content="Workspace paused",
        )
        
        logger.info("Paused workspace")
    
    async def resume(self) -> None:
        """Resume the workspace."""
        if self.state != WorkspaceState.PAUSED:
            logger.warning("Workspace is not paused")
            return
        
        self.state = WorkspaceState.ACTIVE
        
        self._add_message(
            sender_id="system",
            sender_name="System",
            message_type=MessageType.SYSTEM,
            content="Workspace resumed",
        )
        
        logger.info("Resumed workspace")
    
    async def complete(self) -> None:
        """Mark the workspace as completed."""
        self.state = WorkspaceState.COMPLETED
        self.completed_at = datetime.now()
        
        self._add_message(
            sender_id="system",
            sender_name="System",
            message_type=MessageType.SYSTEM,
            content="Workspace completed",
        )
        
        logger.info("Completed workspace")
    
    async def collaborate(
        self,
        task: str,
        orchestrator: Any | None = None,
        timeout: float = 300.0,
    ) -> dict[str, Any]:
        """
        Execute a collaborative task across all agents in the workspace.
        
        Args:
            task: The task description
            orchestrator: Optional orchestrator for intelligent routing
            timeout: Timeout in seconds
            
        Returns:
            Dictionary containing results from all agents
        """
        if self.state != WorkspaceState.ACTIVE:
            raise ValueError("Workspace must be active to collaborate")
        
        if not self._participants:
            raise ValueError("No agents in workspace")
        
        logger.info(f"Starting collaborative task: {task[:100]}...")
        
        # Add task message
        self._add_message(
            sender_id="system",
            sender_name="System",
            message_type=MessageType.SYSTEM,
            content=f"Collaborative task: {task}",
        )
        
        # If orchestrator is provided, use it to select agents
        selected_participants = list(self._participants.values())
        
        if orchestrator is not None:
            try:
                # Register all agents with orchestrator
                for participant in self._participants.values():
                    if participant.agent is not None:
                        orchestrator.register_agent(
                            participant.agent.card,
                            participant.agent_url,
                            participant.participant_id,
                        )
                
                # Get orchestration decision
                result = await orchestrator.orchestrate(task)
                
                # Filter to selected agents
                selected_ids = {s.agent_id for s in result.selected_agents}
                selected_participants = [
                    p for p in self._participants.values()
                    if p.participant_id in selected_ids
                ]
                
                logger.info(f"Orchestrator selected {len(selected_participants)} agents")
            except Exception as e:
                logger.warning(f"Orchestrator failed, using all agents: {e}")
        
        # Execute task with selected agents
        results = {}
        tasks = []
        
        for participant in selected_participants:
            if participant.agent_url:
                client = await self._get_client(participant.participant_id)
                tasks.append(self._execute_agent_task(participant, client, task))
        
        # Wait for all tasks with timeout
        try:
            agent_results = await asyncio.wait_for(
                asyncio.gather(*tasks, return_exceptions=True),
                timeout=timeout,
            )
            
            # Collect results
            for participant, result in zip(selected_participants, agent_results):
                if isinstance(result, Exception):
                    results[participant.name] = f"Error: {str(result)}"
                    logger.error(f"Agent {participant.name} failed: {result}")
                else:
                    results[participant.name] = result
        
        except asyncio.TimeoutError:
            logger.error("Collaboration timed out")
            results["error"] = "Collaboration timed out"
        
        # Add results as artifact
        self.add_artifact(
            name=f"collaboration_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
            artifact_type="result",
            content=results,
            created_by="system",
            metadata={"task": task},
        )
        
        return results
    
    async def _execute_agent_task(
        self,
        participant: WorkspaceParticipant,
        client: Client,
        task: str,
    ) -> str:
        """Execute a task with an agent."""
        try:
            # Add agent task message
            self._add_message(
                sender_id=participant.participant_id,
                sender_name=participant.name,
                message_type=MessageType.STATUS,
                content=f"Processing task: {task[:100]}...",
            )
            
            # Execute task
            result = await client.ask(task)
            
            # Add result message
            self._add_message(
                sender_id=participant.participant_id,
                sender_name=participant.name,
                message_type=MessageType.AGENT,
                content=result,
                metadata={"task": task},
            )
            
            return result
        
        except Exception as e:
            # Add error message
            self._add_message(
                sender_id=participant.participant_id,
                sender_name=participant.name,
                message_type=MessageType.STATUS,
                content=f"Error: {str(e)}",
            )
            raise
    
    async def send_message_to_agent(
        self,
        participant_id: str,
        message: str,
    ) -> str:
        """
        Send a direct message to a specific agent.
        
        Args:
            participant_id: ID of the participant to message
            message: Message content
            
        Returns:
            Response from the agent
        """
        if participant_id not in self._participants:
            raise ValueError(f"Participant not found: {participant_id}")
        
        participant = self._participants[participant_id]
        
        if not participant.agent_url:
            raise ValueError(f"Agent URL not available for {participant.name}")
        
        # Get client
        client = await self._get_client(participant_id)
        
        # Send message
        self._add_message(
            sender_id="system",
            sender_name="System",
            message_type=MessageType.USER,
            content=message,
            metadata={"to": participant.name},
        )
        
        response = await client.ask(message)
        
        # Add response
        self._add_message(
            sender_id=participant.participant_id,
            sender_name=participant.name,
            message_type=MessageType.AGENT,
            content=response,
        )
        
        return response
    
    async def broadcast_message(
        self,
        message: str,
        sender_id: str = "system",
        sender_name: str = "System",
    ) -> dict[str, str]:
        """
        Broadcast a message to all agents in the workspace.
        
        Args:
            message: Message to broadcast
            sender_id: ID of the sender
            sender_name: Name of the sender
            
        Returns:
            Dictionary of responses from all agents
        """
        self._add_message(
            sender_id=sender_id,
            sender_name=sender_name,
            message_type=MessageType.SYSTEM,
            content=f"Broadcast: {message}",
        )
        
        responses = {}
        tasks = []
        
        for participant in self._participants.values():
            if participant.agent_url:
                client = await self._get_client(participant.participant_id)
                tasks.append(self._execute_agent_task(participant, client, message))
        
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        for participant, result in zip(self._participants.values(), results):
            if isinstance(result, Exception):
                responses[participant.name] = f"Error: {str(result)}"
            else:
                responses[participant.name] = result
        
        return responses
    
    def add_artifact(
        self,
        name: str,
        artifact_type: str,
        content: str | bytes | dict[str, Any],
        created_by: str,
        metadata: dict[str, Any] | None = None,
    ) -> str:
        """
        Add an artifact to the workspace.
        
        Args:
            name: Name of the artifact
            artifact_type: Type (file, data, result, etc.)
            content: Artifact content
            created_by: ID of creator
            metadata: Additional metadata
            
        Returns:
            Artifact ID
        """
        artifact_id = str(uuid.uuid4())
        
        artifact = WorkspaceArtifact(
            artifact_id=artifact_id,
            name=name,
            artifact_type=artifact_type,
            content=content,
            created_by=created_by,
            created_at=datetime.now(),
            metadata=metadata or {},
        )
        
        self._artifacts[artifact_id] = artifact
        
        # Save artifact to disk if it's a file
        if artifact_type == "file" and isinstance(content, (str, bytes)):
            artifact_path = self.workspace_dir / "artifacts" / name
            artifact_path.parent.mkdir(parents=True, exist_ok=True)
            
            if isinstance(content, bytes):
                artifact_path.write_bytes(content)
            else:
                artifact_path.write_text(content)
        
        # Add message
        self._add_message(
            sender_id=created_by,
            sender_name=self._get_participant_name(created_by),
            message_type=MessageType.ARTIFACT,
            content=f"Created artifact: {name}",
            metadata={"artifact_id": artifact_id, "artifact_type": artifact_type},
        )
        
        logger.info(f"Added artifact: {name}")
        return artifact_id
    
    def get_artifact(self, artifact_id: str) -> WorkspaceArtifact | None:
        """
        Get an artifact by ID.
        
        Args:
            artifact_id: ID of the artifact
            
        Returns:
            WorkspaceArtifact or None if not found
        """
        return self._artifacts.get(artifact_id)
    
    def get_artifacts(
        self,
        artifact_type: str | None = None,
        created_by: str | None = None,
    ) -> list[WorkspaceArtifact]:
        """
        Get artifacts, optionally filtered.
        
        Args:
            artifact_type: Optional filter by type
            created_by: Optional filter by creator
            
        Returns:
            List of artifacts
        """
        artifacts = list(self._artifacts.values())
        
        if artifact_type is not None:
            artifacts = [a for a in artifacts if a.artifact_type == artifact_type]
        
        if created_by is not None:
            artifacts = [a for a in artifacts if a.created_by == created_by]
        
        return artifacts
    
    def set_shared_state(self, key: str, value: Any) -> None:
        """
        Set a shared state value.
        
        Args:
            key: State key
            value: State value
        """
        self._shared_state[key] = value
        logger.debug(f"Set shared state: {key}")
    
    def get_shared_state(self, key: str, default: Any = None) -> Any:
        """
        Get a shared state value.
        
        Args:
            key: State key
            default: Default value if not found
            
        Returns:
            State value or default
        """
        return self._shared_state.get(key, default)
    
    def get_all_shared_state(self) -> dict[str, Any]:
        """
        Get all shared state.
        
        Returns:
            Dictionary of all shared state
        """
        return self._shared_state.copy()
    
    def get_messages(
        self,
        message_type: MessageType | None = None,
        sender_id: str | None = None,
        limit: int | None = None,
    ) -> list[WorkspaceMessage]:
        """
        Get workspace messages, optionally filtered.
        
        Args:
            message_type: Optional filter by type
            sender_id: Optional filter by sender
            limit: Optional limit on number of messages
            
        Returns:
            List of messages
        """
        messages = self._messages.copy()
        
        if message_type is not None:
            messages = [m for m in messages if m.message_type == message_type]
        
        if sender_id is not None:
            messages = [m for m in messages if m.sender_id == sender_id]
        
        if limit is not None:
            messages = messages[-limit:]
        
        return messages
    
    def on_message(
        self,
        callback: Callable[[WorkspaceMessage], Coroutine[Any, Any, None]],
    ) -> None:
        """
        Register a callback for new messages.
        
        Args:
            callback: Async function to call on new messages
        """
        self._message_callbacks.append(callback)
    
    async def export_workspace(self, export_path: Path | None = None) -> Path:
        """
        Export the workspace state to a JSON file.
        
        Args:
            export_path: Optional custom export path
            
        Returns:
            Path to the exported file
        """
        if export_path is None:
            export_path = self.workspace_dir / f"workspace_export_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        
        export_data = {
            "workspace_id": self.workspace_id,
            "name": self.name,
            "description": self.description,
            "state": self.state.value,
            "created_at": self.created_at.isoformat(),
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
            "participants": self.list_participants(),
            "messages": [m.to_dict() for m in self._messages],
            "artifacts": [a.to_dict() for a in self._artifacts.values()],
            "shared_state": self._shared_state,
        }
        
        export_path.parent.mkdir(parents=True, exist_ok=True)
        export_path.write_text(json.dumps(export_data, indent=2))
        
        logger.info(f"Exported workspace to: {export_path}")
        return export_path
    
    async def close(self) -> None:
        """Close the workspace and clean up resources."""
        self._running = False
        
        # Cancel message processor
        if self._message_processor_task and not self._message_processor_task.done():
            self._message_processor_task.cancel()
            try:
                await self._message_processor_task
            except asyncio.CancelledError:
                pass
        
        # Close all clients
        for client in self._clients.values():
            await client.close()
        self._clients.clear()
        
        # Export workspace if persistence is enabled
        if self.enable_persistence:
            await self.export_workspace()
        
        # Clean up workspace directory if auto_cleanup is enabled
        if self.auto_cleanup and self.workspace_dir.exists():
            shutil.rmtree(self.workspace_dir)
            logger.info(f"Cleaned up workspace directory: {self.workspace_dir}")
        
        logger.info(f"Closed workspace: {self.name}")
    
    async def __aenter__(self) -> "Workspace":
        """Async context manager entry."""
        await self.start()
        return self
    
    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """Async context manager exit."""
        if exc_type is not None:
            self.state = WorkspaceState.ERROR
        await self.close()
    
    # Private methods
    
    def _add_message(
        self,
        sender_id: str,
        sender_name: str,
        message_type: MessageType,
        content: str,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        """Add a message to the workspace."""
        message = WorkspaceMessage(
            message_id=str(uuid.uuid4()),
            timestamp=datetime.now(),
            sender_id=sender_id,
            sender_name=sender_name,
            message_type=message_type,
            content=content,
            metadata=metadata or {},
        )
        
        self._messages.append(message)
        
        # Enforce message limit
        if len(self._messages) > self.max_messages:
            self._messages = self._messages[-self.max_messages:]
        
        # Trigger callbacks (don't await, run in background)
        for callback in self._message_callbacks:
            asyncio.create_task(callback(message))
    
    async def _get_client(self, participant_id: str) -> Client:
        """Get or create a client for a participant."""
        if participant_id not in self._clients:
            participant = self._participants[participant_id]
            
            if not participant.agent_url:
                raise ValueError(f"No agent URL for participant: {participant_id}")
            
            self._clients[participant_id] = Client(
                agent_url=participant.agent_url,
                streaming=True,
            )
        
        return self._clients[participant_id]
    
    def _get_participant_name(self, participant_id: str) -> str:
        """Get participant name by ID."""
        if participant_id == "system":
            return "System"
        
        participant = self._participants.get(participant_id)
        return participant.name if participant else "Unknown"
    
    async def _process_messages(self) -> None:
        """Background task to process messages."""
        while self._running:
            await asyncio.sleep(1)
            
            # Here you could add message processing logic
            # For now, this is just a placeholder for future enhancements
    
    def __repr__(self) -> str:
        """String representation."""
        return (
            f"Workspace(name='{self.name}', "
            f"id='{self.workspace_id}', "
            f"state='{self.state.value}', "
            f"participants={len(self._participants)})"
        )

