#!/usr/bin/env python3
"""
TinyLlama Wrapper for LEGO-MoE
Simple interface to TinyLlama via Ollama
"""

import subprocess
import json
from typing import Optional

class TinyLlamaWrapper:
    """Wrapper for TinyLlama-1.1B via Ollama"""
    
    def __init__(self, model_name: str = "tinyllama"):
        self.model_name = model_name
        self._verify_model()
    
    def _verify_model(self):
        """Check if TinyLlama is available"""
        try:
            result = subprocess.run(
                ["ollama", "list"],
                capture_output=True,
                text=True,
                timeout=5
            )
            
            if self.model_name not in result.stdout:
                print(f"⚠️  {self.model_name} not found. Pulling...")
                self._pull_model()
            else:
                print(f"✓ {self.model_name} available")
        
        except Exception as e:
            print(f"Error checking model: {e}")
    
    def _pull_model(self):
        """Download TinyLlama if not present"""
        try:
            print(f"Downloading {self.model_name}...")
            subprocess.run(
                ["ollama", "pull", self.model_name],
                check=True
            )
            print(f"✓ {self.model_name} downloaded")
        except Exception as e:
            print(f"Error downloading model: {e}")
            raise
    
    def generate(
        self,
        prompt: str,
        temperature: float = 0.7,
        max_tokens: int = 256,
        stop: Optional[list] = None
    ) -> str:
        """
        Generate response from TinyLlama
        
        Args:
            prompt: Input prompt
            temperature: Sampling temperature (0.0-1.0)
            max_tokens: Maximum tokens to generate
            stop: Stop sequences
        
        Returns:
            Generated text
        """
        try:
            # Build Ollama command
            cmd = [
                "ollama", "run", self.model_name,
                prompt
            ]
            
            # Run generation with shorter timeout
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=15  # Reduced from 30 to 15 seconds
            )
            
            if result.returncode == 0:
                return result.stdout.strip()
            else:
                return f"Error: {result.stderr}"
        
        except subprocess.TimeoutExpired:
            return "Error: Generation timeout"
        except Exception as e:
            return f"Error: {str(e)}"
    
    def generate_with_context(
        self,
        query: str,
        expert_outputs: dict,
        temperature: float = 0.7
    ) -> str:
        """
        Generate response with expert context
        
        Args:
            query: User query
            expert_outputs: Dict of expert name -> output
            temperature: Sampling temperature
        
        Returns:
            Generated response
        """
        # Build context
        context = f"Query: {query}\n\n"
        
        if expert_outputs:
            context += "Expert Analysis:\n"
            for expert_name, output in expert_outputs.items():
                context += f"- {expert_name}: {output}\n"
            context += "\n"
        
        context += "Provide a clear, concise response:\n"
        
        # Generate
        return self.generate(context, temperature=temperature)


def test_tinyllama():
    """Test TinyLlama wrapper"""
    print("="*70)
    print("Testing TinyLlama Wrapper")
    print("="*70)
    
    llama = TinyLlamaWrapper()
    
    # Test 1: Simple generation
    print("\nTest 1: Simple query")
    response = llama.generate("What is 2+2?", temperature=0.1, max_tokens=50)
    print(f"Response: {response}")
    
    # Test 2: With expert context
    print("\nTest 2: With expert context")
    expert_outputs = {
        "Math Expert": "4"
    }
    response = llama.generate_with_context(
        "What is 2+2?",
        expert_outputs,
        temperature=0.1
    )
    print(f"Response: {response}")
    
    print("\n" + "="*70)
    print("✓ TinyLlama wrapper working!")
    print("="*70)


if __name__ == "__main__":
    test_tinyllama()