#!/mnt/storage/Development/dictation-service/.venv/bin/python
import os
import sys
import queue
import json
import time
import subprocess
import threading
import sounddevice as sd
from vosk import Model, KaldiRecognizer
import logging
import asyncio
import aiohttp
from openai import AsyncOpenAI
from enum import Enum
from dataclasses import dataclass
from typing import List, Optional
import pyttsx3
import numpy as np

# Setup logging
logging.basicConfig(
    filename="/home/universal/.gemini/tmp/428d098e581799ff7817b2001dd545f7b891975897338dd78498cc16582e004f/debug.log",
    level=logging.DEBUG,
)

# Configuration
SHARED_MODELS_DIR = os.path.expanduser("~/.shared/models/vosk-models")
MODEL_NAME = "vosk-model-en-us-0.22-lgraph"  # Faster model with good accuracy
MODEL_PATH = os.path.join(SHARED_MODELS_DIR, MODEL_NAME)
SAMPLE_RATE = 16000
BLOCK_SIZE = 4000  # Smaller blocks for lower latency
DICTATION_LOCK_FILE = "listening.lock"
CONVERSATION_LOCK_FILE = "conversation.lock"

# VLLM Configuration
VLLM_ENDPOINT = "http://127.0.0.1:8000/v1"
VLLM_MODEL = "Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4"
MAX_CONVERSATION_HISTORY = 10
TTS_ENABLED = True


class AppState(Enum):
    """Application states for dictation and conversation modes"""

    IDLE = "idle"
    DICTATION = "dictation"
    CONVERSATION = "conversation"


@dataclass
class ConversationMessage:
    """Represents a single conversation message"""

    role: str  # "user" or "assistant"
    content: str
    timestamp: float


class TTSManager:
    """Manages text-to-speech functionality"""

    def __init__(self):
        self.engine = None
        self.enabled = TTS_ENABLED
        self._init_engine()

    def _init_engine(self):
        """Initialize TTS engine"""
        if not self.enabled:
            return
        try:
            self.engine = pyttsx3.init()
            # Configure voice properties for more natural speech
            voices = self.engine.getProperty("voices")
            if voices:
                # Try to find a good voice
                for voice in voices:
                    if "english" in voice.name.lower() or "en_" in voice.id.lower():
                        self.engine.setProperty("voice", voice.id)
                        break
            self.engine.setProperty("rate", 150)  # Moderate speech rate
            self.engine.setProperty("volume", 0.8)
            logging.info("TTS engine initialized")
        except Exception as e:
            logging.error(f"Failed to initialize TTS: {e}")
            self.enabled = False

    def speak(self, text: str):
        """Speak text synchronously"""
        if not self.enabled or not self.engine or not text.strip():
            return

        try:
            self.engine.say(text)
            self.engine.runAndWait()
            logging.info(f"TTS spoke: {text[:50]}...")
        except Exception as e:
            logging.error(f"TTS error: {e}")


class VLLMClient:
    """Client for VLLM API communication"""

    def __init__(self, endpoint: str = VLLM_ENDPOINT):
        self.endpoint = endpoint
        self.client = AsyncOpenAI(api_key="vllm-api-key", base_url=endpoint)
        self._test_connection()

    def _test_connection(self):
        """Test connection to VLLM endpoint"""
        try:
            import requests

            response = requests.get(f"{self.endpoint}/models", timeout=2)
            if response.status_code == 200:
                logging.info(f"VLLM endpoint connected: {self.endpoint}")
            else:
                logging.warning(
                    f"VLLM endpoint returned status: {response.status_code}"
                )
        except Exception as e:
            logging.warning(f"VLLM endpoint test failed: {e}")

    async def get_response(self, messages: List[dict]) -> str:
        """Get AI response from VLLM"""
        try:
            response = await self.client.chat.completions.create(
                model=VLLM_MODEL, messages=messages, max_tokens=500, temperature=0.7
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            logging.error(f"VLLM API error: {e}")
            return "Sorry, I'm having trouble connecting right now."


class ConversationManager:
    """Manages conversation state and AI interactions with persistent context"""

    def __init__(self):
        self.conversation_history: List[ConversationMessage] = []
        self.persistent_history_file = "conversation_history.json"
        self.vllm_client = VLLMClient()
        self.tts_manager = TTSManager()
        self.is_speaking = False
        self.max_history = MAX_CONVERSATION_HISTORY
        self.load_persistent_history()

    def load_persistent_history(self):
        """Load conversation history from persistent storage"""
        try:
            if os.path.exists(self.persistent_history_file):
                with open(self.persistent_history_file, "r") as f:
                    data = json.load(f)
                    for msg_data in data:
                        message = ConversationMessage(
                            msg_data["role"], msg_data["content"], msg_data["timestamp"]
                        )
                        self.conversation_history.append(message)
                logging.info(
                    f"Loaded {len(self.conversation_history)} messages from persistent storage"
                )
        except Exception as e:
            logging.error(f"Error loading conversation history: {e}")
            self.conversation_history = []

    def save_persistent_history(self):
        """Save conversation history to persistent storage"""
        try:
            data = []
            for msg in self.conversation_history:
                data.append(
                    {
                        "role": msg.role,
                        "content": msg.content,
                        "timestamp": msg.timestamp,
                    }
                )
            with open(self.persistent_history_file, "w") as f:
                json.dump(data, f, indent=2)
            logging.info("Conversation history saved")
        except Exception as e:
            logging.error(f"Error saving conversation history: {e}")

    def add_message(self, role: str, content: str):
        """Add message to conversation history"""
        message = ConversationMessage(role, content, time.time())
        self.conversation_history.append(message)

        # Keep history within limits
        if len(self.conversation_history) > self.max_history:
            self.conversation_history = self.conversation_history[-self.max_history :]

        # Save to persistent storage
        self.save_persistent_history()

        logging.info(f"Added {role} message: {content[:50]}...")

    def get_messages_for_api(self) -> List[dict]:
        """Get conversation history formatted for API call"""
        messages = []

        # Add system prompt
        messages.append(
            {
                "role": "system",
                "content": "You are a helpful AI assistant in a voice conversation. Be concise and natural in your responses.",
            }
        )

        # Add conversation history
        for msg in self.conversation_history:
            messages.append({"role": msg.role, "content": msg.content})

        return messages

    async def process_user_input(self, text: str):
        """Process user input and generate AI response"""
        if not text.strip():
            return

        # Add user message
        self.add_message("user", text)

        # Show notification
        send_notification("🤖 Processing", "Thinking...", 2000)

        # Mark as speaking to prevent audio interruption
        self.is_speaking = True

        try:
            # Get AI response
            api_messages = self.get_messages_for_api()
            response = await self.vllm_client.get_response(api_messages)

            # Add AI response
            self.add_message("assistant", response)

            # Speak response
            if self.tts_manager.enabled:
                send_notification(
                    "🤖 AI Responding",
                    response[:50] + "..." if len(response) > 50 else response,
                    3000,
                )
                self.tts_manager.speak(response)
            else:
                send_notification("🤖 AI Response", response, 5000)

        except Exception as e:
            logging.error(f"Error processing user input: {e}")
            send_notification("❌ Error", "Failed to process your request", 3000)
        finally:
            self.is_speaking = False

    def start_conversation(self):
        """Start a new conversation session (maintains persistent context)"""
        send_notification(
            "🤖 Conversation Started",
            "Speak to talk with AI! Context: "
            + str(len(self.conversation_history))
            + " messages",
            4000,
        )
        logging.info(
            f"Conversation session started with {len(self.conversation_history)} messages of context"
        )

    def end_conversation(self):
        """End the current conversation session (preserves context for next call)"""
        send_notification(
            "🤖 Conversation Ended", "Context preserved for next call", 3000
        )
        logging.info("Conversation session ended (context preserved for next call)")

    def clear_all_history(self):
        """Clear all conversation history (for fresh start)"""
        self.conversation_history.clear()
        try:
            if os.path.exists(self.persistent_history_file):
                os.remove(self.persistent_history_file)
        except Exception as e:
            logging.error(f"Error removing history file: {e}")
        logging.info("All conversation history cleared")


# Global State (Legacy support)
is_listening = False
q = queue.Queue()
last_partial_text = ""
typing_thread = None
should_type = False

# New State Management
app_state = AppState.IDLE
conversation_manager = None

# Voice Activity Detection (simple implementation)
last_audio_time = 0
speech_threshold = 1.0  # seconds of silence before considering speech ended
last_speech_time = 0


def send_notification(title, message, duration=2000):
    """Sends a system notification"""
    try:
        subprocess.run(
            ["notify-send", "-t", str(duration), "-u", "low", title, message],
            capture_output=True,
            check=True,
        )
    except (FileNotFoundError, subprocess.CalledProcessError):
        pass


def download_model_if_needed():
    """Download model if needed"""
    if not os.path.exists(MODEL_PATH):
        logging.info(f"Model '{MODEL_PATH}' not found. Looking in shared directory...")

        # Check if model exists in shared models directory
        shared_model_path = os.path.join(SHARED_MODELS_DIR, MODEL_NAME)
        if os.path.exists(shared_model_path):
            logging.info(f"Found model in shared directory: {shared_model_path}")
            return

        logging.info(f"Model '{MODEL_NAME}' not found anywhere. Downloading...")
        try:
            # Download to shared models directory
            os.makedirs(SHARED_MODELS_DIR, exist_ok=True)
            subprocess.check_call(
                ["wget", f"https://alphacephei.com/vosk/models/{MODEL_NAME}.zip"],
                cwd=SHARED_MODELS_DIR,
            )
            subprocess.check_call(["unzip", f"{MODEL_NAME}.zip"], cwd=SHARED_MODELS_DIR)
            logging.info(f"Download complete. Model installed at: {MODEL_PATH}")
        except Exception as e:
            logging.error(f"Error downloading model: {e}")
            sys.exit(1)
    else:
        logging.info(f"Using model at: {MODEL_PATH}")


def audio_callback(indata, frames, time, status):
    """Enhanced audio callback with voice activity detection"""
    global last_audio_time

    if status:
        logging.warning(status)

    # Convert indata to a NumPy array for numerical operations
    indata_np = np.frombuffer(indata, dtype=np.int16)

    # Track audio activity for voice activity detection
    if app_state == AppState.CONVERSATION:
        audio_level = np.abs(indata_np).mean()
        if audio_level > 0.01:  # Simple threshold for speech detection
            last_audio_time = time.currentTime

    if app_state in [AppState.DICTATION, AppState.CONVERSATION]:
        q.put(bytes(indata))


def process_partial_text(text):
    """Process partial text based on current mode"""
    global last_partial_text

    if text and text != last_partial_text:
        last_partial_text = text

        if app_state == AppState.DICTATION:
            logging.info(f"💭 {text}")
            # Show brief notification for longer partial text
            if len(text) > 3:
                send_notification(
                    "🎤 Speaking", text[:50] + "..." if len(text) > 50 else text, 1000
                )
        elif app_state == AppState.CONVERSATION:
            logging.info(f"💭 [Conversation] {text}")


async def process_final_text(text):
    """Process final text based on current mode"""
    global last_partial_text

    if not text.strip():
        return

    formatted = text.strip()

    # Filter out spurious single words that are likely false positives
    if len(formatted.split()) == 1 and formatted.lower() in [
        "the",
        "a",
        "an",
        "uh",
        "huh",
        "um",
        "hmm",
    ]:
        logging.info(f"⏭️  Filtered out spurious word: {formatted}")
        return

    # Filter out very short results that are likely noise
    if len(formatted) < 2:
        logging.info(f"⏭️  Filtered out too short: {formatted}")
        return

    # Remove "the" from start and end of transcriptions (common Vosk false positive)
    words = formatted.split()
    spurious_words = {"the", "a", "an"}
    
    # Remove from start
    while words and words[0].lower() in spurious_words:
        removed = words.pop(0)
        logging.info(f"⏭️  Removed spurious word from start: {removed}")
    
    # Remove from end
    while words and words[-1].lower() in spurious_words:
        removed = words.pop()
        logging.info(f"⏭️  Removed spurious word from end: {removed}")
    
    if not words:
        logging.info(f"⏭️  Filtered out - only spurious words: {formatted}")
        return
    
    formatted = " ".join(words)
    formatted = formatted[0].upper() + formatted[1:] if formatted else formatted

    if app_state == AppState.DICTATION:
        logging.info(f"✅ {formatted}")
        send_notification(
            "🎤 Dictation",
            f"Typed: {formatted[:30]}{'...' if len(formatted) > 30 else ''}",
            2000,
        )

        # Type the text immediately
        try:
            subprocess.run(["ydotool", "type", formatted + " "])
            logging.info(f"📝 Typed: {formatted}")
        except Exception as e:
            logging.error(f"Error typing: {e}")
            send_notification(
                "❌ Typing Error", "Could not type text - check ydotool", 3000
            )

    elif app_state == AppState.CONVERSATION:
        logging.info(f"✅ [Conversation] User said: {formatted}")

        # Process through conversation manager
        if conversation_manager and not conversation_manager.is_speaking:
            await conversation_manager.process_user_input(formatted)

    # Clear partial text
    last_partial_text = ""


def continuous_audio_processor():
    """Enhanced background thread with conversation support"""
    recognizer = None
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    
    # Start the event loop in a separate thread
    def run_loop():
        loop.run_forever()
    
    loop_thread = threading.Thread(target=run_loop, daemon=True)
    loop_thread.start()

    while True:
        current_app_state = app_state

        if current_app_state != AppState.IDLE and recognizer is None:
            # Initialize recognizer when we start listening
            try:
                model = Model(MODEL_PATH)
                recognizer = KaldiRecognizer(model, SAMPLE_RATE)
                logging.info("Audio processor initialized")
            except Exception as e:
                logging.error(f"Failed to initialize recognizer: {e}")
                time.sleep(1)
                continue

        elif current_app_state == AppState.IDLE and recognizer is not None:
            # Clean up when we stop
            recognizer = None
            logging.info("Audio processor cleaned up")
            time.sleep(0.1)
            continue

        if current_app_state == AppState.IDLE:
            time.sleep(0.1)
            continue

        # Process audio when active - use shorter timeout for lower latency
        try:
            data = q.get(timeout=0.05)  # Reduced timeout for faster processing

            if recognizer:
                # Feed audio data to recognizer first
                if recognizer.AcceptWaveform(data):
                    # Final result available
                    result = json.loads(recognizer.Result())
                    final_text = result.get("text", "")
                    if final_text:
                        logging.info(f"🎯 Final result received: {final_text}")
                        # Run async processing
                        asyncio.run_coroutine_threadsafe(
                            process_final_text(final_text), loop
                        )
                else:
                    # Check for partial results
                    partial_result = recognizer.PartialResult()
                    if partial_result:
                        partial = json.loads(partial_result)
                        partial_text = partial.get("partial", "")
                        if partial_text:
                            process_partial_text(partial_text)
                
                # Process additional queued audio chunks if available (batch processing)
                try:
                    while True:
                        additional_data = q.get_nowait()
                        if recognizer.AcceptWaveform(additional_data):
                            result = json.loads(recognizer.Result())
                            final_text = result.get("text", "")
                            if final_text:
                                logging.info(f"🎯 Final result received (batch): {final_text}")
                                asyncio.run_coroutine_threadsafe(
                                    process_final_text(final_text), loop
                                )
                except queue.Empty:
                    pass  # No more data available

        except queue.Empty:
            continue
        except Exception as e:
            logging.error(f"Audio processing error: {e}")
            time.sleep(0.1)


def show_streaming_feedback():
    """Show visual feedback when dictation starts"""
    if app_state == AppState.DICTATION:
        send_notification(
            "🎤 Dictation Active",
            "Speak now - text will be typed into focused app!",
            4000,
        )
    elif app_state == AppState.CONVERSATION:
        send_notification("🤖 Conversation Active", "Speak to talk with AI!", 3000)


def main():
    global app_state, conversation_manager

    try:
        logging.info("Starting enhanced AI dictation service")

        # Initialize conversation manager
        conversation_manager = ConversationManager()

        # Model Setup
        download_model_if_needed()
        logging.info("Model ready")

        # Start audio processing thread
        audio_thread = threading.Thread(target=continuous_audio_processor, daemon=True)
        audio_thread.start()
        logging.info("Audio processor thread started")

        logging.info("=== Enhanced AI Dictation Service Ready ===")
        logging.info("Features: Dictation (Alt+D) + AI Conversation (Ctrl+Alt+D)")

        # Test VLLM connection
        send_notification(
            "🚀 AI Dictation Service",
            "Service ready! Press Ctrl+Alt+D to start AI conversation",
            5000,
        )

        # Open audio stream
        with sd.RawInputStream(
            samplerate=SAMPLE_RATE,
            blocksize=BLOCK_SIZE,
            dtype="int16",
            channels=1,
            callback=audio_callback,
        ):
            logging.info("Audio stream opened")

            while True:
                # Check lock files for state changes
                dictation_lock_exists = os.path.exists(DICTATION_LOCK_FILE)
                conversation_lock_exists = os.path.exists(CONVERSATION_LOCK_FILE)

                # Determine desired state
                # Priority: Dictation takes precedence over conversation when both locks exist
                if dictation_lock_exists:
                    desired_state = AppState.DICTATION
                elif conversation_lock_exists:
                    desired_state = AppState.CONVERSATION
                else:
                    desired_state = AppState.IDLE

                # Handle state transitions
                if desired_state != app_state:
                    old_state = app_state
                    app_state = desired_state

                    if app_state == AppState.DICTATION:
                        logging.info("[Dictation] STARTED - Enhanced streaming mode")
                        show_streaming_feedback()
                    elif app_state == AppState.CONVERSATION:
                        logging.info("[Conversation] STARTED - AI conversation mode")
                        conversation_manager.start_conversation()
                        show_streaming_feedback()
                    elif old_state != AppState.IDLE:
                        logging.info(f"[{old_state.value.upper()}] STOPPED")
                        if old_state == AppState.CONVERSATION:
                            conversation_manager.end_conversation()
                        elif old_state == AppState.DICTATION:
                            send_notification(
                                "🛑 Dictation Stopped", "Press Alt+D to resume", 2000
                            )

                # Sleep to prevent busy waiting
                time.sleep(0.05)

    except KeyboardInterrupt:
        logging.info("\nExiting...")
    except Exception as e:
        logging.error(f"Fatal error: {e}")


if __name__ == "__main__":
    main()
