#!/usr/bin/env python3
"""
Lightweight BeastBullet + TinyLlama Demo
Works without torch - demonstrates integration concept
"""

import subprocess
import json
import os
from pathlib import Path

class SimpleBeastBulletDemo:
    """Simplified demo showing BeastBullet + TinyLlama integration"""
    
    def __init__(self, experts_dir="experts/"):
        self.experts_dir = Path(experts_dir)
        self.available_experts = self._discover_experts()
        
    def _discover_experts(self):
        """Discover available expert models"""
        expert_files = list(self.experts_dir.glob("*_expert_v1.0.pt"))
        experts = {}
        for f in expert_files:
            name = f.stem.replace("_v1.0", "")
            experts[name] = {
                "path": str(f),
                "size_mb": f.stat().st_size / (1024 * 1024)
            }
        return experts
    
    def route_query(self, query):
        """Simple routing based on query keywords"""
        query_lower = query.lower()
        selected = []
        
        # Math queries
        if any(w in query_lower for w in ["calculate", "math", "%", "percent", "+"]):
            selected.append("math_expert")
        
        # Logic queries  
        if any(w in query_lower for w in ["if", "then", "therefore", "logic"]):
            selected.append("logic_expert")
        
        # Code queries
        if any(w in query_lower for w in ["code", "function", "python", "javascript"]):
            selected.append("code_generation_expert")
        
        # Default to QA
        if not selected:
            selected.append("question_answering_expert")
        
        # Always add validator
        selected.append("validator_expert")
        
        return selected
    
    def query_tinyllama(self, prompt):
        """Query TinyLlama via Ollama CLI"""
        try:
            result = subprocess.run(
                ["ollama", "run", "tinyllama", prompt],
                capture_output=True,
                text=True,
                timeout=30
            )
            return result.stdout.strip()
        except Exception as e:
            return f"Error querying TinyLlama: {e}"
    
    def process_query(self, user_query):
        """Process query through BeastBullet + TinyLlama pipeline"""
        
        print("=" * 70)
        print(f"Query: {user_query}")
        print("=" * 70)
        
        # Step 1: Route to experts
        selected_experts = self.route_query(user_query)
        print(f"\n🔀 Selected experts: {', '.join(selected_experts)}")
        
        # Step 2: Show expert availability
        print(f"\n📊 Expert status:")
        for expert in selected_experts:
            if expert in self.available_experts:
                info = self.available_experts[expert]
                print(f"  ✅ {expert}: {info['size_mb']:.1f} MB")
            else:
                print(f"  ⚠️  {expert}: Not found")
        
        # Step 3: Build synthesis prompt
        prompt = f"""You are a helpful AI assistant integrated with BeastBullet expert system.

Question: {user_query}

Expert Analysis:
"""
        for expert in selected_experts:
            if expert in self.available_experts:
                prompt += f"\n- {expert}: Available for analysis"
        
        prompt += "\n\nProvide a clear, accurate answer:"
        
        # Step 4: Query TinyLlama
        print(f"\n🤖 Querying TinyLlama via Ollama...")
        answer = self.query_tinyllama(prompt)
        
        # Step 5: Return result
        print(f"\n💡 Answer:")
        print(f"   {answer}")
        print("\n" + "=" * 70)
        
        return {
            "query": user_query,
            "answer": answer,
            "experts_selected": selected_experts,
            "experts_available": len(self.available_experts)
        }


def main():
    print("=" * 70)
    print("🧠 BeastBullet + TinyLlama Integration Demo")
    print("=" * 70)
    
    # Initialize demo
    demo = SimpleBeastBulletDemo("experts/")
    
    print(f"\n📦 Discovered {len(demo.available_experts)} expert models:")
    for name, info in sorted(demo.available_experts.items()):
        print(f"  • {name}: {info['size_mb']:.1f} MB")
    
    print(f"\n✅ TinyLlama available in Ollama")
    
    # 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?",
    ]
    
    for query in test_queries:
        print()
        demo.process_query(query)
    
    print("\n" + "=" * 70)
    print("✅ Demo Complete!")
    print("=" * 70)
    print("\nNote: This is a lightweight demo. For full functionality,")
    print("install torch to load the actual expert models:")
    print("  pip install torch")


if __name__ == "__main__":
    main()