"""Text-to-Speech voice assistant using Edge TTS - EMPATHETIC ARIA (NO STYLES)"""
import os
import time
import tempfile
import threading
import asyncio
from functools import lru_cache

import pygame
import edge_tts


# ========== Helpers & Defaults ==========

@lru_cache(maxsize=1)
def edge_voices_index():
    """
    Return {ShortName: voice_dict} for all Edge TTS voices.
    Cached so we don't hit the network repeatedly.
    """
    try:
        voices = asyncio.run(edge_tts.list_voices())
        return {v["ShortName"]: v for v in voices}
    except Exception as e:
        print(f"[TTS] ⚠️ Failed to load voices: {e}")
        return {}

# ONE consistent empathetic voice - Sara (natural, no modifications)
# Pure Sara voice for all emotions - her natural empathetic tone is perfect!
VOICE_MAP = {
    "Happy":   "en-US-SaraNeural",      # Sara's natural voice
    "Sad":     "en-US-SaraNeural",      # Sara's natural voice
    "Angry":   "en-US-SaraNeural",      # Sara's natural voice
    "Anxious": "en-US-SaraNeural",      # Sara's natural voice
    "Stressed":"en-US-SaraNeural",      # Sara's natural voice
    "Neutral": "en-US-SaraNeural",      # Sara's natural voice
}

def rate_pitch_for_emotion(emotion: str, intensity: float):
    """
    No modifications - Sara's natural voice is perfect as-is!
    Her empathetic tone is already ideal for emotional support.
    """
    # Return neutral for all emotions - let Sara's natural voice shine!
    return "+0%", "+0Hz"


def ssml_wrap(text: str, rate: str = "+0%", pitch: str = "+0Hz"):
    """
    Optional: wrap plain text in minimal SSML for nicer pacing.
    Edge TTS pitch format: +NHz or -NHz
    """
    return (
        f'<speak version="1.0" xml:lang="en-US">'
        f'<prosody rate="{rate}" pitch="{pitch}">'
        f'<break time="80ms"/>{text}<break time="60ms"/>'
        f'</prosody></speak>'
    )


# ========== VoiceAssistant ==========

class VoiceAssistant:
    """
    Empathetic voice assistant using Sara's natural voice.
    ONE consistent, warm, empathetic voice - no modifications needed!
    """

    def __init__(self, voice: str = "en-US-SaraNeural", rate: str = "+0%"):
        self.voice = voice
        self.rate = rate
        self.pitch = "+0Hz"

        self.counter = 0
        self.is_speaking = False
        self.speaking_lock = threading.Lock()
        self.audio_workers = []

        # Init audio output with robust error handling
        print("[TTS] 🔧 Initializing pygame mixer...")
        try:
            pygame.mixer.quit()  # Clean slate
            pygame.mixer.init(frequency=24000, size=-16, channels=2, buffer=2048)
            print(f"[TTS] ✅ Pygame mixer initialized (24kHz, 2ch, 2048 buffer)")
        except Exception as e:
            print(f"[TTS] ⚠️ Pygame init warning: {e}")
            try:
                pygame.mixer.init()
                print(f"[TTS] ✅ Pygame mixer initialized (default settings)")
            except Exception as e2:
                print(f"[TTS] ❌ CRITICAL: Pygame init failed: {e2}")
                print("[TTS] ❌ Audio will not work! Install pygame: pip install pygame")
        
        print(f"[TTS] ✅ Empathetic voice ready: {voice} (Sara - natural empathetic tone)")
        
        # Test Edge TTS connectivity at startup
        print("[TTS] 🌐 Testing Edge TTS connectivity...")
        try:
            voices = asyncio.run(edge_tts.list_voices())
            print(f"[TTS] ✅ Edge TTS connected - {len(voices)} voices available")
        except Exception as e:
            print(f"[TTS] ⚠️ Edge TTS connection issue: {e}")
            print("[TTS] ⚠️ Check internet connection or install: pip install edge-tts")

    # ----- Worker wiring -----

    def register_audio_worker(self, worker):
        """Register audio workers that should be paused during speech (to prevent echo)."""
        self.audio_workers.append(worker)
        worker_name = worker.__class__.__name__
        print(f"[TTS] ✅ Registered audio worker: {worker_name}")

    # ----- Voice controls -----

    def list_voices(self):
        """Return a list of available Edge TTS ShortName voices."""
        idx = edge_voices_index()
        print(f"[TTS] {len(idx)} voices available.")
        return list(idx.keys())

    def set_voice(self, short_name: str):
        """Set the Edge TTS voice by ShortName."""
        if short_name in edge_voices_index():
            self.voice = short_name
            print(f"[TTS] 🎙️ voice → {short_name}")
        else:
            print(f"[TTS] ⚠️ voice '{short_name}' not found; keeping {self.voice}")

    def set_rate(self, rate: str):
        """Set speech rate, e.g., '+10%' or '-5%'."""
        self.rate = rate
        print(f"[TTS] ⏩ rate → {rate}")

    def set_pitch(self, pitch: str):
        """Set speech pitch in Hz, e.g., '+10Hz' or '-5Hz'."""
        self.pitch = pitch
        print(f"[TTS] 🎚️ pitch → {pitch}")

    def apply_emotion_voice(self, emotion: str, intensity: float = 0.5):
        """
        Adapt Aria's voice to match user's emotion through rate/pitch.
        Call this right before speak/speak_async.
        """
        # Set voice (always Aria for consistency)
        self.set_voice(VOICE_MAP.get(emotion, VOICE_MAP["Neutral"]))
        
        # Set rate/pitch for emotional tone
        r, p = rate_pitch_for_emotion(emotion, intensity)
        self.set_rate(r)
        self.set_pitch(p)

    # ----- Playback control -----

    def stop(self):
        """
        Stop playback immediately (for barge-in).
        Also resumes any paused audio workers so mic listens again.
        """
        print("[TTS] 🛑 STOP called")
        try:
            pygame.mixer.music.stop()
            pygame.mixer.music.unload()
        except Exception as e:
            print(f"[TTS] Stop warning: {e}")
        
        with self.speaking_lock:
            self.is_speaking = False
        
        # Resume all workers
        for worker in self.audio_workers:
            if hasattr(worker, 'resume_listening'):
                try:
                    worker.resume_listening()
                except Exception as e:
                    print(f"[TTS] Resume error: {e}")

    # ----- Synthesis & Playback -----

    def _get_unique_filename(self, ext: str = ".mp3"):
        self.counter += 1
        return os.path.join(
            tempfile.gettempdir(),
            f"tts_{self.counter}_{int(time.time() * 1000)}{ext}"
        )

    async def _generate_speech(self, text: str, filename: str):
        """
        Synthesize speech with emotional tone through rate/pitch.
        Note: Edge-TTS no longer supports custom SSML styles.
        """
        try:
            print(f"[TTS] 🔧 Generating: voice={self.voice}, rate={self.rate}, pitch={self.pitch}")
            
            # Use ONLY voice, rate, pitch (NO style parameter)
            communicate = edge_tts.Communicate(
                text,
                voice=self.voice,
                rate=self.rate,
                pitch=self.pitch
            )
            await communicate.save(filename)
            
            # Verify file was created and has content
            if os.path.exists(filename):
                size = os.path.getsize(filename)
                if size > 0:
                    print(f"[TTS] ✅ Generated {size} bytes")
                    return True
                else:
                    print(f"[TTS] ❌ File created but empty!")
                    return False
            else:
                print(f"[TTS] ❌ File not created: {filename}")
                return False
                
        except Exception as e:
            print(f"[TTS] ❌ Generation error: {e}")
            import traceback
            traceback.print_exc()
            return False

    def _play_audio(self, filename: str):
        """Play audio file with pygame mixer"""
        try:
            if not os.path.exists(filename):
                print(f"[TTS] ❌ File doesn't exist: {filename}")
                return False
            
            print(f"[TTS] ▶️  Playing audio...")
            pygame.mixer.music.load(filename)
            pygame.mixer.music.play()
            
            # Wait for playback to complete
            while pygame.mixer.music.get_busy():
                pygame.time.Clock().tick(20)
            
            pygame.mixer.music.unload()
            print(f"[TTS] ✅ Playback complete")
            return True
            
        except Exception as e:
            print(f"[TTS] ❌ Playback error: {e}")
            import traceback
            traceback.print_exc()
            return False

    def speak(self, text: str):
        """
        Speak text synchronously (pauses audio workers to prevent echo).
        Accepts plain text or SSML (<speak>...</speak>).
        """
        if not text or not text.strip():
            print("[TTS] ⚠️ Empty text, skipping")
            return

        print(f"\n[TTS] 🔊 Speaking: '{text[:80]}...'")
        print(f"[TTS] 📊 Registered workers: {len(self.audio_workers)}")

        # Pause mic/ASR/etc. immediately when starting to speak
        paused_count = 0
        for worker in self.audio_workers:
            worker_name = worker.__class__.__name__
            if hasattr(worker, 'pause_listening'):
                try:
                    worker.pause_listening()
                    print(f"[TTS] ⏸️  Paused: {worker_name}")
                    paused_count += 1
                except Exception as e:
                    print(f"[TTS] ⚠️ Failed to pause {worker_name}: {e}")
            else:
                print(f"[TTS] ⚠️ Worker {worker_name} has no pause_listening() method!")
        
        if paused_count == 0 and len(self.audio_workers) > 0:
            print(f"[TTS] ⚠️ WARNING: No workers paused! Echo may occur.")

        with self.speaking_lock:
            self.is_speaking = True

        temp_file = self._get_unique_filename(".mp3")
        success = False

        try:
            # Generate speech
            print(f"[TTS] 🎤 Generating speech...")
            if asyncio.run(self._generate_speech(text, temp_file)):
                # Play audio
                if self._play_audio(temp_file):
                    success = True
                else:
                    print("[TTS] ❌ Playback failed")
                
                # Clean up temp file
                try:
                    if os.path.exists(temp_file):
                        os.remove(temp_file)
                        print(f"[TTS] 🗑️  Temp file removed")
                except Exception as e:
                    print(f"[TTS] ⚠️ Cleanup warning: {e}")
            else:
                print("[TTS] ❌ Speech generation failed")
                
        except Exception as e:
            print(f"[TTS] ❌ Error in speak(): {e}")
            import traceback
            traceback.print_exc()
        finally:
            with self.speaking_lock:
                self.is_speaking = False

            # Tiny safety delay, then resume listening
            time.sleep(0.2)
            
            print("[TTS] 🔄 Resuming workers...")
            resumed_count = 0
            for worker in self.audio_workers:
                worker_name = worker.__class__.__name__
                if hasattr(worker, 'resume_listening'):
                    try:
                        worker.resume_listening()
                        print(f"[TTS] ▶️  Resumed: {worker_name}")
                        resumed_count += 1
                    except Exception as e:
                        print(f"[TTS] ⚠️ Failed to resume {worker_name}: {e}")
            
            if success:
                print("[TTS] ✅ Speak completed successfully\n")
            else:
                print("[TTS] ⚠️ Speak completed with errors\n")

    def speak_async(self, text: str):
        """Speak text asynchronously in a separate thread."""
        threading.Thread(target=self.speak, args=(text,), daemon=True).start()

    def get_is_speaking(self) -> bool:
        """Thread-safe check if TTS is speaking."""
        with self.speaking_lock:
            return self.is_speaking