#!/usr/bin/env python3
"""
live-whisper.py — Self-contained live transcription worker using openai-whisper.

Fallback for transcribe-cli on platforms where faster-whisper / ONNX fails
(e.g. linux-arm64). Auto-creates a Python venv and installs openai-whisper.

Protocol:
  stdin  — raw PCM16 audio (16kHz, mono, 16-bit signed LE)
  stdout — JSON lines:
    {"type":"status","message":"Installing dependencies..."}
    {"type":"status","message":"Loading model..."}
    {"type":"ready"}
    {"type":"transcript","text":"hello world","isFinal":false}
    {"type":"transcript","text":"hello world how are you","isFinal":true}
    {"type":"error","message":"..."}

Usage:
  arecord -f S16_LE -r 16000 -c 1 -t raw -q - | python3 live-whisper.py --model base

Based on the proven ARM-compatible approach from hydra/whisper_asr.
"""

import sys
import os
import json
import subprocess
import importlib
import struct
import time
import threading
from pathlib import Path

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

SCRIPT_DIR = Path(__file__).resolve().parent
VENV = SCRIPT_DIR / ".whisper-venv"
PY = VENV / "bin" / "python"
PIP = VENV / "bin" / "pip"

SAMPLE_RATE = 16000
CHANNELS = 1
SAMPLE_WIDTH = 2  # 16-bit
CHUNK_SECONDS = 3  # Transcribe every N seconds of audio
WINDOW_SECONDS = 10  # Transcribe last N seconds (sliding window)

# ---------------------------------------------------------------------------
# Output helpers (JSON lines to stdout)
# ---------------------------------------------------------------------------

def emit(event: dict):
    """Write a JSON event to stdout and flush."""
    sys.stdout.write(json.dumps(event) + "\n")
    sys.stdout.flush()


def emit_status(msg: str):
    emit({"type": "status", "message": msg})


def emit_error(msg: str):
    emit({"type": "error", "message": msg})


def emit_transcript(text: str, is_final: bool = False):
    emit({"type": "transcript", "text": text, "isFinal": is_final})

# ---------------------------------------------------------------------------
# Venv bootstrap (inspired by hydra/whisper_asr)
# ---------------------------------------------------------------------------

def _in_venv() -> bool:
    return sys.prefix != sys.base_prefix


def _ensure_venv():
    if VENV.exists():
        return
    emit_status("Creating Python venv for Whisper...")
    import venv
    venv.EnvBuilder(with_pip=True).create(str(VENV))
    # Upgrade pip quietly
    subprocess.check_call(
        [str(PY), "-m", "pip", "install", "--upgrade", "pip"],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
    )


def _ensure_deps():
    """Check and install missing dependencies."""
    need = []
    try:
        import numpy
    except ImportError:
        need.append("numpy")
    try:
        import whisper
    except ImportError:
        need.append("openai-whisper")

    if need:
        emit_status(f"Installing: {', '.join(need)}...")
        subprocess.check_call(
            [str(PIP), "install", *need],
            stdout=subprocess.DEVNULL, stderr=subprocess.PIPE,
        )
        # Force reimport
        for mod in ["numpy", "whisper"]:
            if mod in sys.modules:
                del sys.modules[mod]

# ---------------------------------------------------------------------------
# Bootstrap: ensure we're in the venv with deps
# ---------------------------------------------------------------------------

if not _in_venv():
    _ensure_venv()
    # Re-exec this script inside the venv
    os.execv(str(PY), [str(PY)] + sys.argv)

_ensure_deps()

# Now safe to import
import numpy as np

# ---------------------------------------------------------------------------
# Main transcription loop
# ---------------------------------------------------------------------------

def main():
    import argparse
    parser = argparse.ArgumentParser(description="Live Whisper transcription worker")
    parser.add_argument("--model", default="base", help="Whisper model size (tiny/base/small/medium/large)")
    parser.add_argument("--chunk-seconds", type=float, default=CHUNK_SECONDS, help="Transcribe interval")
    parser.add_argument("--window-seconds", type=float, default=WINDOW_SECONDS, help="Sliding window size")
    parser.add_argument("--language", default=None, help="Language code (e.g. en, es, fr). Auto-detect if omitted.")
    args = parser.parse_args()

    import whisper

    # Load model
    emit_status(f"Loading Whisper {args.model} model...")
    try:
        device = "cpu"
        try:
            import torch
            if torch.cuda.is_available():
                device = "cuda"
        except ImportError:
            pass

        model = whisper.load_model(args.model, device=device)
    except Exception as e:
        emit_error(f"Failed to load model: {e}")
        sys.exit(1)

    emit({"type": "ready"})

    # Audio buffer — accumulates PCM16 as float32 @ 16kHz
    audio_buf = np.zeros(0, dtype=np.float32)
    buf_lock = threading.Lock()
    chunk_bytes = int(args.chunk_seconds * SAMPLE_RATE * SAMPLE_WIDTH)
    window_samples = int(args.window_seconds * SAMPLE_RATE)
    last_text = ""
    running = True

    def read_stdin():
        """Read PCM16 from stdin in a background thread."""
        nonlocal audio_buf, running
        try:
            while running:
                data = sys.stdin.buffer.read(chunk_bytes)
                if not data:
                    break  # EOF
                # Convert PCM16 LE to float32 [-1, 1]
                samples = np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0
                with buf_lock:
                    audio_buf = np.concatenate([audio_buf, samples])
        except Exception:
            pass
        finally:
            running = False

    # Start stdin reader thread
    reader = threading.Thread(target=read_stdin, daemon=True)
    reader.start()

    try:
        while running:
            time.sleep(args.chunk_seconds)

            with buf_lock:
                if len(audio_buf) < SAMPLE_RATE:  # Need at least 1s of audio
                    continue
                # Take the last window_seconds of audio
                window = audio_buf[-window_samples:].copy() if len(audio_buf) > window_samples else audio_buf.copy()

            # Transcribe
            try:
                fp16 = (device == "cuda")
                result = model.transcribe(
                    window,
                    fp16=fp16,
                    language=args.language,
                    no_speech_threshold=0.6,
                    condition_on_previous_text=False,
                )
                text = result.get("text", "").strip()
                if text and text != last_text:
                    last_text = text
                    emit_transcript(text, is_final=False)
            except Exception as e:
                emit_error(f"Transcription error: {e}")

    except KeyboardInterrupt:
        pass

    # Final transcription of the full buffer
    with buf_lock:
        full_audio = audio_buf.copy()

    if len(full_audio) >= SAMPLE_RATE:
        try:
            fp16 = (device == "cuda")
            result = model.transcribe(
                full_audio,
                fp16=fp16,
                language=args.language,
            )
            text = result.get("text", "").strip()
            if text:
                emit_transcript(text, is_final=True)
        except Exception:
            pass

    running = False


if __name__ == "__main__":
    main()
