#!/usr/bin/env python3
"""
Simple BeastBullet + TinyLlama Integration Test
Uses existing Ollama TinyLlama installation
"""

import sys
import os

# Add code directory to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'code'))

from integration_example import ExpertLoader, TinyLlamaIntegration
import json

def main():
    print("=" * 70)
    print("🧠 BeastBullet + TinyLlama Integration")
    print("=" * 70)
    
    # Step 1: Load BeastBullet experts
    print("\n📦 Loading BeastBullet experts...")
    loader = ExpertLoader("experts/")
    
    try:
        # Load encoder (if available)
        if os.path.exists("experts/shared_encoder_v1.0.pt"):
            loader.load_encoder()
        else:
            print("⚠️  Shared encoder not found, using placeholder")
            loader.encoder = None
    except Exception as e:
        print(f"⚠️  Could not load encoder: {e}")
        loader.encoder = None
    
    # Load all expert models
    print("\n🔧 Loading expert models...")
    loader.load_all_experts()
    
    # Step 2: Create TinyLlama integration
    print("\n🔗 Creating TinyLlama integration...")
    integration = TinyLlamaIntegration(
        experts=loader.experts,
        encoder=loader.encoder
    )
    
    # Step 3: Load Sonnet-quality configuration
    print("\n⚙️  Loading configuration...")
    try:
        with open("configs/config_sonnet_quality.json", "r") as f:
            config = json.load(f)
            synthesis_config = config.get("synthesis", {})
    except Exception as e:
        print(f"⚠️  Could not load config: {e}, using defaults")
        synthesis_config = {"temperature": 0.3, "max_tokens": 512}
    
    # Step 4: Test queries
    print("\n" + "=" * 70)
    print("📝 Running Test Queries")
    print("=" * 70)
    
    test_queries = [
        "What is 15% of 240?",
        "If all cats are mammals, and Fluffy is a cat, what is Fluffy?",
        "Explain what is machine learning in simple terms",
    ]
    
    for i, query in enumerate(test_queries, 1):
        print(f"\n{'─' * 70}")
        print(f"Query {i}: {query}")
        print('─' * 70)
        
        try:
            result = integration.query(query, synthesis_config)
            
            print(f"\n💡 Answer:")
            print(f"   {result['answer']}")
            print(f"\n📊 Details:")
            print(f"   • Confidence: {result['confidence']:.2%}")
            print(f"   • Experts used: {', '.join(result['experts_used'])}")
            print(f"\n🔍 Reasoning trace:")
            for step in result['reasoning_trace']:
                print(f"   - {step}")
                
        except Exception as e:
            print(f"❌ Error processing query: {e}")
            import traceback
            traceback.print_exc()
    
    print("\n" + "=" * 70)
    print("✅ Integration test complete!")
    print("=" * 70)

if __name__ == "__main__":
    main()