import os
import asyncio
from typing import List, Tuple, Optional, Dict, Union, Any
from pathlib import Path
import json
from datetime import datetime, UTC
from pydantic import BaseModel, Field, ConfigDict

from .agent import Agent, RAGAgent
from .rag import RAGSystem, Document
from .thread import Thread, ThreadManager, Message
from .file import FileManager
from .interpreter import CodeInterpreter
from .flow import CommunicationFlow

class AgencyConfig(BaseModel):
    """
    Configuration for an agency.
    
    Attributes:
        name: Name of the agency
        description: Description of the agency's purpose
        storage_path: Path for persistent storage
        shared_instructions: Path to shared instructions file
        default_temperature: Default temperature for agents
        default_max_tokens: Default max tokens for agents
        use_code_interpreter: Whether to enable code interpreter
        use_rag: Whether to enable RAG capabilities
        use_file_storage: Whether to enable file storage
    """
    name: str
    description: str
    storage_path: Optional[str] = None
    shared_instructions: Optional[str] = None
    default_temperature: float = 0.7
    default_max_tokens: int = 4096
    use_code_interpreter: bool = False
    use_rag: bool = False
    use_file_storage: bool = False

class Agency(BaseModel):
    """
    Manages a collection of agents and their communication flows.
    
    The Agency class is responsible for:
    1. Managing agent initialization and configuration
    2. Coordinating communication between agents
    3. Maintaining shared knowledge and resources
    4. Monitoring and logging agent activities
    """
    model_config = ConfigDict(arbitrary_types_allowed=True)
    
    name: str = Field(..., description="Name of the agency")
    description: str = Field(..., description="Description of what the agency does")
    agents: Dict[str, Agent] = Field(default_factory=dict)
    flows: List[CommunicationFlow] = Field(default_factory=list)
    shared_instructions: Optional[str] = None
    storage_path: Optional[str] = None
    shared_rag: Optional[RAGSystem] = None
    activity_log: List[Dict[str, Any]] = Field(default_factory=list)
    
    def __init__(self, **data):
        super().__init__(**data)
        self._initialize_agency()
        
    def _initialize_agency(self):
        """Initialize agency components."""
        # Set up storage path
        if not self.storage_path:
            self.storage_path = str(Path.home() / ".multi-swarm" / "agencies" / self.name)
            
        # Initialize shared RAG system
        self.shared_rag = RAGSystem(
            storage_path=Path(self.storage_path) / "shared_knowledge"
        )
        
        # Load shared instructions if provided
        if self.shared_instructions and os.path.exists(self.shared_instructions):
            with open(self.shared_instructions, "r") as f:
                self.shared_instructions = f.read()
                
    def add_agent(
        self,
        agent: Agent,
        flows: Optional[List[Tuple[str, str]]] = None
    ):
        """Add an agent to the agency."""
        self.agents[agent.name] = agent
        
        # Add communication flows if provided
        if flows:
            for source, target in flows:
                self.flows.append(
                    CommunicationFlow(
                        source=source,
                        target=target,
                        description=f"Flow from {source} to {target}"
                    )
                )
                
        self._log_activity(
            "agent_added",
            {"agent_name": agent.name, "flows_added": len(flows) if flows else 0}
        )
        
    def remove_agent(self, agent_name: str):
        """Remove an agent from the agency."""
        if agent_name in self.agents:
            del self.agents[agent_name]
            # Remove associated flows
            self.flows = [
                f for f in self.flows
                if f.source != agent_name and f.target != agent_name
            ]
            self._log_activity("agent_removed", {"agent_name": agent_name})
            
    def get_agent(self, agent_name: str) -> Optional[Agent]:
        """Get an agent by name."""
        return self.agents.get(agent_name)
        
    def list_agents(self) -> List[Dict[str, Any]]:
        """List all agents and their basic information."""
        return [
            {
                "name": name,
                "description": agent.description,
                "capabilities": [
                    "rag" if agent.use_rag else None,
                    "code_interpreter" if agent.use_code_interpreter else None,
                    "file_storage" if agent.use_file_storage else None
                ]
            }
            for name, agent in self.agents.items()
        ]
        
    def can_communicate(self, source: str, target: str) -> bool:
        """Check if two agents can communicate."""
        return any(
            f.source == source and f.target == target and f.enabled
            for f in self.flows
        )
        
    def process_message(
        self,
        message: str,
        agent_name: str,
        thread: Optional[Thread] = None
    ) -> Tuple[str, Thread]:
        """Process a message through an agent."""
        agent = self.get_agent(agent_name)
        if not agent:
            raise ValueError(f"Agent not found: {agent_name}")
            
        # Create or use existing thread
        if thread is None:
            thread = Thread()
            
        # Add message to thread
        thread.add_message(Message(role="user", content=message))
        
        # Process with agent
        response = agent.process_message(thread)
        thread.add_message(Message(role="assistant", content=response))
        
        self._log_activity(
            "message_processed",
            {
                "agent": agent_name,
                "thread_id": thread.id,
                "message_count": len(thread.messages)
            }
        )
        
        return response, thread
        
    def delegate_message(
        self,
        message: str,
        source_agent: str,
        target_agent: str,
        thread: Optional[Thread] = None
    ) -> Tuple[str, Thread]:
        """Delegate a message from one agent to another."""
        if not self.can_communicate(source_agent, target_agent):
            raise ValueError(
                f"Communication not allowed from {source_agent} to {target_agent}"
            )
            
        return self.process_message(message, target_agent, thread)
        
    def add_shared_knowledge(
        self,
        content: str,
        metadata: Optional[Dict] = None
    ) -> Document:
        """Add knowledge to the shared RAG system."""
        if metadata is None:
            metadata = {}
        metadata["added_at"] = datetime.now(UTC).isoformat()
        
        document = self.shared_rag.add_document(content, metadata)
        self._log_activity(
            "knowledge_added",
            {"document_id": document.id, "metadata": metadata}
        )
        return document
        
    def search_shared_knowledge(
        self,
        query: str,
        k: int = 5
    ) -> List[Document]:
        """Search the shared knowledge base."""
        results = self.shared_rag.search(query, k=k)
        self._log_activity(
            "knowledge_searched",
            {"query": query, "results_count": len(results)}
        )
        return results
        
    def _log_activity(self, activity_type: str, details: Dict[str, Any]):
        """Log agency activity."""
        self.activity_log.append({
            "timestamp": datetime.now(UTC).isoformat(),
            "type": activity_type,
            "details": details
        })
        
    def get_activity_log(
        self,
        activity_type: Optional[str] = None,
        start_time: Optional[datetime] = None,
        end_time: Optional[datetime] = None
    ) -> List[Dict[str, Any]]:
        """Get filtered activity log."""
        filtered_log = self.activity_log
        
        if activity_type:
            filtered_log = [
                entry for entry in filtered_log
                if entry["type"] == activity_type
            ]
            
        if start_time:
            filtered_log = [
                entry for entry in filtered_log
                if datetime.fromisoformat(entry["timestamp"]) >= start_time
            ]
            
        if end_time:
            filtered_log = [
                entry for entry in filtered_log
                if datetime.fromisoformat(entry["timestamp"]) <= end_time
            ]
            
        return filtered_log
        
    def save_state(self):
        """Save agency state to disk."""
        if not self.storage_path:
            return
            
        state_path = Path(self.storage_path) / "agency_state.json"
        state = {
            "name": self.name,
            "description": self.description,
            "flows": [flow.model_dump() for flow in self.flows],
            "activity_log": self.activity_log
        }
        
        os.makedirs(os.path.dirname(state_path), exist_ok=True)
        with open(state_path, "w") as f:
            json.dump(state, f)
            
    def load_state(self):
        """Load agency state from disk."""
        if not self.storage_path:
            return
            
        state_path = Path(self.storage_path) / "agency_state.json"
        if not state_path.exists():
            return
            
        with open(state_path, "r") as f:
            state = json.load(f)
            
        self.name = state["name"]
        self.description = state["description"]
        self.flows = [
            CommunicationFlow(**flow_data)
            for flow_data in state["flows"]
        ]
        self.activity_log = state["activity_log"]
    
    def run_demo(self):
        """Print welcome message and instructions for the demo."""
        entry_agent = self.get_agent(self.config.name)
        print(f"\nWelcome to the {entry_agent.config.name} Agency Demo!")
        print("\nType 'exit' to end the conversation.")
        print("\nEnter your message:")
        self.demo_loop()
    
    def demo_loop(self):
        """Run the interactive demo loop."""
        try:
            while True:
                user_input = input("> ")
                if user_input.lower() == 'exit':
                    print("\nThank you for using the agency. Goodbye!")
                    break
                
                response = self.process_message(user_input)
                entry_agent = self.get_agent(self.config.name)
                print(f"\n{entry_agent.config.name}: {response}")
        
        except KeyboardInterrupt:
            print("\n\nDemo terminated by user. Goodbye!")
        except Exception as e:
            print(f"\nAn error occurred: {str(e)}")
    
    def process_message(self, message: str) -> str:
        """Process a message through the agency."""
        entry_agent = self.get_agent(self.config.name)
        if not entry_agent:
            return None
        
        # Create a thread for this message
        thread = self.thread_manager.create_thread()
        thread.add_message({
            "role": "user",
            "content": message,
            "agent_name": "user"
        })
        
        # Process with the entry agent
        response = entry_agent._process_with_llm(thread)
        return response 
    
    def run(self, prompt: str) -> str:
        """
        Process a prompt through the agency's workflow.
        
        Args:
            prompt: The input prompt to process
            
        Returns:
            The final response from the agency
        """
        if not prompt:
            return {"error": "Empty prompt provided"}
            
        try:
            # Get the entry point agent (first agent in the list)
            entry_agent = list(self.agents.values())[0]
            
            # Get the first flow's thread
            first_flow = next(iter(self.flows.values()))
            thread = self.thread_manager.get_thread(first_flow.thread_id)
            
            # Add the initial prompt to the thread
            thread.add_message(
                role="user",
                content=prompt,
                agent_name="user"
            )
            
            # Process with entry agent
            response = entry_agent._process_with_llm(thread)
            thread.add_message(
                role="assistant",
                content=response,
                agent_name=entry_agent.config.name
            )
            
            # Process through connected agents based on the response
            for source, target in self.flows.keys():
                if source == entry_agent.config.name:
                    target_agent = self.agents[target]
                    target_response = target_agent._process_with_llm(thread)
                    thread.add_message(
                        role="assistant",
                        content=target_response,
                        agent_name=target
                    )
            
            # Return the combined response
            return self._combine_responses(thread)
            
        except Exception as e:
            return {"error": f"Error processing prompt: {str(e)}"}
    
    def _combine_responses(self, thread: Thread) -> Dict:
        """Combine responses from all agents into a final response."""
        combined = {}
        
        for message in thread.messages:
            if message.role == "assistant":
                try:
                    # Try to parse as JSON
                    response_data = json.loads(message.content)
                    combined.update(response_data)
                except:
                    # If not JSON, add as string
                    combined[f"{message.metadata['agent_name']}_response"] = message.content
        
        return combined 