"""
Command-line interface for Multi-Swarm Framework.
"""

import click
import os
from pathlib import Path
import json
from typing import Dict, Any, Optional
from datetime import datetime, UTC

from .core.agent import Agent
from .core.agency import Agency, CommunicationFlow
from .core.rag import RAGSystem

def load_config(config_path: str) -> Dict[str, Any]:
    """Load configuration from a JSON file."""
    with open(config_path, "r") as f:
        return json.load(f)

def create_agent_from_config(config: Dict[str, Any]) -> Agent:
    """Create an agent from configuration."""
    return Agent(**config)

def create_agency_from_config(config: Dict[str, Any]) -> Agency:
    """Create an agency from configuration."""
    # Create agents
    agents = {}
    for agent_config in config.get("agents", []):
        agent = create_agent_from_config(agent_config)
        agents[agent.name] = agent
        
    # Create flows
    flows = []
    for flow_config in config.get("flows", []):
        flows.append(CommunicationFlow(**flow_config))
        
    # Create agency
    return Agency(
        name=config["name"],
        description=config["description"],
        agents=agents,
        flows=flows,
        shared_instructions=config.get("shared_instructions"),
        storage_path=config.get("storage_path")
    )

@click.group()
def cli():
    """Multi-Swarm Framework CLI."""
    pass

@cli.command()
@click.argument("name")
@click.option("--description", "-d", help="Description of the agent")
@click.option("--use-rag/--no-rag", default=True, help="Enable/disable RAG")
@click.option("--use-code/--no-code", default=False, help="Enable/disable code interpreter")
@click.option("--use-files/--no-files", default=False, help="Enable/disable file storage")
@click.option("--llm", default="claude", help="LLM provider (claude/gemini)")
@click.option("--model", help="Model name for the LLM provider")
@click.option("--temperature", default=0.7, help="Temperature for LLM generation")
@click.option("--output", "-o", help="Output path for agent configuration")
def create_agent(
    name: str,
    description: str,
    use_rag: bool,
    use_code: bool,
    use_files: bool,
    llm: str,
    model: Optional[str],
    temperature: float,
    output: Optional[str]
):
    """Create a new agent."""
    # Set default model based on provider
    if not model:
        model = "claude-3-sonnet" if llm == "claude" else "gemini-pro"
        
    # Create agent configuration
    config = {
        "name": name,
        "description": description,
        "use_rag": use_rag,
        "use_code_interpreter": use_code,
        "use_file_storage": use_files,
        "llm_provider": llm,
        "provider_config": {
            "model": model,
            "api_version": "2024-03" if llm == "claude" else "2024-01",
            "max_tokens": 4096
        },
        "temperature": temperature
    }
    
    # Create agent
    agent = Agent(**config)
    
    # Save configuration if output path provided
    if output:
        os.makedirs(os.path.dirname(output), exist_ok=True)
        with open(output, "w") as f:
            json.dump(config, f, indent=2)
            
    click.echo(f"Created agent: {name}")
    click.echo(f"Configuration saved to: {output}" if output else "")

@cli.command()
@click.argument("name")
@click.option("--description", "-d", help="Description of the agency")
@click.option("--config", "-c", help="Path to agency configuration file")
@click.option("--instructions", "-i", help="Path to shared instructions file")
@click.option("--storage", "-s", help="Storage path for agency data")
@click.option("--output", "-o", help="Output path for agency configuration")
def create_agency(
    name: str,
    description: Optional[str],
    config: Optional[str],
    instructions: Optional[str],
    storage: Optional[str],
    output: Optional[str]
):
    """Create a new agency."""
    if config:
        # Load from configuration file
        agency_config = load_config(config)
    else:
        # Create minimal configuration
        agency_config = {
            "name": name,
            "description": description or f"Agency: {name}",
            "agents": [],
            "flows": [],
            "shared_instructions": instructions,
            "storage_path": storage
        }
        
    # Create agency
    agency = create_agency_from_config(agency_config)
    
    # Save configuration if output path provided
    if output:
        os.makedirs(os.path.dirname(output), exist_ok=True)
        with open(output, "w") as f:
            json.dump(agency_config, f, indent=2)
            
    click.echo(f"Created agency: {name}")
    click.echo(f"Configuration saved to: {output}" if output else "")

@cli.command()
@click.argument("agency_config")
@click.option("--agent", "-a", help="Agent to process the message")
@click.option("--message", "-m", help="Message to process")
@click.option("--interactive/--no-interactive", default=True, help="Enable interactive mode")
def chat(
    agency_config: str,
    agent: Optional[str],
    message: Optional[str],
    interactive: bool
):
    """Chat with an agency or specific agent."""
    # Load agency
    config = load_config(agency_config)
    agency = create_agency_from_config(config)
    
    if not interactive:
        # Single message mode
        if not agent or not message:
            raise click.UsageError("Agent and message required in non-interactive mode")
            
        response, _ = agency.process_message(message, agent)
        click.echo(response)
    else:
        # Interactive mode
        click.echo("Multi-Swarm Interactive Chat")
        click.echo("Type 'exit' to quit, 'agents' to list available agents")
        click.echo("Format: @agent_name message")
        
        while True:
            try:
                user_input = click.prompt("You")
                
                if user_input.lower() == "exit":
                    break
                    
                if user_input.lower() == "agents":
                    agents = agency.list_agents()
                    for agent_info in agents:
                        click.echo(
                            f"{agent_info['name']}: {agent_info['description']}"
                        )
                    continue
                    
                if not user_input.startswith("@"):
                    click.echo("Please specify an agent using @agent_name")
                    continue
                    
                # Parse agent and message
                parts = user_input[1:].split(" ", 1)
                if len(parts) != 2:
                    click.echo("Invalid format. Use: @agent_name message")
                    continue
                    
                agent_name, message = parts
                
                # Process message
                response, _ = agency.process_message(message, agent_name)
                click.echo(f"\n{agent_name}: {response}\n")
                
            except KeyboardInterrupt:
                break
            except Exception as e:
                click.echo(f"Error: {str(e)}")

@cli.command()
@click.argument("agency_config")
@click.option("--query", "-q", help="Search query")
@click.option("--k", default=5, help="Number of results to return")
@click.option("--agent", "-a", help="Specific agent to search")
def search(
    agency_config: str,
    query: str,
    k: int,
    agent: Optional[str]
):
    """Search agency knowledge base."""
    # Load agency
    config = load_config(agency_config)
    agency = create_agency_from_config(config)
    
    if agent:
        # Search specific agent
        agent_obj = agency.get_agent(agent)
        if not agent_obj:
            raise click.UsageError(f"Agent not found: {agent}")
            
        results = agent_obj.search_knowledge(query, k=k)
    else:
        # Search shared knowledge
        results = agency.search_shared_knowledge(query, k=k)
        
    # Display results
    click.echo(f"\nSearch results for: {query}\n")
    for i, result in enumerate(results, 1):
        click.echo(f"{i}. Score: {result.score:.4f}")
        click.echo(f"Content: {result.document.content[:200]}...")
        click.echo(f"Created: {result.document.created_at}")
        if result.reflection:
            click.echo(f"Reflection: {result.reflection}")
        click.echo()

@cli.command()
@click.argument("agency_config")
@click.option("--content", "-c", help="Content to add")
@click.option("--file", "-f", help="File containing content to add")
@click.option("--agent", "-a", help="Specific agent to add knowledge to")
def add_knowledge(
    agency_config: str,
    content: Optional[str],
    file: Optional[str],
    agent: Optional[str]
):
    """Add knowledge to agency or specific agent."""
    if not content and not file:
        raise click.UsageError("Either content or file must be provided")
        
    # Load content from file if specified
    if file:
        with open(file, "r") as f:
            content = f.read()
            
    # Load agency
    config = load_config(agency_config)
    agency = create_agency_from_config(config)
    
    # Add knowledge
    metadata = {
        "added_at": datetime.now(UTC).isoformat(),
        "source": file if file else "direct"
    }
    
    if agent:
        # Add to specific agent
        agent_obj = agency.get_agent(agent)
        if not agent_obj:
            raise click.UsageError(f"Agent not found: {agent}")
            
        document = agent_obj.add_knowledge(content, metadata)
        click.echo(f"Added knowledge to agent: {agent}")
    else:
        # Add to shared knowledge
        document = agency.add_shared_knowledge(content, metadata)
        click.echo("Added knowledge to agency")
        
    click.echo(f"Document ID: {document.id}")

if __name__ == "__main__":
    cli() 