#!/usr/bin/env python3
"""
Simple test of SmoLagent functionality
"""
import pandas as pd
from smolagents import CodeAgent, DuckDuckGoSearchTool
import warnings
warnings.filterwarnings('ignore')

def test_basic_smolagent():
    """Test basic SmoLagent setup"""
    print("🧪 Testing SmoLagent Setup...")
    
    try:
        # Test 1: Import check
        print("✅ Imports successful")
        
        # Test 2: Create agent without model (should fail gracefully)
        try:
            agent = CodeAgent(tools=[DuckDuckGoSearchTool()])
            print("✅ Agent created without model")
        except Exception as e:
            print(f"❌ Agent creation failed: {e}")
            print("💡 This is expected - CodeAgent needs a model parameter")
        
        # Test 3: Try with a simple model setup
        try:
            from smolagents import HfApiModel
            print("🔄 Trying HuggingFace model...")
            model = HfApiModel(model_id="microsoft/DialoGPT-medium")
            agent = CodeAgent(
                tools=[DuckDuckGoSearchTool()],
                model=model
            )
            print("✅ Agent created successfully with HuggingFace model!")
            
            # Test a simple query
            response = agent.run("What is 2 + 2?")
            print(f"🤖 Agent response: {response}")
            
        except Exception as e:
            print(f"❌ HuggingFace model failed: {e}")
            
            # Test 4: Try Ollama
            try:
                from smolagents import OllamaModel
                print("🔄 Trying Ollama model...")
                model = OllamaModel(model_id="llama2", base_url="http://localhost:11434")
                agent = CodeAgent(
                    tools=[DuckDuckGoSearchTool()],
                    model=model
                )
                print("✅ Agent created successfully with Ollama!")
                
                # Test a simple query
                response = agent.run("What is 2 + 2?")
                print(f"🤖 Agent response: {response}")
                
            except Exception as e:
                print(f"❌ Ollama model failed: {e}")
                print("💡 Make sure Ollama is running with: ollama serve")
    
    except Exception as e:
        print(f"❌ Test failed: {e}")

def test_with_data():
    """Test SmoLagent with actual CSV data"""
    csv_file_path = "C:/Users/Cosmo/Desktop/NTU Peak Singtel/outsystems_sample_logs_6months.csv"
    
    try:
        # Load data
        print("\n📊 Loading CSV data...")
        df = pd.read_csv(csv_file_path)
        print(f"✅ Data loaded: {df.shape[0]} rows, {df.shape[1]} columns")
        print(f"📋 Columns: {list(df.columns)}")
        
        # Basic analysis
        error_count = df[df['LogLevel'] == 'Error'].shape[0]
        print(f"🚨 Error entries: {error_count}")
        
        return df
        
    except Exception as e:
        print(f"❌ Data loading failed: {e}")
        return None

if __name__ == "__main__":
    print("=" * 50)
    print("🤖 SMOLAGENT TEST SUITE")
    print("=" * 50)
    
    # Test basic functionality
    test_basic_smolagent()
    
    # Test with data
    df = test_with_data()
    
    print("\n" + "=" * 50)
    print("✅ Test completed!")
    print("=" * 50)