import whisperx
import gc
import os
import sys

HF_TOKEN="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"
DEVICE = "cuda"
COMPUTE_TYPE = "float16"
BATCH_SIZE = 4

if len(sys.argv) < 2:
    print("Verwendung: python transkribiere_datei.py <audiodatei> [...]")
    sys.exit(1)

audio_files = sys.argv[1:]
print(f"{len(audio_files)} Datei(en) zu transkribieren\n")

print("Lade Whisper large-v3...")
model = whisperx.load_model("large-v3", DEVICE, compute_type=COMPUTE_TYPE, language="de")

print("Lade Diarization-Pipeline...")
diarize_model = whisperx.diarize.DiarizationPipeline(token=HF_TOKEN, device=DEVICE)

for audio_path in audio_files:
    base = os.path.splitext(audio_path)[0]
    out_path = base + ".txt"
    fname = os.path.basename(audio_path)

    print(f"\n{'='*60}")
    print(f"Verarbeite: {fname}")
    print(f"{'='*60}")

    audio = whisperx.load_audio(audio_path)

    print("  1/4 Transkription...")
    result = model.transcribe(audio, batch_size=BATCH_SIZE, language="de")

    print("  2/4 Alignment...")
    model_a, metadata = whisperx.load_align_model(language_code="de", device=DEVICE)
    result = whisperx.align(result["segments"], model_a, metadata, audio, DEVICE,
                            return_char_alignments=False)
    del model_a; gc.collect()

    print("  3/4 Speaker Diarization...")
    diarize_segments = diarize_model(audio)
    result = whisperx.assign_word_speakers(diarize_segments, result)

    print("  4/4 Speichere Transkript...")
    with open(out_path, "w", encoding="utf-8") as f:
        f.write(f"Transkript: {fname}\n")
        f.write("=" * 60 + "\n\n")
        current_speaker = None
        for seg in result["segments"]:
            speaker = seg.get("speaker", "UNBEKANNT")
            text = seg["text"].strip()
            start = seg["start"]
            end = seg["end"]
            if speaker != current_speaker:
                f.write(f"\n[{speaker}]\n")
                current_speaker = speaker
            f.write(f"  [{start:6.1f}s – {end:5.1f}s]  {text}\n")

    print(f"  -> Gespeichert: {out_path}")

print("\nAlle Dateien fertig!")