#!/usr/bin/env python3
"""
Pard0x Multi-Agent System Administrator Assistant
A supervisor-based system that coordinates specialized agents for system administration tasks.
"""

import sys
import warnings
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langgraph_supervisor import create_supervisor
from agents import (
    create_os_detector_worker,
    create_logs_analyzer_worker,
    create_performance_analyzer_worker
)
from custom_tools import print_poem

# Suppress the shell tool warning since worker agents use it intentionally for sysadmin tasks
warnings.filterwarnings("ignore", message="The shell tool has no safeguards by default. Use at your own risk.")


def print_welcome():
    """Print welcome message and system capabilities."""
    print("\n" + "="*80)
    print("🤖 Welcome to Pard0x Multi-Agent System Administrator Assistant!")
    print("="*80)
    print("\nI coordinate a team of specialized agents to help you with system administration tasks:")
    print("  • 🖥️ OS Detector - System identification and environment analysis")
    print("  • 📊 Logs Analyzer - Log investigation and error diagnosis")
    print("  • ⚡ Performance Analyzer - Resource monitoring and optimization")
    print("  • 🎭 Morale Booster - Motivational poems for tough debugging sessions!")
    print("\n" + "-"*80)


def print_examples():
    """Print example queries."""
    print("\n💡 Example queries you can try:")
    print("  - 'What operating system is this server running?'")
    print("  - 'Check the system logs for any errors in the last hour'")
    print("  - 'Analyze current system performance and identify bottlenecks'")
    print("  - 'My web server is down, help me troubleshoot'")
    print("  - 'Write me a motivational poem about debugging'")
    print("\n" + "-"*80)


def create_sysadmin_supervisor():
    """Create the main supervisor that coordinates between specialized agents."""
    
    # Get the base model
    model = ChatOpenAI(model="gpt-4o-mini", temperature=0)
    
    # Create specialized workers
    os_detector = create_os_detector_worker()
    logs_analyzer = create_logs_analyzer_worker()
    performance_analyzer = create_performance_analyzer_worker()
    
    # Create the supervisor with our agents
    supervisor = create_supervisor(
        agents=[os_detector, logs_analyzer, performance_analyzer],
        model=model,
        prompt="""You are Pard0x, an expert System Administrator Supervisor coordinating a team of specialized agents.

Your team consists of:
1. **OS Detector**: Identifies system information, environment, and configuration
2. **Logs Analyzer**: Investigates system and application logs for issues  
3. **Performance Analyzer**: Monitors and diagnoses performance problems

Your role:
1. **Task Analysis**: Understand the user's request and determine which agent(s) to engage
2. **Coordination**: Delegate tasks to appropriate agents based on their specialties
3. **Synthesis**: Combine insights from multiple agents into coherent solutions
4. **Direct Action**: Handle simple tasks yourself without delegation
5. **Morale Boost**: Use the poem tool to encourage users during tough debugging sessions

Decision guidelines:
- For system identification or environment questions → OS Detector
- For error investigation or log analysis → Logs Analyzer
- For performance issues or resource problems → Performance Analyzer
- For complex issues, engage multiple agents in sequence
- For simple queries or when agents aren't needed, respond directly

Communication style:
- Be professional yet approachable
- Provide clear explanations of your delegation decisions
- Synthesize agent findings into actionable recommendations
- Add a touch of humor when appropriate (especially with poems!)

Remember: Your goal is to solve system problems efficiently by leveraging your team's specialized skills while maintaining a positive debugging experience!

IMPORTANT: You do NOT have direct access to shell commands or SSH. You must delegate all system tasks to your specialized agents.""",
        tools=[print_poem]  # Supervisor only has poem tool - no shell/SSH access
    )
    
    return supervisor.compile()


def process_query(app, query: str, conversation_history: list) -> None:
    """Process a user query through the supervisor system with conversation history."""
    print(f"\n🔄 Processing your request: '{query}'")
    print("-" * 80)
    
    # Convert conversation history to LangChain message format
    messages = []
    for msg in conversation_history:
        if msg["role"] == "user":
            messages.append(HumanMessage(content=msg["content"]))
        else:  # assistant
            # Create an AI message - LangGraph will handle this format
            messages.append({"role": "assistant", "content": msg["content"]})
    
    # Add the new user message
    messages.append(HumanMessage(content=query))
    
    # Stream the response
    collected_responses = []
    
    for chunk in app.stream(
        {"messages": messages},
        stream_mode="values"
    ):
        chunk["messages"][-1].pretty_print()
        # Collect AI responses
        if chunk["messages"] and hasattr(chunk["messages"][-1], 'type') and chunk["messages"][-1].type == "ai":
            collected_responses.append(chunk["messages"][-1].content)
    
    # Add both user and assistant messages to history
    conversation_history.append({"role": "user", "content": query})
    if collected_responses:
        # Use the last response (most complete)
        conversation_history.append({"role": "assistant", "content": collected_responses[-1]})


def main():
    """Main interaction loop with conversation history."""
    print_welcome()
    print_examples()
    
    # Create the supervisor system
    print("\n🚀 Initializing the multi-agent system...")
    try:
        app = create_sysadmin_supervisor()
        print("✅ System ready!\n")
    except Exception as e:
        print(f"❌ Failed to initialize system: {str(e)}")
        sys.exit(1)
    
    # Initialize conversation history
    conversation_history = []
    
    # Interactive loop
    print("💬 Enter your queries below (type 'exit' to quit, 'help' for examples):\n")
    
    while True:
        try:
            query = input("You: ").strip()
            
            if not query:
                continue
                
            if query.lower() in ['exit', 'quit', 'q']:
                print("\n👋 Thanks for using Pard0x! Stay curious and keep debugging!")
                break
                
            if query.lower() in ['help', 'h', '?']:
                print_examples()
                continue
                
            if query.lower() in ['history', 'show history']:
                print("\n📜 Conversation History:")
                print("-" * 40)
                for i, msg in enumerate(conversation_history):
                    role = "You" if msg["role"] == "user" else "Assistant"
                    print(f"{i+1}. {role}: {msg['content'][:100]}{'...' if len(msg['content']) > 100 else ''}")
                print("-" * 40)
                continue
            
            process_query(app, query, conversation_history)
            
        except KeyboardInterrupt:
            print("\n\n👋 Goodbye! Keep those systems running smoothly!")
            break
        except Exception as e:
            print(f"\n❌ Unexpected error: {str(e)}")
            print("Please try again with a different query.")


if __name__ == "__main__":
    main()