#!/usr/bin/env python3
"""
Simple Novita AI Chat Application
"""

import os
import requests
import json
import time

class NovitaAIChat:
    def __init__(self, api_key):
        self.api_key = api_key
        self.base_url = "https://api.novita.ai/openai"
        self.headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
        self.conversation_history = []
        self.current_model = "meta-llama/llama-3.2-1b-instruct"
    
    def get_available_models(self):
        """Get list of available models"""
        try:
            response = requests.get(f"{self.base_url}/models", headers=self.headers, timeout=10)
            if response.status_code == 200:
                models = response.json()
                return models.get('data', [])
            else:
                print(f"❌ Error getting models: {response.status_code}")
                return []
        except Exception as e:
            print(f"❌ Error: {e}")
            return []
    
    def chat_completion(self, message, model=None):
        """Send message to Novita AI and get response"""
        if model is None:
            model = self.current_model
        
        # Add user message to history
        self.conversation_history.append({"role": "user", "content": message})
        
        # Prepare payload
        payload = {
            "model": model,
            "messages": self.conversation_history,
            "max_tokens": 500,
            "temperature": 0.7,
            "top_p": 0.9
        }
        
        try:
            print("🤖 Thinking...", end="", flush=True)
            response = requests.post(
                f"{self.base_url}/chat/completions", 
                headers=self.headers, 
                json=payload, 
                timeout=60
            )
            
            if response.status_code == 200:
                result = response.json()
                assistant_message = result.get('choices', [{}])[0].get('message', {}).get('content', '')
                
                # Add assistant response to history
                self.conversation_history.append({"role": "assistant", "content": assistant_message})
                
                print("\r" + " " * 20 + "\r", end="")  # Clear "Thinking..." message
                return assistant_message
            else:
                print(f"\r❌ Error: {response.status_code} - {response.text}")
                return None
                
        except Exception as e:
            print(f"\r❌ Error: {e}")
            return None
    
    def change_model(self, model_id):
        """Change the current model"""
        self.current_model = model_id
        print(f"✅ Model changed to: {model_id}")
    
    def clear_history(self):
        """Clear conversation history"""
        self.conversation_history = []
        print("✅ Conversation history cleared")
    
    def show_models(self):
        """Show available models"""
        models = self.get_available_models()
        if models:
            print("\n📋 Available Models:")
            print("-" * 50)
            for i, model in enumerate(models[:20], 1):  # Show first 20 models
                model_id = model.get('id', 'Unknown')
                print(f"{i:2d}. {model_id}")
            print("-" * 50)
            print(f"Current model: {self.current_model}")
        else:
            print("❌ Could not fetch models")

def main():
    print("🚀 Novita AI Chat Application")
    print("=" * 50)
    
    # Check API key
    api_key = os.getenv('NOVITA_API_KEY')
    if not api_key:
        print("❌ NOVITA_API_KEY not found")
        api_key = input("Enter your Novita AI API key: ").strip()
        if not api_key:
            print("❌ API key required")
            return
        os.environ['NOVITA_API_KEY'] = api_key
    
    # Initialize chat
    chat = NovitaAIChat(api_key)
    
    # Test connection
    print("🔍 Testing connection...")
    models = chat.get_available_models()
    if not models:
        print("❌ Could not connect to Novita AI")
        return
    
    print(f"✅ Connected! Found {len(models)} models")
    
    # Show current model
    print(f"🤖 Current model: {chat.current_model}")
    
    # Main chat loop
    print("\n💬 Start chatting! Type 'help' for commands, 'quit' to exit")
    print("-" * 50)
    
    while True:
        try:
            user_input = input("\n👤 You: ").strip()
            
            if not user_input:
                continue
            
            # Handle commands
            if user_input.lower() in ['quit', 'exit', 'q']:
                print("👋 Goodbye!")
                break
            elif user_input.lower() == 'help':
                print("\n📋 Available Commands:")
                print("  help          - Show this help")
                print("  models        - Show available models")
                print("  change <id>   - Change model (e.g., change 5)")
                print("  clear         - Clear conversation history")
                print("  quit/exit/q   - Exit the application")
                print("  <any text>    - Send message to AI")
                continue
            elif user_input.lower() == 'models':
                chat.show_models()
                continue
            elif user_input.lower() == 'clear':
                chat.clear_history()
                continue
            elif user_input.lower().startswith('change '):
                try:
                    model_num = int(user_input.split()[1])
                    models = chat.get_available_models()
                    if 1 <= model_num <= len(models):
                        new_model = models[model_num - 1].get('id')
                        chat.change_model(new_model)
                    else:
                        print(f"❌ Invalid model number. Use 1-{len(models)}")
                except (ValueError, IndexError):
                    print("❌ Usage: change <number>")
                continue
            
            # Send message to AI
            response = chat.chat_completion(user_input)
            if response:
                print(f"\n🤖 Assistant: {response}")
            
        except KeyboardInterrupt:
            print("\n👋 Goodbye!")
            break
        except Exception as e:
            print(f"❌ Error: {e}")

if __name__ == "__main__":
    main()