#!/usr/bin/env python3
"""
Minimal Modern LangChain 1.0+ Agent for Skyrelis Testing

This is a stripped-down agent designed to test Skyrelis monitoring
compatibility with modern LangChain patterns without external dependencies.
"""

import os
from typing import List, Optional
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Modern LangChain imports
try:
    from langchain_core.tools import tool
    from langchain_openai import ChatOpenAI
    from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
    LANGCHAIN_AVAILABLE = True
    print("✅ LangChain 1.0+ loaded successfully")
except ImportError as e:
    print(f"❌ LangChain import error: {e}")
    LANGCHAIN_AVAILABLE = False

# Skyrelis monitoring
try:
    from skyrelis import observe
    SKYRELIS_AVAILABLE = True
    print("✅ Skyrelis library loaded successfully")
except ImportError as e:
    print(f"❌ Skyrelis import error: {e}")
    SKYRELIS_AVAILABLE = False


class MinimalModernAgent:
    """
    Minimal modern LangChain 1.0+ agent for Skyrelis testing
    
    Features:
    - Modern LangChain patterns (no AgentExecutor)
    - Simple built-in tools
    - Direct LLM function calling
    - Skyrelis monitoring integration
    """
    
    def __init__(self, api_key: str = None, monitor_url: str = None):
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
        self.monitor_url = monitor_url or os.getenv("MONITOR_URL")
        
        # Initialize components
        self.llm = None
        self.llm_with_tools = None
        self.tools = []
        
        # Setup
        self._init_llm()
        self._setup_tools()
        
        print(f"🤖 Minimal Modern Agent initialized")
        print(f"   LLM Available: {self.llm is not None}")
        print(f"   Tools Count: {len(self.tools)}")
        print(f"   Skyrelis Monitor: {self.monitor_url is not None}")

    def _init_llm(self):
        """Initialize OpenAI LLM"""
        if not LANGCHAIN_AVAILABLE or not self.api_key:
            print("⚠️  LLM initialization skipped (missing LangChain or API key)")
            return
            
        try:
            self.llm = ChatOpenAI(
                model="gpt-4o-mini",
                temperature=0,
                api_key=self.api_key
            )
            print("✅ OpenAI LLM initialized")
        except Exception as e:
            print(f"❌ LLM initialization failed: {e}")

    def _setup_tools(self):
        """Setup simple tools"""
        if not LANGCHAIN_AVAILABLE:
            return
            
        # Create simple tools
        self.tools = [
            self._create_calculator_tool(),
            self._create_info_tool()
        ]
        
        # Bind tools to LLM
        if self.llm:
            try:
                self.llm_with_tools = self.llm.bind_tools(self.tools)
                print("✅ Tools bound to LLM successfully")
            except Exception as e:
                print(f"❌ Tool binding failed: {e}")

    def _create_calculator_tool(self):
        """Simple calculator tool"""
        @tool
        def simple_calculator(expression: str) -> str:
            """Perform simple arithmetic calculations.
            
            Args:
                expression: Mathematical expression like "2 + 3" or "10 * 5"
                
            Returns:
                The result of the calculation
            """
            try:
                # Simple safe evaluation for basic math
                allowed_chars = set('0123456789+-*/.() ')
                if not all(c in allowed_chars for c in expression):
                    return "Error: Only basic arithmetic operations allowed"
                
                result = eval(expression)
                return f"Result: {result}"
            except Exception as e:
                return f"Error: {str(e)}"
        
        return simple_calculator

    def _create_info_tool(self):
        """Simple information tool"""
        @tool
        def agent_info(query: str) -> str:
            """Get information about this agent and its capabilities.
            
            Args:
                query: Question about the agent
                
            Returns:
                Information about the agent
            """
            info_db = {
                "name": "Minimal Modern Agent",
                "version": "1.0",
                "langchain": "1.0+",
                "capabilities": "Basic math, information queries",
                "monitoring": "Skyrelis integration",
                "purpose": "Testing modern LangChain + Skyrelis compatibility"
            }
            
            query_lower = query.lower()
            
            if "name" in query_lower:
                return f"Agent name: {info_db['name']}"
            elif "version" in query_lower:
                return f"Version: {info_db['version']}"
            elif "capabilities" in query_lower or "what can" in query_lower:
                return f"Capabilities: {info_db['capabilities']}"
            elif "langchain" in query_lower:
                return f"LangChain version: {info_db['langchain']}"
            elif "monitoring" in query_lower or "skyrelis" in query_lower:
                return f"Monitoring: {info_db['monitoring']}"
            else:
                return f"This is {info_db['name']} v{info_db['version']} - {info_db['purpose']}"
        
        return agent_info

    def ask(self, user_msg: str) -> str:
        """
        Process a user query with optional Skyrelis monitoring
        
        Args:
            user_msg: User's question
            
        Returns:
            Agent's response
        """
        if not user_msg or not user_msg.strip():
            return "Please provide a question."
        
        print(f"\n🔍 Processing: {user_msg}")
        
        # Try modern LangChain with Skyrelis monitoring
        if LANGCHAIN_AVAILABLE and self.llm_with_tools:
            try:
                return self._process_with_llm(user_msg)
            except Exception as e:
                print(f"❌ LLM processing failed: {e}")
                return self._fallback_response(user_msg)
        else:
            return self._fallback_response(user_msg)

    def _process_with_llm(self, user_msg: str) -> str:
        """Process with LLM and optional Skyrelis monitoring"""
        
        # Create messages
        system_msg = SystemMessage(content="""You are a helpful AI assistant with access to tools.
        
Rules:
- Be concise and helpful
- Use tools when appropriate
- Provide clear, direct answers""")
        
        user_message = HumanMessage(content=user_msg)
        messages = [system_msg, user_message]
        
        # Process with or without Skyrelis monitoring
        if SKYRELIS_AVAILABLE and self.monitor_url:
            print("🔒 Using Skyrelis monitoring")
            return self._process_with_skyrelis_monitoring(messages)
        else:
            print("📡 Processing without monitoring")
            return self._process_without_monitoring(messages)

    def _process_with_skyrelis_monitoring(self, messages: List) -> str:
        """Process with Skyrelis monitoring"""
        try:
            # Method 1: Direct function monitoring
            monitored_invoke = observe(
                remote_observer_url=self.monitor_url,
                agent_name="minimal_modern_agent"
            )(self.llm_with_tools.invoke)
            
            response = monitored_invoke(messages)
            
        except Exception as e:
            print(f"⚠️  Skyrelis monitoring failed: {e}")
            print("   Falling back to direct LLM call")
            response = self.llm_with_tools.invoke(messages)
        
        return self._handle_response(response, messages)

    def _process_without_monitoring(self, messages: List) -> str:
        """Process without monitoring"""
        response = self.llm_with_tools.invoke(messages)
        return self._handle_response(response, messages)

    def _handle_response(self, response, messages: List) -> str:
        """Handle LLM response and tool calls"""
        
        # Check for tool calls
        if hasattr(response, 'tool_calls') and response.tool_calls:
            print(f"🛠️  Processing {len(response.tool_calls)} tool call(s)")
            
            # Add AI response to messages
            messages.append(response)
            
            # Execute tool calls
            for tool_call in response.tool_calls:
                tool_name = tool_call["name"]
                tool_args = tool_call["args"]
                print(f"   Executing: {tool_name} with {tool_args}")
                
                # Find and execute tool
                tool_result = None
                for tool in self.tools:
                    if tool.name == tool_name:
                        tool_result = tool.invoke(tool_args)
                        break
                
                if tool_result:
                    print(f"   Result: {tool_result}")
                    # Add tool result to messages
                    from langchain_core.messages import ToolMessage
                    messages.append(ToolMessage(
                        content=str(tool_result),
                        tool_call_id=tool_call["id"]
                    ))
            
            # Get final response
            if SKYRELIS_AVAILABLE and self.monitor_url:
                try:
                    monitored_final = observe(
                        remote_observer_url=self.monitor_url,
                        agent_name="minimal_modern_agent_final"
                    )(self.llm.invoke)
                    final_response = monitored_final(messages)
                except Exception as e:
                    print(f"⚠️  Final Skyrelis monitoring failed: {e}")
                    final_response = self.llm.invoke(messages)
            else:
                final_response = self.llm.invoke(messages)
            
            return final_response.content
        else:
            print("ℹ️  Direct response (no tools needed)")
            return response.content

    def _fallback_response(self, user_msg: str) -> str:
        """Simple fallback when LLM is not available"""
        user_lower = user_msg.lower()
        
        if any(word in user_lower for word in ['calculate', 'math', '+', '-', '*', '/']):
            return "I can help with calculations, but I need my LLM to be properly initialized."
        elif any(word in user_lower for word in ['info', 'about', 'what', 'who']):
            return "I'm a Minimal Modern Agent designed to test Skyrelis monitoring with LangChain 1.0+."
        else:
            return "I'm a test agent for modern LangChain + Skyrelis integration. My LLM is not currently available."

    def get_status(self) -> dict:
        """Get agent status"""
        return {
            "agent_type": "MinimalModernAgent",
            "langchain_available": LANGCHAIN_AVAILABLE,
            "skyrelis_available": SKYRELIS_AVAILABLE,
            "llm_available": self.llm is not None,
            "tools_count": len(self.tools),
            "monitor_url": self.monitor_url is not None,
            "api_key_set": self.api_key is not None
        }


def main():
    """Test the minimal agent"""
    print("🧪 Testing Minimal Modern Agent with Skyrelis")
    print("=" * 60)
    
    # Create agent
    agent = MinimalModernAgent()
    
    # Show status
    status = agent.get_status()
    print(f"\n📊 Agent Status:")
    for key, value in status.items():
        print(f"   {key}: {value}")
    
    # Test queries
    test_queries = [
        "What is your name?",
        "Calculate 15 + 27",
        "What are your capabilities?",
        "Tell me about Skyrelis monitoring"
    ]
    
    print(f"\n🔍 Testing Queries:")
    for query in test_queries:
        print(f"\n" + "-" * 50)
        response = agent.ask(query)
        print(f"✅ Response: {response}")
    
    print(f"\n" + "=" * 60)
    print("✅ Minimal agent testing completed!")


if __name__ == "__main__":
    main()
