import whisperx
import gc
import os

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

AUDIO_DIR = os.path.expanduser(
    "~/VRtual X Dropbox/Tim B. Frank/Apps/Tims Hermes/Projekte/battenfeld-cincinnati/2026-06-30_Audionotizen"
)
OUTPUT_DIR = os.path.expanduser("~/Transkripte/battenfeld-cincinnati")
os.makedirs(OUTPUT_DIR, exist_ok=True)

files = sorted([f for f in os.listdir(AUDIO_DIR) if f.endswith(".mp3")])
print(f"{len(files)} Dateien gefunden: {files}\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 fname in files:
    audio_path = os.path.join(AUDIO_DIR, fname)
    out_name = fname.replace(".mp3", ".txt")
    out_path = os.path.join(OUTPUT_DIR, out_name)

    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!")