import asyncio
import logging
from backend.voice.tts import TTSPipeline

logger = logging.getLogger(__name__)

try:
    from openwakeword.model import Model
    # Try to load the interrupt model, fallback if it doesn't exist to prevent crash
    interrupt_model = Model(wakeword_models=["models/interrupt/quiet_stop_shutup.onnx"])
except Exception as e:
    logger.warning(f"JARVIS 10X: Failed to load openwakeword interrupt model: {e}")
    interrupt_model = None

class CoachingState:
    def __init__(self):
        self.muted = False
        
coaching_state = CoachingState()
tts = TTSPipeline()

async def listen_for_interrupt(audio_chunk: bytes):
    if not interrupt_model:
        return
        
    try:
        prediction = interrupt_model.predict(audio_chunk)
        if prediction.get("quiet_stop_shutup", 0.0) > 0.5:
            logger.info("JARVIS 10X: 'Quiet!' interrupt detected. Muting coaching.")
            await mute_coaching_immediately()
    except Exception as e:
        logger.error(f"Error in interrupt prediction: {e}")

async def mute_coaching_immediately():
    coaching_state.muted = True
    await stop_any_in_progress_tts()
    # Auto-unmute after a cooldown (e.g. 5 minutes)
    asyncio.create_task(schedule_unmute_after(seconds=300))

async def stop_any_in_progress_tts():
    # If TTS has a hard_stop method
    try:
        if hasattr(tts, "hard_stop"):
            await tts.hard_stop()
    except Exception as e:
        logger.error(f"Failed to hard stop TTS: {e}")

async def schedule_unmute_after(seconds: int):
    await asyncio.sleep(seconds)
    coaching_state.muted = False
    logger.info("JARVIS 10X: Coaching cooldown finished. Unmuted.")