#!/usr/bin/env python3
"""
Test script for Tool Calling Module
====================================
Tests sequential multi-tool execution.
"""

import sys
import os
sys.path.insert(0, os.path.dirname(__file__))

from tool_caller import OllamaToolCaller, ToolResult
import json


def test_single_tool():
    """Test single tool call"""
    print("=" * 60)
    print("TEST 1: Single Tool Call (PostgreSQL)")
    print("=" * 60)
    
    caller = OllamaToolCaller()
    caller.register_default_tools()
    
    query = "How many videos are in the database?"
    print(f"Query: {query}")
    print("-" * 60)
    
    result = caller.run(query)
    print(f"Result:\n{result}")
    print()


def test_multi_tool():
    """Test multi-tool sequential call"""
    print("=" * 60)
    print("TEST 2: Multi-Tool Sequential (PostgreSQL → Qdrant)")
    print("=" * 60)
    
    caller = OllamaToolCaller()
    caller.register_default_tools()
    
    query = "Find videos about dogs, then search for similar content in the vector database"
    print(f"Query: {query}")
    print("-" * 60)
    
    result = caller.run(query)
    print(f"Result:\n{result}")
    print()


def test_tool_direct():
    """Test direct tool execution"""
    print("=" * 60)
    print("TEST 3: Direct Tool Execution")
    print("=" * 60)
    
    caller = OllamaToolCaller()
    caller.register_default_tools()
    
    # Test PostgreSQL directly
    print("Testing PostgreSQL tool directly:")
    result = caller.registry.execute("query_postgres", {
        "query": "SELECT COUNT(*) as count FROM videos"
    })
    print(f"  Success: {result.success}")
    print(f"  Data: {result.data}")
    print(f"  Time: {result.execution_time_ms:.1f}ms")
    print()
    
    # Test Bash directly
    print("Testing Bash tool directly:")
    result = caller.registry.execute("execute_bash", {
        "command": "echo 'Hello from Tool Caller!' && date"
    })
    print(f"  Success: {result.success}")
    print(f"  Data: {result.data}")
    print(f"  Time: {result.execution_time_ms:.1f}ms")
    print()


def test_bash_safety():
    """Test bash command safety"""
    print("=" * 60)
    print("TEST 4: Bash Safety Check")
    print("=" * 60)
    
    caller = OllamaToolCaller()
    caller.register_default_tools()
    
    # Test blocked command
    print("Testing blocked command (rm -rf /):")
    result = caller.registry.execute("execute_bash", {
        "command": "rm -rf /"
    })
    print(f"  Success: {result.success}")
    print(f"  Error: {result.error}")
    print()
    
    # Test safe command
    print("Testing safe command:")
    result = caller.registry.execute("execute_bash", {
        "command": "ls -la /tmp | head -5"
    })
    print(f"  Success: {result.success}")
    print(f"  Data: {result.data}")
    print()


def main():
    """Run all tests"""
    print("\n" + "=" * 60)
    print("TOOL CALLING MODULE - TEST SUITE")
    print("=" * 60 + "\n")
    
    try:
        # Test 1: Single tool
        test_single_tool()
        
        # Test 2: Multi tool
        test_multi_tool()
        
        # Test 3: Direct execution
        test_tool_direct()
        
        # Test 4: Safety
        test_bash_safety()
        
        print("=" * 60)
        print("ALL TESTS COMPLETED")
        print("=" * 60)
        
    except Exception as e:
        print(f"\nERROR: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()