#!/usr/bin/env python

import argparse

from pyannote.audio import Pipeline

def srt_timestamp(milis):
    h = milis // (1000 * 60 * 60)
    milis %= 1000 * 60 * 60
    m = milis // (1000 * 60)
    milis %= 1000 * 60
    s = milis // 1000
    milis %= 1000
    return f"{h:02}:{m:02}:{s:02},{milis:03}"

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("input")
    parser.add_argument("-o", "--output", default="output.srt")
    args = parser.parse_args()

    pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-community-1",
                                        token="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")
    output = pipeline(args.input, num_speakers=2)
    for turn, speaker in output.speaker_diarization:
        print(f"{speaker} speaks between t={turn.start:.3f}s and t={turn.end:.3f}s")

    with open(args.output, "w") as srt_file:
        for i, (turn, speaker) in enumerate(output.speaker_diarization, 1):
            srt_file.write(f"{i}\n"
                           f"{srt_timestamp(int(turn.start * 1000))} --> "
                           f"{srt_timestamp(int(turn.end * 1000))}\n"
                           f"{speaker}\n\n")

if __name__ == "__main__":
    main()