"""Eigenständiger Transkriptions-Prozess (wird per `python -m notarius.transcription.cli`
als Subprozess gestartet, siehe worker.py).

Läuft absichtlich als separater Prozess statt als Hintergrunddienst: Whisper- und
Diarization-Modelle werden erst hier geladen und geben beim Prozessende automatisch
ihren VRAM/RAM wieder frei. Notarius selbst hält also außerhalb einer aktiven
Transkription keine schweren ML-Modelle im Speicher.

Kommuniziert den Fortschritt über stdout in einem einfachen Zeilenprotokoll, das
worker.py parst:
    STAGE:<name>      - neue Verarbeitungsstufe erreicht
    ERROR:<message>   - nicht behebbarer Fehler, Prozess beendet sich mit Exitcode 1
"""

from __future__ import annotations

import argparse
import gc
import sys
from pathlib import Path

from notarius.keyring_store import get_hf_token
from notarius.transcription.formatting import format_transcript_markdown

DEVICE = "cuda"
COMPUTE_TYPE = "float16"
BATCH_SIZE = 4
WHISPER_MODEL_NAME = "large-v3"


def _emit_stage(name: str) -> None:
    print(f"STAGE:{name}", flush=True)


def _emit_error(message: str) -> None:
    print(f"ERROR:{message}", flush=True)


def transcribe(audio_path: Path, output_path: Path, title: str, language: str) -> None:
    # Import erst hier (statt auf Modulebene), damit `python -m notarius.transcription.cli --help`
    # bzw. der GUI-Prozess selbst nicht das schwere whisperx/torch mitimportieren muss.
    import whisperx

    hf_token = get_hf_token()
    if not hf_token:
        _emit_error(
            "Kein HuggingFace-Token hinterlegt (Settings). Für die Sprechererkennung "
            "wird ein Token mit akzeptierten Pyannote-Nutzungsbedingungen benötigt."
        )
        sys.exit(1)

    _emit_stage("loading_models")
    model = whisperx.load_model(WHISPER_MODEL_NAME, DEVICE, compute_type=COMPUTE_TYPE, language=language)
    diarize_model = whisperx.diarize.DiarizationPipeline(token=hf_token, device=DEVICE)

    audio = whisperx.load_audio(str(audio_path))

    _emit_stage("transcribing")
    result = model.transcribe(audio, batch_size=BATCH_SIZE, language=language)

    _emit_stage("aligning")
    model_a, metadata = whisperx.load_align_model(language_code=language, device=DEVICE)
    result = whisperx.align(
        result["segments"], model_a, metadata, audio, DEVICE, return_char_alignments=False
    )
    del model_a
    gc.collect()

    _emit_stage("diarizing")
    diarize_segments = diarize_model(audio)
    result = whisperx.assign_word_speakers(diarize_segments, result)

    _emit_stage("writing")
    markdown = format_transcript_markdown(title, result["segments"])
    output_path.write_text(markdown, encoding="utf-8")

    _emit_stage("done")


def main() -> None:
    parser = argparse.ArgumentParser(description="Notarius Transkriptions-Worker")
    parser.add_argument("audio_path", type=Path)
    parser.add_argument("output_path", type=Path)
    parser.add_argument("--title", default=None, help="Titel für die Transkript-Überschrift")
    parser.add_argument("--language", default="de")
    args = parser.parse_args()

    title = args.title or args.audio_path.stem

    try:
        transcribe(args.audio_path, args.output_path, title, args.language)
    except Exception as exc:  # noqa: BLE001 - jeder Fehler muss dem Elternprozess gemeldet werden
        _emit_error(str(exc))
        sys.exit(1)


if __name__ == "__main__":
    main()