#!/usr/bin/env python3
"""STT-CLI – Client für stt_service.py (oder lokal ohne Dienst).

Nur Transkripte gehen auf stdout (eine Zeile pro Äußerung bzw. Datei); alle
Statusmeldungen auf stderr. Dadurch direkt in Skripten nutzbar:
    text=$(python stt_cli.py --once)

Beispiele:
    python stt_cli.py aufnahme.wav                 # Datei an den Dienst schicken
    python stt_cli.py *.mp3 --json                 # mehrere Dateien, JSON-Zeilen
    python stt_cli.py                              # Mikrofon → fortlaufend transkribieren
    python stt_cli.py --once                       # nach der ersten Äußerung beenden
    python stt_cli.py --until-enter                # Pausen erlaubt, Enter beendet → ein Text
    python stt_cli.py -d respeaker                 # ReSpeaker (rechter ASR-Kanal automatisch)
    python stt_cli.py --local aufnahme.wav         # ohne Dienst (lädt Modell selbst)
    python stt_cli.py --until-enter --llm-correct  # zusätzlich per lokalem LLM korrigieren
    python stt_cli.py --list-devices

Exit-Codes: 0 ok, 1 Fehler, 2 keine Sprache erkannt (--once/--timeout, Dateien)
"""

from __future__ import annotations

import argparse
import json
import os
import queue
import sys
import threading
import time
from pathlib import Path

DEFAULT_URL = os.environ.get("STT_URL", "http://127.0.0.1:9998")
BLOCK_SEC = 0.1


def status(msg: str) -> None:
    print(msg, file=sys.stderr, flush=True)


def emit(result: dict, as_json: bool, source: str | None = None) -> None:
    if as_json:
        if source:
            result = {"source": source, **result}
        print(json.dumps(result, ensure_ascii=False), flush=True)
    else:
        print(result.get("text", ""), flush=True)


class Output:
    """Gibt Ergebnisse in Reihenfolge aus — optional nach LLM-Korrektur (--llm-correct).

    Die Korrektur läuft in einem eigenen Thread, damit Aufnahme und Empfang nie warten.
    """

    def __init__(self, args):
        self.as_json = args.json
        self.verbose = args.verbose
        self.corrector = None
        self._executor = None
        if getattr(args, "llm_correct", False):
            from concurrent.futures import ThreadPoolExecutor

            import stt_llm

            cfg = stt_llm.LlmConfig.from_env()
            if args.llm_model:
                cfg.model = args.llm_model
            if args.llm_url:
                cfg.url = args.llm_url
            self.corrector = stt_llm.LlmCorrector(cfg)
            self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="llm")

    def warmup(self) -> None:
        """LLM laden, während der Nutzer schon spricht."""
        if self.corrector is not None:
            status(f"LLM-Korrektur aktiv: {self.corrector.config.model} (wird im Hintergrund geladen)")
            self._executor.submit(self.corrector.warmup)

    def put(self, result: dict, source: str | None = None) -> None:
        if self.corrector is None:
            self._emit(result, source)
        else:
            self._executor.submit(self._correct_and_emit, result, source)

    def _correct_and_emit(self, result: dict, source: str | None) -> None:
        raw = result.get("text", "")
        if raw:
            c = self.corrector.correct(raw)
            result = {**result, "text": c.text, "raw_text": raw, "llm": c.info()}
            if self.verbose or c.status in ("rejected", "error"):
                labels = {"corrected": "korrigiert", "unchanged": "unverändert",
                          "rejected": "verworfen, Rohtext behalten", "error": "Fehler, Rohtext behalten"}
                note = f" — {c.detail}" if c.detail else ""
                status(f"  LLM: {labels.get(c.status, c.status)} ({c.seconds:.2f}s){note}")
                if self.verbose and c.changed:
                    status(f"  roh: {raw}")
        self._emit(result, source)

    def _emit(self, result: dict, source: str | None) -> None:
        emit(result if self.as_json else {"text": result.get("text", "")}, self.as_json, source)

    def close(self) -> None:
        if self._executor is not None:
            self._executor.shutdown(wait=True)


def emit_joined(parts: list[dict], out: Output) -> str:
    """--until-enter: alle Abschnitte als EIN Ergebnis ausgeben."""
    text = " ".join(p["text"] for p in parts if p.get("text")).strip()
    out.put({"text": text, "parts": len(parts),
             "duration": round(sum(p.get("duration", 0.0) for p in parts), 2),
             "processing_s": round(sum(p.get("processing_s", 0.0) for p in parts), 3)})
    return text


def start_enter_watcher() -> threading.Event:
    """Setzt das Event, sobald auf dem Terminal Enter gedrückt wird."""
    pressed = threading.Event()

    def wait():
        try:
            sys.stdin.readline()
        except Exception:
            pass
        pressed.set()

    threading.Thread(target=wait, daemon=True).start()
    return pressed


# === Geräte ===================================================================

def list_devices() -> None:
    import sounddevice as sd

    default_in = sd.default.device[0]
    for i, dev in enumerate(sd.query_devices()):
        if dev["max_input_channels"] > 0:
            mark = "*" if i == default_in else " "
            print(f"{mark}[{i:>2}] {dev['name']}  ({dev['max_input_channels']} Kanäle, "
                  f"{int(dev['default_samplerate'])} Hz)")


SYSTEM_DEVICES = ("pipewire", "pulse")  # ALSA-Plugins, die dem System-Default folgen


def pulse_default_source() -> str:
    """Name der System-Standardquelle (PipeWire/PulseAudio), '' wenn unbekannt."""
    import subprocess

    try:
        return subprocess.run(["pactl", "get-default-source"], capture_output=True,
                              text=True, timeout=3).stdout.strip()
    except Exception:
        return ""


def resolve_input(spec: str | None) -> tuple[int | None, dict]:
    """None → System-Default über PipeWire/Pulse; Zahl → Index; Text → Namenssuche.

    PortAudios eigener Default ist auf PipeWire-Systemen oft die erste ALSA-Karte
    (nicht das eingestellte Mikrofon) — daher bevorzugt das 'pulse'-Gerät.
    """
    import sounddevice as sd

    devices = sd.query_devices()
    if spec is None:
        for name in SYSTEM_DEVICES:
            for i, dev in enumerate(devices):
                if dev["name"] == name and dev["max_input_channels"] > 0:
                    return i, dev
        idx = sd.default.device[0]
        return (None, sd.query_devices(kind="input")) if idx is None or idx < 0 else (idx, devices[idx])
    if spec.isdigit():
        idx = int(spec)
        if idx < len(devices) and devices[idx]["max_input_channels"] > 0:
            return idx, devices[idx]
        raise SystemExit(f"Gerät [{idx}] existiert nicht oder hat keine Eingänge (--list-devices)")
    for i, dev in enumerate(devices):
        if spec.lower() in dev["name"].lower() and dev["max_input_channels"] > 0:
            return i, dev
    raise SystemExit(f"Kein Eingabegerät mit '{spec}' gefunden (--list-devices)")


def is_respeaker(name: str) -> bool:
    return any(k in name.lower() for k in ("xvf3800", "respeaker"))


class Microphone:
    """Nimmt im Callback auf und legt s16le-Mono-Blöcke in eine Queue — der
    Audio-Thread wird nie durch Netzwerk oder Inferenz aufgehalten."""

    def __init__(self, device: str | None, channel: int | None):
        import sounddevice as sd

        self.index, info = resolve_input(device)
        self.name = info["name"]
        self.channels = min(int(info["max_input_channels"]), 2)
        if self.name in SYSTEM_DEVICES:
            # PipeWire resampelt selbst → direkt in Whisper-Rate aufnehmen
            self.rate = 16000
            source = pulse_default_source()
            if source:
                self.name = f"{self.name} → {source}"
        else:
            self.rate = int(info["default_samplerate"])
        if channel is None and is_respeaker(self.name) and self.channels >= 2:
            channel = 1  # XVF3800: rechter Kanal = ASR-optimierter Ausgang
        self.channel = channel
        self.blocks: queue.Queue[bytes] = queue.Queue(maxsize=600)  # ~60 s Puffer
        self.dropped = 0
        self._stream = sd.InputStream(
            device=self.index, samplerate=self.rate, channels=self.channels, dtype="int16",
            blocksize=int(self.rate * BLOCK_SEC), callback=self._callback,
        )

    def _callback(self, indata, frames, time_info, flags):
        import numpy as np

        data = indata[:, self.channel] if self.channel is not None else indata.mean(axis=1)
        try:
            self.blocks.put_nowait(np.ascontiguousarray(data, dtype="<i2").tobytes())
        except queue.Full:
            self.dropped += 1

    def __enter__(self):
        ch = f", Kanal {self.channel}" if self.channel is not None else ""
        status(f"Mikrofon: {self.name} ({self.rate} Hz, {self.channels} Kanäle{ch})")
        self._stream.start()
        return self

    def __exit__(self, *exc):
        self._stream.stop()
        self._stream.close()


# === Dienst-Modus =============================================================

def service_file(url: str, path: Path, args) -> dict:
    import requests

    params = {k: v for k, v in (("language", args.language), ("prompt", args.prompt)) if v}
    with open(path, "rb") as fh:
        r = requests.post(f"{url}/transcribe", params=params, data=fh, timeout=args.http_timeout)
    if r.status_code != 200:
        raise RuntimeError(f"HTTP {r.status_code}: {r.text[:300]}")
    return r.json()


def service_mic(url: str, args, out: Output) -> int:
    from websockets.sync.client import connect

    ws_url = url.replace("http://", "ws://").replace("https://", "wss://").rstrip("/") + "/ws/stream"
    got_text = threading.Event()
    stop = threading.Event()
    result = {"code": 0}
    parts: list[dict] = []  # --until-enter: Abschnitte sammeln statt sofort ausgeben

    with Microphone(args.device, args.channel) as mic, connect(ws_url, max_size=None) as ws:
        ws.recv()  # ready
        start = {"type": "start", "sample_rate": mic.rate, "language": args.language,
                 "prompt": args.prompt}
        if args.silence_ms:
            start["silence_ms"] = args.silence_ms
        ws.send(json.dumps(start))
        ws.recv()  # ready (neu konfiguriert)
        enter = start_enter_watcher() if args.until_enter else None
        status("Aufnahme läuft — Pausen sind erlaubt. [Enter] beendet." if enter
               else "Bereit — sprich. (Strg+C beendet)")

        def receiver():
            try:
                for raw in ws:
                    event = json.loads(raw)
                    etype = event.get("type")
                    if etype == "speech_start" and args.verbose:
                        status(f"  … Sprache erkannt ({event['t']:.1f}s)")
                    elif etype == "transcript":
                        if not event["text"]:
                            continue
                        if args.verbose:
                            status(f"  ({event['duration']:.1f}s Audio, {event['processing_s']:.2f}s Rechenzeit)")
                        got_text.set()
                        if enter is not None:
                            parts.append(event)
                            if args.verbose:
                                status(f"  + {event['text']}")
                            continue
                        out.put(event)
                        if args.once:
                            stop.set()
                            return
                    elif etype == "done" and enter is not None:
                        stop.set()  # nach Enter: alle Abschnitte sind da
                        return
                    elif etype == "error":
                        status(f"  Fehler vom Dienst: {event.get('detail')}")
                stop.set()  # Dienst hat die Verbindung regulär geschlossen
            except Exception as exc:  # Verbindung weg
                if not stop.is_set():
                    status(f"Verbindung beendet: {exc}")
                    result["code"] = 1
                stop.set()

        rx = threading.Thread(target=receiver, daemon=True)
        rx.start()
        deadline = time.monotonic() + args.timeout if args.timeout else None
        end_sent = False
        try:
            while not stop.is_set():
                timed_out = bool(deadline and time.monotonic() > deadline)
                if enter is not None and (enter.is_set() or timed_out):
                    if not end_sent:
                        status("Zeitlimit erreicht." if timed_out and not enter.is_set()
                               else "Aufnahme beendet — transkribiere den Rest …")
                        ws.send(json.dumps({"type": "end"}))  # laufenden Abschnitt abschließen
                        end_sent = True
                    stop.wait(0.2)  # auf "done" warten, kein Audio mehr senden
                    continue
                if timed_out:
                    status("Zeitlimit erreicht.")
                    break
                try:
                    ws.send(mic.blocks.get(timeout=0.2))
                except queue.Empty:
                    continue
        except KeyboardInterrupt:
            status("\nBeendet.")
        finally:
            if mic.dropped:
                status(f"Warnung: {mic.dropped} Audioblöcke verworfen (Netz zu langsam?)")
            if not stop.is_set():
                try:
                    ws.send(json.dumps({"type": "close"}))
                    rx.join(timeout=60)  # ausstehende Transkripte noch abholen
                except Exception:
                    pass
    if result["code"]:
        return result["code"]
    if enter is not None:
        return 0 if emit_joined(parts, out) else 2
    return 0 if got_text.is_set() or not (args.once or args.timeout) else 2


# === Lokaler Modus (ohne Dienst) ==============================================

def _local_engine(args):
    import stt_core as core

    cfg = core.EngineConfig.from_env()
    if args.model:
        cfg.model = cfg.cpu_model = args.model
    if args.compute_device:
        cfg.device = args.compute_device
    engine = core.SttEngine(cfg)
    status("Lade Modell …")
    engine.load()
    status(f"Modell {engine.active_model} auf {engine.active_device} bereit ({engine.load_seconds:.1f}s)")
    return engine


def local_files(paths: list[Path], args, out: Output) -> int:
    import stt_core as core

    engine = _local_engine(args)
    code = 2
    for path in paths:
        audio = core.decode_audio_bytes(path.read_bytes())
        tr = engine.transcribe(audio, language=args.language, initial_prompt=args.prompt)
        out.put(tr.to_dict(), str(path))
        code = 0 if tr.text else code
    return code


def local_mic(args, out: Output) -> int:
    from concurrent.futures import ThreadPoolExecutor

    import stt_core as core

    engine = _local_engine(args)
    vad = core.VadConfig.from_env()
    if args.silence_ms:
        vad.silence_ms = args.silence_ms
    segmenter = core.UtteranceSegmenter(vad)
    executor = ThreadPoolExecutor(max_workers=1)  # Inferenz getrennt von der Aufnahme
    done = threading.Event()
    got_text = threading.Event()
    futures = []  # --until-enter: Reihenfolge der Abschnitte erhalten

    def on_result(future):
        tr = future.result()
        if not tr.text:
            return
        if args.until_enter:
            got_text.set()
            if args.verbose:
                status(f"  + {tr.text}")
            return
        if not (args.once and got_text.is_set()):
            out.put(tr.to_dict(with_segments=False))
            got_text.set()
            if args.once:
                done.set()

    def submit(event):
        fut = executor.submit(engine.transcribe, event.audio, language=args.language,
                              initial_prompt=args.prompt)
        fut.add_done_callback(on_result)
        futures.append(fut)

    enter = None
    with Microphone(args.device, args.channel) as mic:
        resampler = core.StreamResampler(mic.rate)
        if args.until_enter:
            enter = start_enter_watcher()
            status("Aufnahme läuft — Pausen sind erlaubt. [Enter] beendet.")
        else:
            status("Bereit — sprich. (Strg+C beendet)")
        deadline = time.monotonic() + args.timeout if args.timeout else None
        try:
            while not done.is_set():
                if enter is not None and enter.is_set():
                    status("Aufnahme beendet — transkribiere den Rest …")
                    break
                if deadline and time.monotonic() > deadline:
                    status("Zeitlimit erreicht.")
                    break
                try:
                    block = mic.blocks.get(timeout=0.2)
                except queue.Empty:
                    continue
                audio = resampler.process(core.pcm16_to_float(block))
                for event in segmenter.feed(audio):
                    if event.kind == "speech_start" and args.verbose:
                        status(f"  … Sprache erkannt ({event.time_s:.1f}s)")
                    elif event.kind == "utterance":
                        submit(event)
        except KeyboardInterrupt:
            status("\nBeendet.")
        for event in segmenter.flush():
            submit(event)
    executor.shutdown(wait=True)
    if args.until_enter:
        parts = [f.result().to_dict(with_segments=False) for f in futures]
        return 0 if emit_joined(parts, out) else 2
    return 0 if got_text.is_set() or not (args.once or args.timeout) else 2


# === Main =====================================================================

def main() -> int:
    p = argparse.ArgumentParser(
        description="STT-CLI: Datei oder Mikrofon → Text (über stt_service oder lokal)",
        formatter_class=argparse.RawDescriptionHelpFormatter, epilog=__doc__.split("Beispiele:")[1],
    )
    p.add_argument("files", nargs="*", type=Path, help="Audiodateien (WAV, MP3, OGG, …); ohne = Mikrofon")
    p.add_argument("--url", default=DEFAULT_URL, help=f"Dienst-URL (Standard: {DEFAULT_URL}, env STT_URL)")
    p.add_argument("--local", action="store_true", help="ohne Dienst: Modell im Prozess laden")
    p.add_argument("-l", "--language", default=None, help="Sprachcode (de, en, …) oder 'auto'; Standard: Dienst-Default")
    p.add_argument("--prompt", default=None, help="Kontext/Eigennamen für Whisper")
    p.add_argument("--json", action="store_true", help="JSON-Zeilen statt reinem Text")
    p.add_argument("-d", "--device", default=None, help="Eingabegerät: Index oder Namensteil (Standard: System-Default)")
    p.add_argument("--channel", type=int, default=None, help="Kanal wählen (ReSpeaker: automatisch 1)")
    p.add_argument("--once", action="store_true", help="nach der ersten erkannten Äußerung beenden")
    p.add_argument("--until-enter", action="store_true",
                   help="Aufnahme mit beliebig langen Pausen, bis Enter gedrückt wird; Ausgabe als EIN Text")
    p.add_argument("--timeout", type=float, default=None, help="Mikrofon-Modus nach N Sekunden beenden")
    p.add_argument("--silence-ms", type=int, default=None, help="Stille bis Äußerungsende (Standard: 700)")
    p.add_argument("--model", default=None, help="nur --local: Whisper-Modell")
    p.add_argument("--compute-device", choices=["auto", "cuda", "cpu"], default=None, help="nur --local")
    p.add_argument("--http-timeout", type=float, default=300.0)
    p.add_argument("--llm-correct", action="store_true",
                   help="Ergebnis zusätzlich per lokalem LLM korrigieren (langsamer; Rohtext bleibt bei Zweifel)")
    p.add_argument("--llm-model", default=None, help="LLM-Modell (Standard: gemma4:12b, env STT_LLM_MODEL)")
    p.add_argument("--llm-url", default=None,
                   help="LLM-Server (Standard: http://127.0.0.1:11434 = Ollama, env STT_LLM_URL)")
    p.add_argument("--list-devices", action="store_true", help="Eingabegeräte auflisten")
    p.add_argument("-v", "--verbose", action="store_true", help="Statusdetails auf stderr")
    args = p.parse_args()

    if args.until_enter:
        if args.files or args.once:
            status("Fehler: --until-enter ist nur für das Mikrofon und nicht zusammen mit --once nutzbar")
            return 1
        if not sys.stdin.isatty():
            status("Fehler: --until-enter braucht ein Terminal für die Enter-Taste (stdin ist umgeleitet)")
            return 1

    if args.list_devices:
        list_devices()
        return 0
    out = Output(args)
    try:
        if args.files:
            missing = [f for f in args.files if not f.is_file()]
            if missing:
                status(f"Datei nicht gefunden: {', '.join(map(str, missing))}")
                return 1
            if args.local:
                return local_files(args.files, args, out)
            code = 2
            for path in args.files:
                result = service_file(args.url, path, args)
                out.put(result, str(path))
                code = 0 if result["text"] else code
            return code
        out.warmup()
        return local_mic(args, out) if args.local else service_mic(args.url, args, out)
    except KeyboardInterrupt:
        return 130
    except SystemExit:
        raise
    except Exception as exc:
        hint = ""
        if "Connection refused" in str(exc) or "ConnectionRefused" in type(exc).__name__:
            hint = f"\nLäuft der Dienst? (systemctl --user status stt-service, {args.url}/health) — oder --local nutzen."
        status(f"Fehler: {exc}{hint}")
        return 1
    finally:
        out.close()  # ausstehende LLM-Korrekturen noch ausgeben


if __name__ == "__main__":
    sys.exit(main())