#!/usr/bin/env python3
"""
Model Manager for Police Bot AI Runtime
Helps manage and switch between different Ollama models
"""

import os
import sys
import requests
import json
from config import config

def print_banner():
    """Print banner"""
    print("=" * 60)
    print("🤖 Police Bot AI - Model Manager")
    print("=" * 60)

def check_ollama_status():
    """Check if Ollama is running and get available models"""
    try:
        response = requests.get("http://localhost:11434/api/tags", timeout=5)
        if response.status_code == 200:
            data = response.json()
            models = data.get("models", [])
            return True, models
        else:
            return False, []
    except:
        return False, []

def list_available_models():
    """List all available models in Ollama"""
    print("\n📋 Available Models in Ollama:")
    print("-" * 40)
    
    ollama_running, models = check_ollama_status()
    
    if not ollama_running:
        print("❌ Ollama is not running")
        print("Please start Ollama first: ollama serve")
        return []
    
    if not models:
        print("⚠️  No models found in Ollama")
        print("You can pull models using:")
        print("  ollama pull mistral")
        print("  ollama pull llama3")
        print("  ollama pull police-bot")
        return []
    
    available_models = []
    for i, model in enumerate(models, 1):
        model_name = model.get("name", "Unknown")
        model_size = model.get("size", 0)
        size_mb = model_size / (1024 * 1024) if model_size > 0 else 0
        
        print(f"{i:2d}. {model_name:<20} ({size_mb:.0f} MB)")
        available_models.append(model_name)
    
    return available_models

def list_configured_models():
    """List models configured in the system"""
    print("\n⚙️  Configured Models in System:")
    print("-" * 40)
    
    configured_models = config.list_available_models()
    current_model = config.OLLAMA_MODEL
    
    for i, model in enumerate(configured_models, 1):
        status = "✓" if model == current_model else " "
        print(f"{i:2d}. [{status}] {model}")
    
    return configured_models

def switch_model():
    """Switch to a different model"""
    print("\n🔄 Switch Model")
    print("-" * 40)
    
    configured_models = config.list_available_models()
    current_model = config.OLLAMA_MODEL
    
    print(f"Current model: {current_model}")
    print("\nAvailable models:")
    
    for i, model in enumerate(configured_models, 1):
        print(f"{i}. {model}")
    
    try:
        choice = input(f"\nEnter choice (1-{len(configured_models)}): ").strip()
        if choice:
            model_index = int(choice) - 1
            if 0 <= model_index < len(configured_models):
                selected_model = configured_models[model_index]
                
                # Check if model is available in Ollama
                ollama_running, models = check_ollama_status()
                if ollama_running:
                    model_names = [m.get("name") for m in models]
                    if selected_model not in model_names:
                        print(f"⚠️  Warning: {selected_model} is not available in Ollama")
                        pull = input("Would you like to pull it? (y/n): ").lower().strip()
                        if pull == 'y':
                            print(f"Pulling {selected_model}...")
                            os.system(f"ollama pull {selected_model}")
                
                # Switch model
                if config.switch_model(selected_model):
                    print(f"✅ Switched to {selected_model}")
                    
                    # Save to environment file
                    save_model_to_env(selected_model)
                else:
                    print(f"❌ Failed to switch to {selected_model}")
            else:
                print("❌ Invalid choice")
        else:
            print("❌ No choice made")
    except ValueError:
        print("❌ Invalid input")

def save_model_to_env(model_name):
    """Save model selection to .env file"""
    env_file = ".env"
    env_content = f"OLLAMA_MODEL={model_name}\n"
    
    try:
        with open(env_file, 'w') as f:
            f.write(env_content)
        print(f"✅ Model selection saved to {env_file}")
    except Exception as e:
        print(f"⚠️  Could not save to {env_file}: {e}")

def test_model():
    """Test the current model"""
    print("\n🧪 Test Current Model")
    print("-" * 40)
    
    current_model = config.OLLAMA_MODEL
    print(f"Testing model: {current_model}")
    
    # Check if Ollama is running
    ollama_running, models = check_ollama_status()
    if not ollama_running:
        print("❌ Ollama is not running")
        return
    
    # Check if model is available
    model_names = [m.get("name") for m in models]
    if current_model not in model_names:
        print(f"❌ {current_model} is not available in Ollama")
        print(f"Available models: {', '.join(model_names)}")
        return
    
    # Test the model
    try:
        from police_runtime import PoliceBotRuntime
        
        runtime = PoliceBotRuntime()
        test_prompt = "Hello, how are you today?"
        
        print(f"Sending test prompt: '{test_prompt}'")
        response = runtime.get_llama3_response(test_prompt)
        
        if response:
            print(f"✅ Model responded: {response[:100]}...")
        else:
            print("❌ No response from model")
            
    except Exception as e:
        print(f"❌ Error testing model: {e}")

def show_model_info():
    """Show information about the current model"""
    print("\nℹ️  Current Model Information")
    print("-" * 40)
    
    current_model = config.OLLAMA_MODEL
    model_config = config.get_current_model_config()
    
    print(f"Model: {current_model}")
    print(f"Temperature: {model_config['temperature']}")
    print(f"Top P: {model_config['top_p']}")
    print(f"Max Tokens: {model_config['max_tokens']}")
    
    print(f"\nSystem Prompt Preview:")
    print("-" * 20)
    prompt_preview = model_config['system_prompt'][:200] + "..." if len(model_config['system_prompt']) > 200 else model_config['system_prompt']
    print(prompt_preview)

def pull_model():
    """Pull a model from Ollama"""
    print("\n📥 Pull Model from Ollama")
    print("-" * 40)
    
    model_name = input("Enter model name to pull (e.g., mistral, llama3): ").strip()
    if model_name:
        print(f"Pulling {model_name}...")
        result = os.system(f"ollama pull {model_name}")
        
        if result == 0:
            print(f"✅ Successfully pulled {model_name}")
        else:
            print(f"❌ Failed to pull {model_name}")
    else:
        print("❌ No model name provided")

def main():
    """Main function"""
    print_banner()
    
    while True:
        print("\nOptions:")
        print("1. List available models in Ollama")
        print("2. List configured models")
        print("3. Switch model")
        print("4. Test current model")
        print("5. Show model information")
        print("6. Pull model from Ollama")
        print("7. Exit")
        
        try:
            choice = input("\nEnter choice (1-7): ").strip()
            
            if choice == '1':
                list_available_models()
            elif choice == '2':
                list_configured_models()
            elif choice == '3':
                switch_model()
            elif choice == '4':
                test_model()
            elif choice == '5':
                show_model_info()
            elif choice == '6':
                pull_model()
            elif choice == '7':
                print("Goodbye! 👋")
                break
            else:
                print("❌ Invalid choice")
                
        except KeyboardInterrupt:
            print("\nGoodbye! 👋")
            break
        except Exception as e:
            print(f"❌ Error: {e}")

if __name__ == "__main__":
    main()