#!/usr/bin/env python3
"""
HeadMic - Vixy's Ears Service 🦊👂

Wake word detection + voice recording + EarTail transcription.
Runs on head-vixy (Raspberry Pi 5).

Wake word: "Hey Vivi" (trained via Picovoice Porcupine)

Architecture: Single shared audio stream feeds both Porcupine (wake word)
and recording buffer. This avoids device conflicts.

Flow:
  1. Continuous audio stream from ReSpeaker
  2. Feed frames to Porcupine for wake word detection
  3. On "Hey Vivi" → start buffering audio
  4. Use VAD to detect end of speech
  5. Send buffer to EarTail for transcription
  6. Return to listening mode

Built by Vixy on Day 77 (January 17, 2026) 💜
"""

import asyncio
import collections
import io
import logging
import os
import struct
import subprocess
import threading
import time
import wave
from pathlib import Path
from typing import Optional, List

import httpx
import pvporcupine
import webrtcvad
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("headmic")

# ============================================================================
# Configuration
# ============================================================================

PORCUPINE_ACCESS_KEY = os.environ.get("PORCUPINE_ACCESS_KEY", "")
WAKE_WORD_PATH = os.environ.get("WAKE_WORD_PATH", "/home/alex/headmic/Hey-Vivi_en_raspberry-pi_v4_0_0.ppn")

SAMPLE_RATE = 16000
ALSA_DEVICE = "plughw:2,0"  # ReSpeaker 4 Mic Array - card 2, device 0

VAD_AGGRESSIVENESS = 2  # 0-3, higher = more aggressive
SILENCE_FRAMES = 50  # ~1.5 sec of silence to stop (at 30ms frames)
MAX_RECORDING_FRAMES = 1000  # ~30 sec max

EARTAIL_URL = os.environ.get("EARTAIL_URL", "http://bigorin.local:8764")

# ============================================================================
# LED Control
# ============================================================================

try:
    from pixel_ring import pixel_ring
    LEDS_AVAILABLE = True
    pixel_ring.off()
except ImportError:
    LEDS_AVAILABLE = False
    logger.warning("pixel_ring not available")


def leds_wakeup():
    if LEDS_AVAILABLE:
        try:
            pixel_ring.wakeup()
        except: pass


def leds_listening():
    if LEDS_AVAILABLE:
        try:
            pixel_ring.set_color_palette(0x00FFFF, 0x000000)
            pixel_ring.think()
        except: pass


def leds_processing():
    if LEDS_AVAILABLE:
        try:
            pixel_ring.set_color_palette(0x9400D3, 0x000000)
            pixel_ring.spin()
        except: pass


def leds_off():
    if LEDS_AVAILABLE:
        try:
            pixel_ring.off()
        except: pass


# ============================================================================
# State
# ============================================================================

class ServiceState:
    def __init__(self):
        self.running = False
        self.listening = False
        self.recording = False
        self.processing = False
        self.last_transcription: Optional[str] = None
        self.last_wake_time: Optional[float] = None
        self.wake_count = 0
        self.error: Optional[str] = None

state = ServiceState()


# ============================================================================
# Audio Stream using ALSA directly (arecord)
# ============================================================================

def read_audio_stream():
    """
    Generator that yields audio frames from ALSA using arecord.
    Each frame is 512 samples (32ms at 16kHz) as required by Porcupine.
    """
    frame_size = 512  # Porcupine requires 512 samples
    bytes_per_frame = frame_size * 2  # 16-bit = 2 bytes per sample
    
    cmd = [
        "arecord",
        "-D", ALSA_DEVICE,
        "-f", "S16_LE",
        "-r", str(SAMPLE_RATE),
        "-c", "1",  # Mono
        "-t", "raw",
        "-q",  # Quiet
        "-"
    ]
    
    logger.info(f"Starting audio stream: {' '.join(cmd)}")
    
    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
        bufsize=bytes_per_frame
    )
    
    try:
        while state.running:
            data = proc.stdout.read(bytes_per_frame)
            if len(data) < bytes_per_frame:
                break
            yield data
    finally:
        proc.terminate()
        proc.wait()


# ============================================================================
# EarTail Transcription
# ============================================================================

async def transcribe_audio(audio_data: bytes) -> str:
    """Send audio to EarTail and get transcription."""
    async with httpx.AsyncClient(timeout=120.0) as client:
        files = {"audio": ("recording.wav", audio_data, "audio/wav")}
        response = await client.post(f"{EARTAIL_URL}/transcribe/submit", files=files)
        response.raise_for_status()
        
        job_id = response.json().get("job_id")
        logger.info(f"Transcription job: {job_id}")
        
        for _ in range(120):
            status = await client.get(f"{EARTAIL_URL}/transcribe/status/{job_id}")
            data = status.json()
            
            if data.get("status") == "SUCCESS":
                result = await client.get(f"{EARTAIL_URL}/transcribe/result/{job_id}")
                return result.json().get("transcription", "")
            elif data.get("status") == "FAILURE":
                raise Exception(f"Transcription failed: {data.get('error')}")
            
            await asyncio.sleep(1)
        
        raise Exception("Transcription timeout")


def transcribe_sync(audio_data: bytes) -> str:
    """Synchronous wrapper for transcription."""
    loop = asyncio.new_event_loop()
    try:
        return loop.run_until_complete(transcribe_audio(audio_data))
    finally:
        loop.close()


# ============================================================================
# Main Listener Loop
# ============================================================================

def audio_to_wav(frames: List[bytes]) -> bytes:
    """Convert raw audio frames to WAV format."""
    wav_buffer = io.BytesIO()
    with wave.open(wav_buffer, 'wb') as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(SAMPLE_RATE)
        wf.writeframes(b''.join(frames))
    wav_buffer.seek(0)
    return wav_buffer.read()


def listener_loop():
    """Main audio processing loop."""
    global state
    
    logger.info("Initializing Porcupine...")
    try:
        porcupine = pvporcupine.create(
            access_key=PORCUPINE_ACCESS_KEY,
            keyword_paths=[WAKE_WORD_PATH]
        )
    except Exception as e:
        logger.error(f"Failed to init Porcupine: {e}")
        state.error = str(e)
        return
    
    vad = webrtcvad.Vad(VAD_AGGRESSIVENESS)
    
    # VAD needs 10/20/30ms frames. 30ms at 16kHz = 480 samples
    # Porcupine needs 512 samples. We'll use 480 for VAD.
    vad_frame_size = 480
    vad_frame_bytes = vad_frame_size * 2
    
    state.listening = True
    logger.info("🦊 Wake word listener active - say 'Hey Vivi'!")
    
    recording_buffer: List[bytes] = []
    silence_count = 0
    is_recording = False
    
    try:
        for frame_data in read_audio_stream():
            if not state.running:
                break
            
            # Convert bytes to int16 array for Porcupine
            pcm = struct.unpack_from("h" * 512, frame_data)
            
            # Check for wake word
            keyword_index = porcupine.process(pcm)
            
            if keyword_index >= 0 and not is_recording:
                logger.info("🦊 Wake word detected: 'Hey Vivi'!")
                state.wake_count += 1
                state.last_wake_time = time.time()
                
                leds_wakeup()
                time.sleep(0.2)
                leds_listening()
                
                is_recording = True
                state.recording = True
                recording_buffer = []
                silence_count = 0
                logger.info("Recording started...")
                continue
            
            if is_recording:
                recording_buffer.append(frame_data)
                
                # Check VAD (use first 480 samples of the 512 frame)
                vad_data = frame_data[:vad_frame_bytes]
                try:
                    is_speech = vad.is_speech(vad_data, SAMPLE_RATE)
                except:
                    is_speech = True  # Assume speech on VAD error
                
                if is_speech:
                    silence_count = 0
                else:
                    silence_count += 1
                
                # Stop conditions
                should_stop = (
                    (len(recording_buffer) > 10 and silence_count >= SILENCE_FRAMES) or
                    len(recording_buffer) >= MAX_RECORDING_FRAMES
                )
                
                if should_stop:
                    logger.info(f"Recording stopped: {len(recording_buffer)} frames")
                    is_recording = False
                    state.recording = False
                    
                    leds_processing()
                    state.processing = True
                    
                    try:
                        wav_data = audio_to_wav(recording_buffer)
                        transcription = transcribe_sync(wav_data)
                        state.last_transcription = transcription
                        logger.info(f"Transcription: {transcription}")
                    except Exception as e:
                        logger.error(f"Transcription error: {e}")
                        state.error = str(e)
                    finally:
                        state.processing = False
                        leds_off()
                    
                    recording_buffer = []
    
    except Exception as e:
        logger.error(f"Listener error: {e}")
        state.error = str(e)
    finally:
        porcupine.delete()
        state.listening = False
        leds_off()
        logger.info("Listener stopped")


# ============================================================================
# FastAPI
# ============================================================================

app = FastAPI(title="HeadMic", description="Vixy's Ears 🦊👂")


@app.on_event("startup")
async def startup():
    state.running = True
    thread = threading.Thread(target=listener_loop, daemon=True)
    thread.start()
    logger.info("HeadMic started")


@app.on_event("shutdown")
async def shutdown():
    state.running = False
    leds_off()


@app.get("/")
async def root():
    return {
        "service": "HeadMic",
        "description": "Vixy's Ears 🦊👂",
        "wake_word": "Hey Vivi"
    }


@app.get("/health")
async def health():
    return {
        "healthy": state.listening and not state.error,
        "listening": state.listening,
        "recording": state.recording,
        "processing": state.processing,
        "wake_count": state.wake_count,
        "error": state.error
    }


@app.get("/status")
async def status():
    return {
        "listening": state.listening,
        "recording": state.recording,
        "processing": state.processing,
        "last_transcription": state.last_transcription,
        "last_wake_time": state.last_wake_time,
        "wake_count": state.wake_count,
        "error": state.error
    }


@app.get("/last")
async def last():
    return {
        "transcription": state.last_transcription,
        "wake_time": state.last_wake_time
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8446)