main.py

#!/usr/bin/env python3
"""
kh_Assistant - Basic Chat Version with voice implementation
A local AI assistant with PyQt6 GUI
"""

import sys
import subprocess
import threading
from PyQt6.QtWidgets import (QApplication, QWidget, QVBoxLayout, 
                             QHBoxLayout, QTextEdit, QLineEdit, 
                             QPushButton, QLabel, QComboBox)
from PyQt6.QtCore import Qt, QTimer
from PyQt6.QtGui import QFont

# Import our smart modules
from voice_worker import VoiceWorker
from ai_manager import AIManager


class kh_Assistant(QWidget):
    def __init__(self):
        super().__init__()
        
        # Initialize AI Manager
        self.ai_manager = AIManager()
        
        # Voice state
        self.voice_worker = None
        self.is_listening = False

        self.setup_ui()
        
        # Welcome message
        self.add_message("kh", "Hello. I am kh. Let me check available AI models...")
        
        # Auto-detect models after UI is shown (non-blocking)
        QTimer.singleShot(100, self.refresh_models)
        
    def setup_ui(self):
        """Setup the user interface"""
        self.setWindowTitle("kh_Assistant")
        self.setGeometry(100, 100, 500, 600)
        
        # Catppuccin Mocha theme stylesheet
        self.setStyleSheet("""
            QWidget { 
                background: #1e1e2e;
                color: #cdd6f4;
                font-family: 'Segoe UI', sans-serif; 
            }
            QTextEdit { 
                background: #181825;
                color: #cdd6f4;
                border: 1px solid #45475a;
                border-radius: 8px; 
                padding: 8px; 
                selection-background-color: #585b70;
                selection-color: #cdd6f4;
            }
            QLineEdit { 
                background: #313244;
                color: #cdd6f4;
                border: 1px solid #45475a;
                border-radius: 6px; 
                padding: 6px; 
                selection-background-color: #585b70;
                selection-color: #cdd6f4;
            }
            QLineEdit::placeholder {
                color: #6c7086;
            }
            QPushButton { 
                background: #89b4fa;
                color: #1e1e2e;
                border: none;
                border-radius: 6px; 
                padding: 6px 12px; 
                font-weight: bold; 
            }
            QPushButton:hover { 
                background: #b4befe;
            }
            QPushButton:pressed { 
                background: #74c7ec;
            }
            QPushButton:disabled { 
                background: #45475a;
                color: #6c7086;
            }
            QLabel { 
                padding: 5px; 
                color: #cdd6f4;
            }
            QComboBox {
                background: #313244;
                color: #cdd6f4;
                border: 1px solid #45475a;
                border-radius: 4px;
                padding: 4px 8px;
            }
            QComboBox::drop-down {
                border: none;
                width: 20px;
            }
            QComboBox QAbstractItemView {
                background: #181825;
                color: #cdd6f4;
                border: 1px solid #45475a;
                selection-background-color: #585b70;
            }
            QScrollBar:vertical {
                background: #181825;
                width: 10px;
                border-radius: 5px;
            }
            QScrollBar::handle:vertical {
                background: #585b70;
                border-radius: 5px;
                min-height: 20px;
            }
            QScrollBar::handle:vertical:hover {
                background: #6c7086;
            }
            QScrollBar::add-line:vertical, QScrollBar::sub-line:vertical {
                height: 0px;
            }
        """)
        
        # Create main layout
        layout = QVBoxLayout()
        layout.setContentsMargins(12, 12, 12, 12)
        layout.setSpacing(10)
        
        # Header
        header = QLabel("kh_Assistant v1.1")
        header.setStyleSheet("font-size: 18px; font-weight: bold; color: #cba6f7;")
        header.setAlignment(Qt.AlignmentFlag.AlignCenter)
        layout.addWidget(header)
        
        # Model Selector
        model_row = QHBoxLayout()
        
        model_label = QLabel("Model:")
        model_label.setStyleSheet("color: #a6adc8;")
        model_row.addWidget(model_label)
        
        self.model_selector = QComboBox()
        self.model_selector.currentTextChanged.connect(self.on_model_changed)
        self.model_selector.setMinimumWidth(200)
        model_row.addWidget(self.model_selector)
        
        # Refresh button
        refresh_btn = QPushButton("Refresh")
        refresh_btn.setFixedSize(70, 30)
        refresh_btn.setToolTip("Refresh model list")
        refresh_btn.clicked.connect(self.refresh_models)
        model_row.addWidget(refresh_btn)
        
        model_row.addStretch()
        layout.addLayout(model_row)
        
        # Status bar
        self.status = QLabel("Status: Checking Ollama...")
        self.status.setStyleSheet("color: #f9e2af; font-size: 12px;")
        layout.addWidget(self.status)

        # Chat display
        self.chat_box = QTextEdit()
        self.chat_box.setReadOnly(True)
        self.chat_box.setPlaceholderText("Conversation will appear here...")
        layout.addWidget(self.chat_box)
        
        # Input row
        input_layout = QHBoxLayout()
        
        self.input_box = QLineEdit()
        self.input_box.setPlaceholderText("Type a message...")
        self.input_box.returnPressed.connect(self.send_message)
        input_layout.addWidget(self.input_box)
        
        self.send_button = QPushButton("Send")
        self.send_button.clicked.connect(self.send_message)
        input_layout.addWidget(self.send_button)
        
        # Voice button
        self.voice_button = QPushButton("Voice")
        self.voice_button.setFixedSize(60, 35)
        self.voice_button.clicked.connect(self.toggle_voice)
        self.voice_button.setToolTip("Click to speak")
        input_layout.addWidget(self.voice_button)
        
        layout.addLayout(input_layout)
        self.setLayout(layout)
    
    def add_message(self, sender, text):
        """Add a message to the chat display"""
        # Catppuccin colors: Blue for user, Green for assistant
        color = "#89b4fa" if sender == "You" else "#a6e3a1"
        self.chat_box.append(f'<div style="color:{color}; margin: 4px 0;"><b>{sender}:</b> {text}</div>')
        # Auto-scroll to bottom
        scrollbar = self.chat_box.verticalScrollBar()
        scrollbar.setValue(scrollbar.maximum())
    
    def send_message(self):
        """Handle sending a message to AI"""
        user_text = self.input_box.text().strip()
        if not user_text:
            return
        
        # Check if Ollama is running
        if not self.ai_manager.is_ollama_running:
            self.add_message("kh", "Ollama is not running. Click Refresh to check again.")
            return
        
        # Display user message
        self.add_message("You", user_text)
        self.input_box.clear()
        
        # Update status
        self.status.setText(f"Status: {self.ai_manager.current_model} is processing...")
        self.status.setStyleSheet("color: #f9e2af; font-size: 12px;")
        self.send_button.setEnabled(False)
        self.voice_button.setEnabled(False)
        
        # Get AI response in background
        threading.Thread(target=self.get_ai_response, args=(user_text,), daemon=True).start()
    
    def get_ai_response(self, user_text):
        """Get response from AI using AIManager"""
        result = self.ai_manager.query_ai(user_text)
        
        if result["success"]:
            # Display response
            self.add_message("kh", result["response"])
            
            # Show response time
            duration_ms = result.get("total_duration", 0)
            duration_sec = round(duration_ms / 1000, 2)
            self.status.setText(f"Status: Ready ({duration_sec}s)")
            self.status.setStyleSheet("color: #a6e3a1; font-size: 12px;")
        else:
            # Show error
            self.add_message("kh", f"Error: {result['error']}")
            self.status.setText("Status: Error")
            self.status.setStyleSheet("color: #f38ba8; font-size: 12px;")
            
        self.send_button.setEnabled(True)
        self.voice_button.setEnabled(True)
    
    def refresh_models(self):
        """Refresh the list of available AI models"""
        self.status.setText("Status: Refreshing models...")
        self.status.setStyleSheet("color: #89b4fa; font-size: 12px;")
        self.model_selector.clear()
        self.model_selector.addItem("Loading...")
        self.model_selector.setEnabled(False)

        # Run in background thread
        threading.Thread(target=self._refresh_models_background, daemon=True).start()
    
    def _refresh_models_background(self):
        """Background thread: check Ollama and fetch models"""
        # Check Ollama
        ollama_ok = self.ai_manager.check_ollama_status()
        
        if not ollama_ok:
            # Update UI safely from main thread
            self._ui_update_ollama_error()
            return
        
        # Fetch models
        models = self.ai_manager.get_available_models()
        
        # Update UI safely from main thread
        self._ui_update_model_list(models)
    
    def _ui_update_ollama_error(self):
        """Update UI when Ollama is not available (MUST run on main thread)"""
        self.status.setText("Status: Ollama not running")
        self.status.setStyleSheet("color: #f38ba8; font-size: 12px;")
        self.model_selector.clear()
        self.model_selector.addItem("Ollama not running")
        self.model_selector.setEnabled(False)
        self.add_message("kh", "Ollama is not running. Run 'ollama serve' in a terminal.")
    
    def _ui_update_model_list(self, models):
        """Update model selector with found models (MUST run on main thread)"""
        self.model_selector.clear()
        self.model_selector.setEnabled(True)
        
        if not models:
            self.model_selector.addItem("No models found")
            self.status.setText("Status: No models installed")
            self.status.setStyleSheet("color: #f9e2af; font-size: 12px;")
            self.add_message("kh", "No AI models found. Run: ollama pull llama3.1:8b")
            return
        
        # Add models to dropdown
        for model in models:
            display_name = f"{model['name']} ({model['size_gb']} GB)"
            self.model_selector.addItem(display_name, model['name'])
        
        # Auto-select recommended model
        recommended = self.ai_manager.get_recommended_model()
        if recommended:
            index = self.model_selector.findData(recommended)
            if index >= 0:
                self.model_selector.setCurrentIndex(index)
                self.ai_manager.set_model(recommended)
        
        # Update status
        self.status.setText(f"Status: {self.ai_manager.get_model_info_text()}")
        self.status.setStyleSheet("color: #a6e3a1; font-size: 12px;")
        
        # Welcome message
        count = len(models)
        selected = self.ai_manager.current_model
        self.add_message("kh", f"Found {count} models. Selected: {selected}. Ready to chat.")
    


    def on_model_changed(self, display_name):
        """Handle model selection change"""
        model_name = self.model_selector.currentData()
        if model_name:
            self.ai_manager.set_model(model_name)
            self.status.setText(f"Status: Using {model_name}")
            self.status.setStyleSheet("color: #a6e3a1; font-size: 12px;")

    # === Voice Methods ===
    def toggle_voice(self):
        """Toggle voice input on/off"""
        if self.is_listening:
            self.stop_voice()
        else:
            self.start_voice()
    
    def start_voice(self):
        """Start listening for speech"""
        if self.voice_worker and self.voice_worker.isRunning():
            return
        
        # Update UI
        self.is_listening = True
        self.voice_button.setText("Stop")
        self.voice_button.setToolTip("Click to stop listening")
        self.status.setText("Status: Listening...")
        self.send_button.setEnabled(False)
        self.voice_button.setEnabled(True)
        
        # Start voice worker in background thread
        self.voice_worker = VoiceWorker()
        self.voice_worker.text_ready.connect(self.on_voice_text)
        self.voice_worker.error_occurred.connect(self.on_voice_error)
        self.voice_worker.start()
    
    def stop_voice(self):
        """Stop listening"""
        if self.voice_worker:
            self.voice_worker.stop()
        
        # Update UI
        self.is_listening = False
        self.voice_button.setText("Voice")
        self.voice_button.setToolTip("Click to speak")
        self.status.setText("Status: Ready")
        self.send_button.setEnabled(True)
        self.voice_button.setEnabled(True)
    
    def on_voice_text(self, text):
        """Called when speech is recognized"""
        self.stop_voice()
        self.input_box.setText(text)
        self.send_message()  # Auto-send
    
    def on_voice_error(self, error_message):
        """Called when voice recognition fails"""
        self.stop_voice()
        self.add_message("kh", f"Voice error: {error_message}")


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = kh_Assistant()
    window.show()
    sys.exit(app.exec())