import asyncio
import logging
import numpy as np
import io
import wave
import os
from datetime import datetime

# We reuse the audio_ws logic paths to keep audio storage centralized
from backend.voice.vad import SileroVAD
from backend.voice.stt import STTPipeline
from backend.ws.agent_ws import ws_manager

AUDIO_LOGS_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "voice", "audio_logs")
os.makedirs(AUDIO_LOGS_DIR, exist_ok=True)
TRANSCRIPT_LOG_PATH = os.path.join(AUDIO_LOGS_DIR, "transcripts_log.txt")

class PCMicService:
    def __init__(self):
        # Defer heavy ML init (torch.hub, WhisperModel) to run() so module import
        # doesn't block uvicorn startup under Tauri — avoids the ~60s health-check timeout.
        self.vad = None
        self.stt = None
        self.is_running = False
        self.interrupt_flag = False
        self.is_muted = False
        
    def mute(self):
        self.is_muted = True
        
    def unmute(self):
        self.is_muted = False
        self.interrupt() # Clear any garbage audio captured right before mute
        
    def interrupt(self):
        self.interrupt_flag = True

    async def run(self):
        try:
            import pyaudio
        except ImportError:
            logging.warning("PyAudio not installed. PC local microphone disabled. Run: pip install pyaudio")
            return

        # Lazy-init ML models in thread pool — keeps asyncio event loop free during heavy downloads
        if self.vad is None:
            self.vad = await asyncio.to_thread(SileroVAD)
        if self.stt is None:
            self.stt = await asyncio.to_thread(STTPipeline)

        self.is_running = True
        p = pyaudio.PyAudio()
        CHUNK = 512

        try:
            stream = p.open(format=pyaudio.paInt16,
                            channels=1,
                            rate=16000,
                            input=True,
                            frames_per_buffer=CHUNK)
        except Exception as e:
            logging.error(f"Failed to open PC Microphone: {e}")
            return

        logging.info("PC Local Microphone Service started. Listening continuously...")
        
        speech_buffer = bytearray()
        is_speaking = False
        silence_chunks = 0
        MAX_SILENCE_CHUNKS = 15

        while self.is_running:
            try:
                if self.interrupt_flag:
                    speech_buffer.clear()
                    is_speaking = False
                    silence_chunks = 0
                    self.interrupt_flag = False

                # Read 512 frames (1024 bytes) from the mic
                # Use a non-blocking asyncio sleep to yield to event loop
                await asyncio.sleep(0.001)
                
                # We read in chunks. For pure async we should use run_in_executor but this is simple for 512 frames
                data = stream.read(CHUNK, exception_on_overflow=False)
                
                if self.is_muted:
                    continue
                    

                audio_np = np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0
                
                vad_res = self.vad.process_chunk(audio_np)
                
                # --- JARVIS 10X Emergency Playback ---
                from backend.voice.emergency_playback import active_recording_session
                active_recording_session.append(data)
                
                if vad_res:
                    is_speaking = True
                    silence_chunks = 0
                    speech_buffer.extend(data)
                elif is_speaking:
                    speech_buffer.extend(data)
                    silence_chunks += 1
                    
                    if silence_chunks >= MAX_SILENCE_CHUNKS:
                        is_speaking = False
                        silence_chunks = 0
                        
                        if len(speech_buffer) > 16000: # at least 0.5s
                            logging.info("PC Mic utterance complete. Transcribing...")
                            wav_io = io.BytesIO()
                            with wave.open(wav_io, 'wb') as wf:
                                wf.setnchannels(1)
                                wf.setsampwidth(2)
                                wf.setframerate(16000)
                                wf.writeframes(speech_buffer)
                                
                            text = await self.stt.transcribe(wav_io.getvalue())
                            if text.strip():
                                timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
                                wav_filename = f"pc_mic_{timestamp}.wav"
                                wav_filepath = os.path.join(AUDIO_LOGS_DIR, wav_filename)
                                
                                with open(wav_filepath, "wb") as f:
                                    f.write(wav_io.getvalue())
                                    
                                with open(TRANSCRIPT_LOG_PATH, "a", encoding="utf-8") as f:
                                    f.write(f"[{timestamp}] [File: {wav_filename}] {text.strip()}\n")
                                    
                                logging.info(f"PC Mic transcript: {text.strip()}")
                                
                                from modules.assistant_identity import set_mode
                                lower_text = text.lower()
                                if "friday" in lower_text:
                                    set_mode("friday")
                                elif "jarvis" in lower_text:
                                    set_mode("jarvis")
                                    
                                if "jarvis" in lower_text or "friday" in lower_text:
                                    logging.info(f"Wake word detected from PC mic. Triggering Agent with: {text.strip()}")
                                    await ws_manager.handle_client_event(
                                        websocket=None,
                                        data={"type": "chat", "text": text.strip()}
                                    )
                                    
                        speech_buffer = bytearray()
            except asyncio.CancelledError:
                logging.info("PC Mic Loop cancelled, shutting down cleanly.")
                self.is_running = False
                break
            except Exception as e:
                logging.error(f"PC Mic read error: {e}")
                await asyncio.sleep(1)

        stream.stop_stream()
        stream.close()
        p.terminate()

pc_mic_service = PCMicService()

async def start_pc_mic_loop():
    import os
    if os.environ.get("CLOUD_ENV", "false").lower() == "true":
        return
    await pc_mic_service.run()