#!/usr/bin/env python3
"""
STT-Skript: Nimmt Audio vom ReSpeaker XVF3800 auf und transkribiert es
mit faster-whisper (lokal, CPU-only).

Features:
    - Ansage "Aufnahme laeuft" ueber den ReSpeaker-Lautsprecher (espeak-ng TTS)
    - Kontinuierliche Aufnahme bis das Stoppwort "Ende" erkannt wird
    - Echtzeit-Stoppworterkennung mit Whisper tiny-Modell
    - Finale Transkription mit Whisper medium-Modell (hoehere Genauigkeit)
    - LLM-Korrektur mit gemma3:latest (Ollama) fuer grammatikalisch korrekten Text

Verwendung:
    python3 stt_respeaker.py              # Aufnehmen bis "Ende" gesagt wird
    python3 stt_respeaker.py en           # Sprache Englisch, Stoppwort "end"
    python3 stt_respeaker.py --file audio.wav   # Vorhandene Datei transkribieren

Voraussetzungen:
    pip install faster-whisper sounddevice numpy scipy requests
    espeak-ng installiert
    ollama + gemma3:latest installiert

Beim ersten Aufruf werden die Whisper-Modelle automatisch heruntergeladen.
"""

import sys
import os
import wave
import time
import tempfile
import argparse
import subprocess
import numpy as np
import sounddevice as sd
from scipy.signal import resample_poly
from faster_whisper import WhisperModel
import requests

# Output-Buffering deaktivieren
sys.stdout.reconfigure(line_buffering=True)

# === Konfiguration ===========================================================
DEVICE_SR = 48000
WHISPER_SR = 16000
CHANNELS = 2
DTYPE = "float32"
MODEL_SIZE = "large-v3"  # Auf GPU: beste Genauigkeit, ~0.1s pro 30s Audio
STOPWORD_MODEL = "large-v3"  # Gleiche Modell auf GPU, keine zusaetzliche Ladezeit
DEVICE_COMPUTE = "cuda"  # RTX 3090
COMPUTE_TYPE = "float16"  # Beste GPU-Praezision
GPU_DEVICE_INDEX = 1  # Zweite RTX 3090 (erste ist durch Ollama belegt)
DEFAULT_LANG = "de"
DEVICE_KEYWORD = "XVF3800"
INITIAL_PROMPT = "Nachrichten, Politik, Friedrich Merz, CDU, CSU, SPD, Olaf Scholz, Markus Söder, Kanzler, Merz-Akte"

# LLM-Korrektur (Ollama)
LLM_MODEL = "mistral-small3.2:24b"
LLM_URL = "http://localhost:11434/api/chat"
LLM_SYSTEM_PROMPT = {
    "de": (
        "Du bist ein Korrektor fuer deutsche Sprachtranskripte. "
        "Korrigiere Grammatik, Rechtschreibung, korrigiere falsch verstandene Woerter und ergaenze fehlende Woerter "
        "zu einem sinnvollen Text. "
        "Achte auf korrekte Namen: Friedrich Merz (CDU), Markus Soeder (CSU), "
        "Olaf Scholz (SPD). "
        "Entferne das Stoppwort 'Ende' am Schluss falls vorhanden. "
        "Antworte NUR mit dem korrigierten Text, ohne Einleitung oder Kommentar."
    ),
    "en": (
        "You are a proofreader for English speech transcripts. "
        "Correct grammar, spelling, correct wrong unerstood words and add missing words to make coherent text. "
        "Remove the stopword 'end' at the end. "
        "Reply ONLY with the corrected text, no preamble."
    ),
}

STOPWORDS = {
    "de": ["ende", "end", "stop"],
    "en": ["end", "stop", "done", "finish"],
    "fr": ["fin", "stop", "termine"],
    "es": ["fin", "stop", "terminado"],
}

CHUNK_SEC = 2.0
MAX_DURATION = 60
VAD_THRESHOLD = 0.008  # Angepasst fuer ASR-Kanal (rechter Kanal)


def find_respeaker_device():
    """Sucht den ReSpeaker XVF3800 in der Liste der Audiogeraete."""
    devices = sd.query_devices()
    for i, dev in enumerate(devices):
        name_lower = dev["name"].lower()
        if (("xvf3800" in name_lower or "respeaker" in name_lower or "array" in name_lower)
                and dev["max_input_channels"] > 0):
            print(f"ReSpeaker gefunden: [{i}] {dev['name']}"
                  f" ({dev['max_input_channels']} In, "
                  f"{dev['max_output_channels']} Out)")
            return i
    print(f"Kein Geraet mit '{DEVICE_KEYWORD}' gefunden.")
    print("Verfuegbare Eingabegeraete:")
    for i, dev in enumerate(devices):
        if dev["max_input_channels"] > 0:
            print(f"   [{i}] {dev['name']}")
    return None


def play_announcement(text, device_idx, lang="de"):
    """Erzeugt TTS mit espeak-ng und spielt es ueber den ReSpeaker ab."""
    print(f"Ansage: '{text}'")

    tmp_wav = tempfile.NamedTemporaryFile(
        suffix=".wav", delete=False, prefix="tts_"
    )
    tmp_wav.close()

    espeak_lang = {"de": "de", "en": "en", "fr": "fr", "es": "es"}.get(lang, "de")

    subprocess.run(
        ["espeak-ng", "-v", espeak_lang, "-s", "150",
         "-w", tmp_wav.name, text],
        check=True, capture_output=True
    )

    tmp_loud = tempfile.NamedTemporaryFile(
        suffix=".wav", delete=False, prefix="tts_loud_"
    )
    tmp_loud.close()

    subprocess.run(
        ["ffmpeg", "-y", "-i", tmp_wav.name,
         "-filter:a", "volume=3.0",
         "-ar", str(DEVICE_SR), "-ac", str(CHANNELS),
         tmp_loud.name],
        check=True, capture_output=True
    )

    sink = ("alsa_output.usb-Seeed_Studio_reSpeaker_"
            "XVF3800_4-Mic_Array_114993700261100055-00.analog-stereo")
    subprocess.run(
        ["paplay", "--device", sink,
         "--volume", "1.5", tmp_loud.name],
        check=True, capture_output=True
    )

    os.unlink(tmp_wav.name)
    os.unlink(tmp_loud.name)
    print("   (Ansage beendet)\n")


def resample_to_16k(audio_48k_mono):
    return resample_poly(audio_48k_mono, 1, 3)


def stereo_to_mono(audio):
    """Extrahiert den rechten Kanal (ASR-Output des XVF3800).
    Der rechte Kanal ist der fuer Spracherkennung optimierte ASR-Output
    (AEC residual / auto-select beam). Der linke Kanal ist der verarbeitete
    Beamforming-Output. Beide zu mischen verschlechtert die Qualitaet."""
    if audio.ndim > 1:
        # Rechter Kanal = Index 1 (ASR-Output, optimiert fuer Spracherkennung)
        return audio[:, 1]
    return audio


def compute_rms(audio_float):
    return float(np.sqrt(np.mean(audio_float ** 2)))


def check_stopword(audio_16k, model, language, stopword_list):
    rms = compute_rms(audio_16k)
    if rms < VAD_THRESHOLD:
        return False, "", rms

    segments, _ = model.transcribe(
        audio_16k,
        language=language,
        beam_size=1,
        vad_filter=True,
        temperature=0.0,
    )
    text = " ".join(seg.text.strip().lower() for seg in segments)

    for word in stopword_list:
        if word in text:
            return True, text, rms
    return False, text, rms


def record_until_stopword(device_idx, language, stopword_list):
    """Nimmt kontinuierlich auf bis Stoppwort erkannt wird."""
    print(f"Lade Whisper-Modell '{STOPWORD_MODEL}' (GPU)...")
    sw_model = WhisperModel(
        STOPWORD_MODEL, device=DEVICE_COMPUTE, compute_type=COMPUTE_TYPE,
        device_index=GPU_DEVICE_INDEX,
    )
    print(f"Aufnahme laeuft — sage '{stopword_list[0]}' zum Stoppen.")
    print(f"   (Notbremse nach {MAX_DURATION}s)\n")

    chunk_samples = int(CHUNK_SEC * DEVICE_SR)
    all_audio_48k = []
    total_duration = 0.0

    stream = sd.InputStream(
        samplerate=DEVICE_SR,
        channels=CHANNELS,
        dtype=DTYPE,
        device=device_idx,
        blocksize=chunk_samples,
    )
    stream.start()

    try:
        while total_duration < MAX_DURATION:
            chunk, overflowed = stream.read(chunk_samples)
            if overflowed:
                print("  Buffer Overflow!")

            chunk_mono = stereo_to_mono(chunk)
            chunk_16k = resample_to_16k(chunk_mono)
            all_audio_48k.append(chunk_mono)

            total_duration += CHUNK_SEC
            rms = compute_rms(chunk_16k)

            if rms >= VAD_THRESHOLD:
                is_stop, text, rms = check_stopword(
                    chunk_16k, sw_model, language, stopword_list
                )
                bar = "#" * min(int(rms * 200), 30)
                status = f"  [{total_duration:5.1f}s] RMS:{rms:.3f} {bar}"
                if text:
                    status += f"  '{text}'"
                print(status)

                if is_stop:
                    print(f"\n>>> Stoppwort erkannt: '{text}'")
                    break
            else:
                print(f"  [{total_duration:5.1f}s] RMS:{rms:.3f} (Stille)")
    finally:
        stream.stop()
        stream.close()

    full_audio_48k = np.concatenate(all_audio_48k)
    full_audio_16k = resample_to_16k(full_audio_48k)

    # Software-Verstaerkung: Audio normalisieren auf -3 dBFS Peak
    # Der XVF3800 liefert sehr leises Signal, daher nachverstaerken
    peak = float(np.max(np.abs(full_audio_16k)))
    if peak > 0.001:
        target_peak = 0.7  # -3 dBFS
        gain = target_peak / peak
        full_audio_16k = full_audio_16k * gain
        full_audio_16k = np.clip(full_audio_16k, -1.0, 1.0)
        print(f"  Lautstaerke-Anpassung: Gain x{gain:.1f} "
              f"(Peak {peak:.3f} -> {target_peak:.3f})")
    else:
        print("  Warnung: Sehr leises Signal, keine Verstaerkung.")

    print(f"\nAufnahme beendet. Gesamtdauer: "
          f"{len(full_audio_16k)/WHISPER_SR:.1f}s\n")
    return full_audio_16k


def save_wav(audio_data, filename):
    audio_int16 = (audio_data * 32767).astype(np.int16)
    with wave.open(filename, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(WHISPER_SR)
        wf.writeframes(audio_int16.tobytes())


def transcribe(audio_data, language="de", model_size=MODEL_SIZE):
    """Transkribiert Audio mit faster-whisper."""
    print(f"Lade Whisper-Modell '{model_size}'"
          f" (device={DEVICE_COMPUTE}, compute={COMPUTE_TYPE})...")
    print("   (Beim ersten Aufruf wird das Modell heruntergeladen.)\n")

    model = WhisperModel(
        model_size,
        device=DEVICE_COMPUTE,
        compute_type=COMPUTE_TYPE,
        device_index=GPU_DEVICE_INDEX,
    )

    print("Transkribiere...\n")
    transcribe_kwargs = dict(
        language=language,
        beam_size=5,
        vad_filter=True,
        best_of=5,
        temperature=[0.0, 0.2, 0.4, 0.6],
    )
    if INITIAL_PROMPT:
        transcribe_kwargs["initial_prompt"] = INITIAL_PROMPT

    segments, info = model.transcribe(audio_data, **transcribe_kwargs)

    print(f"   Sprache erkannt: {info.language}"
          f" ({info.language_probability:.1%})")
    print(f"   Dauer: {info.duration:.1f}s\n")

    full_text = ""
    print("--- Whisper-Transkription (roh) ----------------------")
    for seg in segments:
        timestamp = f"[{seg.start:5.1f}s - {seg.end:5.1f}s]"
        print(f"  {timestamp}  {seg.text}")
        full_text += seg.text + " "
    print("------------------------------------------------------\n")

    return full_text.strip()


def llm_correct(text, language="de"):
    """Korrigiert den transkribierten Text mit einem lokalen LLM (Ollama Chat API)."""
    if not text:
        return text

    system_prompt = LLM_SYSTEM_PROMPT.get(language, LLM_SYSTEM_PROMPT["de"])

    print(f"LLM-Korrektur mit '{LLM_MODEL}' (Ollama)...\n")

    try:
        response = requests.post(
            LLM_URL,
            json={
                "model": LLM_MODEL,
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": text},
                ],
                "stream": False,
                "options": {
                    "temperature": 0.0,
                    "num_predict": 500,
                },
            },
            timeout=180,
        )
        response.raise_for_status()
        result = response.json().get("message", {}).get("content", "").strip()

        if result:
            print("--- LLM-korrigierter Text ----------------------------")
            print(f"  {result}")
            print("------------------------------------------------------\n")
            return result
        else:
            print("LLM lieferte leere Antwort — verwende Whisper-Text.\n")
            return text

    except requests.exceptions.RequestException as e:
        print(f"Fehler bei LLM-Korrektur: {e}")
        print("Verwende unkorrigierten Whisper-Text.\n")
        return text


def transcribe_file(filepath, language, model_size=MODEL_SIZE):
    """Transkribiert eine vorhandene WAV-Datei."""
    print(f"Lade Datei: {filepath}\n")
    with wave.open(filepath, "rb") as wf:
        frames = wf.readframes(wf.getnframes())
        audio_int16 = np.frombuffer(frames, dtype=np.int16)
        audio_float = audio_int16.astype(np.float32) / 32767.0
        file_sr = wf.getframerate()

    print(f"   Dauer: {len(audio_float) / file_sr:.1f}s")
    print(f"   Samplerate: {file_sr} Hz\n")

    if file_sr != WHISPER_SR:
        print(f"   Resample: {file_sr} -> {WHISPER_SR} Hz")
        from math import gcd
        g = gcd(file_sr, WHISPER_SR)
        up = WHISPER_SR // g
        down = file_sr // g
        audio_float = resample_poly(audio_float, up, down)

    whisper_text = transcribe(audio_float, language, model_size)

    if whisper_text:
        corrected = llm_correct(whisper_text, language)
        return corrected
    return whisper_text


def main():
    parser = argparse.ArgumentParser(
        description="STT mit ReSpeaker XVF3800 + faster-whisper + LLM-Korrektur"
    )
    parser.add_argument(
        "language",
        type=str,
        nargs="?",
        default=DEFAULT_LANG,
        help="Sprachcode: de, en, fr, es, ... (Standard: de)",
    )
    parser.add_argument(
        "--file",
        type=str,
        default=None,
        help="Vorhandene WAV-Datei transkribieren statt live aufzunehmen",
    )
    parser.add_argument(
        "--model",
        type=str,
        default=MODEL_SIZE,
        help=f"Whisper-Modell: tiny|base|small|medium|large-v3"
             f" (Standard: {MODEL_SIZE})",
    )
    parser.add_argument(
        "--no-llm",
        action="store_true",
        help="LLM-Korrektur ueberspringen (nur Whisper-Text ausgeben)",
    )
    args = parser.parse_args()

    print("=" * 55)
    print("  ReSpeaker XVF3800 -> Whisper + LLM-Korrektur")
    print("=" * 55)

    # Datei-Modus
    if args.file:
        if not os.path.isfile(args.file):
            print(f"Datei nicht gefunden: {args.file}")
            sys.exit(1)
        text = transcribe_file(args.file, args.language, args.model)
        if text:
            print(f"\n=== FINALES ERGEBNIS ===\n{text}\n")
        else:
            print("Keine Sprache erkannt.")
        return

    # Live-Aufnahme-Modus
    device_idx = find_respeaker_device()
    if device_idx is None:
        print("Kein ReSpeaker XVF3800 gefunden.")
        sys.exit(1)

    stopword_list = STOPWORDS.get(args.language, STOPWORDS["de"])

    # Ansage
    play_announcement("Aufnahme laeuft.", device_idx, args.language)

    # Aufnehmen bis Stoppwort
    audio = record_until_stopword(device_idx, args.language, stopword_list)

    # WAV zwischenspeichern
    tmp_wav = tempfile.NamedTemporaryFile(
        suffix=".wav", delete=False, prefix="respeaker_"
    )
    save_wav(audio, tmp_wav.name)
    print(f"Aufnahme gespeichert: {tmp_wav.name}\n")

    # Whisper-Transkription
    whisper_text = transcribe(audio, args.language, args.model)

    # LLM-Korrektur
    if whisper_text and not args.no_llm:
        final_text = llm_correct(whisper_text, args.language)
    else:
        final_text = whisper_text

    # Aufnahme NICHT loeschen — fuer spaetere Wiederverwendung
    print(f"Aufnahme behalten: {tmp_wav.name}\n")

    if final_text:
        print(f"\n=== FINALES ERGEBNIS ===\n{final_text}\n")
        print("Fertig!")
    else:
        print("Keine Sprache erkannt — versuche es mit lauterer Stimme.")


if __name__ == "__main__":
    main()