#!/usr/bin/env python

import argparse
import dataclasses

import numpy as np
from sklearn.cluster import AgglomerativeClustering
from sklearn.preprocessing import normalize
from speechbrain.dataio import audio_io
from speechbrain.dataio.preprocess import AudioNormalizer
from speechbrain.inference.speaker import EncoderClassifier
import torch
import whisper

@dataclasses.dataclass
class Segment:
    start: np.float32
    end: np.float32
    text: str

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(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("input")
    parser.add_argument("-o", "--output", default="output.srt",
                        help="Output subtitle (SRT) file path")
    parser.add_argument("-s", "--speakers", type=int, default=2,
                        help="Number of speakers in the audio")
    parser.add_argument("-m", "--model", default='turbo', choices=whisper.available_models(),
                        help="Whisper model to use")
    parser.add_argument("--gpu", action="store_true", help="run on GPU if available")
    args = parser.parse_args()


    has_gpu = torch.cuda.is_available()
    if args.gpu and not has_gpu:
        print("GPU not available. Falling back to CPU")
    dev = torch.device("cuda" if has_gpu and args.gpu else "cpu")

    model = whisper.load_model("turbo").to(dev)
    result = model.transcribe(args.input, word_timestamps=True, language="en")
    segments = []
    for seg in result["segments"]:
        words = seg["words"]
        start = words[0]["start"]
        end = words[-1]["end"]
        segments.append(Segment(start, end, seg["text"].strip()))

    print("Transcription done. Starting speaker classification")

    classifier = EncoderClassifier.from_hparams(
        source="speechbrain/spkrec-ecapa-voxceleb",
        run_opts={"device": str(dev)},
    )

    signal, sample_rate = audio_io.load(args.input, channels_first=False)
    audio_normalizer = AudioNormalizer()
    signal = audio_normalizer(signal, sample_rate)
    sample_rate = audio_normalizer.sample_rate

    embedding_segments = []
    embeddings = []
    for seg in segments:
        if seg.end - seg.start < 0.1:
            # Segment is too short to bother trying to classify
            continue

        start_sample = int(seg.start * sample_rate)
        end_sample = int(seg.end * sample_rate)

        chunk = signal[start_sample:end_sample]

        # normalize
        chunk /= chunk.abs().max()

        # SpeechBrain expects [batch, time]
        emb = classifier.encode_batch(chunk.unsqueeze(0).to(dev))
        emb = emb.squeeze().detach().cpu().numpy()

        embeddings.append(emb)
        embedding_segments.append(seg)
    embeddings = np.vstack(embeddings)
    embeddings = normalize(embeddings)

    n_speakers = 2
    clustering = AgglomerativeClustering(
        n_clusters=n_speakers,
        metric="cosine",
        linkage="average"
    )
    labels = clustering.fit_predict(embeddings)

    for seg, label in zip(embedding_segments, labels):
        seg.text = f"Speaker {label + 1}: {seg.text}"

    print("Final transcription:")
    with open(args.output, "w") as srt_file:
        for i, seg in enumerate(segments, 1):
            start = srt_timestamp(round(seg.start * 1000))
            end = srt_timestamp(round(seg.end * 1000))
            srt_file.write(f"{i}\n{start} --> {end}\n{seg.text}\n\n")
            print(start, end, seg.text)

if __name__ == "__main__":
    main()