#!/usr/bin/env python3
"""
MaxClaw Agent - A configurable agent harness using DeepAgents
Supports multiple LLM providers via LangChain's init_chat_model
Enhanced with unified provider tools and custom tools
"""
import os
import sys
from typing import Optional, List
from deepagents import create_deep_agent
from langchain.chat_models import init_chat_model
from langchain_core.messages import HumanMessage, SystemMessage
from langchain.tools import Tool

# Import custom tools
try:
    from tools import get_proxmox_tools
    from tools.provider_tools import get_provider_tools, get_all_provider_tools
    TOOLS_AVAILABLE = True
except ImportError:
    TOOLS_AVAILABLE = False
    print("Warning: Custom tools not available. Install required dependencies.")


def load_model_from_config():
    """
    Load the LLM model based on configuration.
    Configuration can come from environment variables or a config file.
    For simplicity, we'll use environment variables.
    
    Expected environment variables:
    - OPENAI_API_KEY: For OpenAI models
    - OLLAMA_BASE_URL: For Ollama (e.g., https://ollama.solarpi.net)
    - MODEL_PROVIDER: One of "openai", "ollama", "openrouter", etc.
    - MODEL_NAME: The specific model name (e.g., "gpt-4o", "gemma2:9b")
    """
    provider = os.getenv("MODEL_PROVIDER", "openai").lower()
    model_name = os.getenv("MODEL_NAME", "gpt-4o")
    
    # Configure based on provider
    if provider == "openai":
        # Uses OPENAI_API_KEY from environment
        return init_chat_model(f"openai:{model_name}")
    elif provider == "ollama":
        # For Ollama, we need to specify the base URL
        base_url = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
        return init_chat_model(f"ollama:{model_name}", base_url=base_url)
    elif provider == "openrouter":
        # For OpenRouter, we need OPENROUTER_API_KEY and optionally base URL
        # Note: LangChain's init_chat_model for openrouter might require specific setup
        # As of now, we assume it's supported via openai-compatible API
        base_url = os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")
        api_key = os.getenv("OPENROUTER_API_KEY")
        if not api_key:
            raise ValueError("OPENROUTER_API_KEY environment variable is required for OpenRouter")
        # Using openai-compatible API
        return init_chat_model(
            f"openai:{model_name}", 
            base_url=base_url, 
            api_key=api_key
        )
    else:
        # Fallback to openai for unknown providers (or add more as needed)
        return init_chat_model(f"openai:{model_name}")


def get_custom_tools() -> List[Tool]:
    """
    Get custom tools for the MaxClaw agent
    
    Returns:
        List of custom tools
    """
    custom_tools = []
    
    if not TOOLS_AVAILABLE:
        return custom_tools
    
    # Add Proxmox tools if configured
    proxmox_host = os.getenv("PROXMOX_HOST")
    if proxmox_host:
        proxmox_user = os.getenv("PROXMOX_USER", "root")
        proxmox_password = os.getenv("PROXMOX_PASSWORD")
        proxmox_token = os.getenv("PROXMOX_TOKEN")
        
        try:
            proxmox_tools = get_proxmox_tools(
                host=proxmox_host,
                user=proxmox_user,
                password=proxmox_password,
                api_token=proxmox_token
            )
            custom_tools.extend(proxmox_tools)
            print(f"✅ Added {len(proxmox_tools)} Proxmox tools")
        except Exception as e:
            print(f"⚠️  Could not load Proxmox tools: {e}")
    
    # Add unified provider tools for all configured providers
    try:
        provider_tools = get_all_provider_tools()
        custom_tools.extend(provider_tools)
        print(f"✅ Added {len(provider_tools)} provider tools")
    except Exception as e:
        print(f"⚠️  Could not load provider tools: {e}")
    
    return custom_tools


def create_maxclaw_agent(system_prompt: Optional[str] = None):
    """
    Create and return a DeepAgent instance with configurable model.
    """
    if system_prompt is None:
        system_prompt = """You are MaxClaw, a helpful AI agent that can plan, use tools, 
        manage context, and delegate to subagents. You are designed to be configurable 
        to work with various LLM providers.
        
        You have access to:
        1. Unified provider tools for managing LLM providers (list models, test connections, etc.)
        2. Custom tools for managing Proxmox virtual machines
        3. Standard DeepAgents tools (planning, filesystem, shell, subagents)
        
        Use these tools when asked about:
        - LLM provider management (listing models, testing connections, setting defaults)
        - System administration and virtualization
        - AI model management and configuration"""
    
    model = load_model_from_config()
    
    # Get custom tools
    custom_tools = get_custom_tools()
    
    # Create the agent with custom tools
    agent = create_deep_agent(
        model=model,
        tools=custom_tools,  # Add custom tools to the agent
        system_prompt=system_prompt
    )
    
    return agent


def main():
    """
    Main entry point for the MaxClaw agent.
    """
    print("🚀 Initializing MaxClaw Agent...")
    
    try:
        agent = create_maxclaw_agent()
        print("✅ Agent initialized successfully!")
        
        # Show available tools
        print("\n🔧 Available tools:")
        print("  • Provider tools: list models, test connections, set defaults")
        print("  • Proxmox tools: manage VMs, snapshots, resources")
        print("  • Standard tools: planning, filesystem, shell, subagents")
        
        print("\n💬 You can now interact with the agent. Type 'exit' to quit.")
        print("💡 Try asking:")
        print("   - 'List available models for my provider'")
        print("   - 'Test connection to Ollama'")
        print("   - 'Show me Proxmox VMs'")
        print("   - 'What's the default model for OpenAI?'")
        
        while True:
            user_input = input("\n👤 You: ").strip()
            if user_input.lower() in ['exit', 'quit']:
                print("👋 Goodbye!")
                break
            
            if not user_input:
                continue
                
            # Invoke the agent
            try:
                result = agent.invoke({"messages": [HumanMessage(content=user_input)]})
                # Extract the AI's response
                ai_message = result["messages"][-1]
                print(f"\n🤖 MaxClaw: {ai_message.content}")
            except Exception as e:
                print(f"\n❌ Error: {str(e)}")
                
    except Exception as e:
        print(f"❌ Failed to initialize agent: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()