from pydub import AudioSegment

class Diarizer:
    """
    Simple stereo-based diarization using RMS energy comparison.
    Assigns Speaker0 or Speaker1 for each segment and optionally
    merges consecutive segments from the same speaker.
    """

    def __init__(self, speaker_names=("Speaker0", "Speaker1"), merge_turns=True):
        self.speaker0, self.speaker1 = speaker_names
        self.merge_turns = merge_turns

    def diarize_segments(self, audio_file, segments):
        """
        audio_file: path to wav/mp3/flac etc
        segments: list of dict or objects with start, end, text
        """
        audio = AudioSegment.from_file(audio_file)

        # mono fallback → no diarization possible
        if audio.channels < 2:
            out = []
            for seg in segments:
                out.append({
                    "start": seg["start"],
                    "end": seg["end"],
                    "text": seg["text"].strip(),
                    "speaker": "Unknown"
                })
            return out

        l_channel, r_channel = audio.split_to_mono()

        diarized = []
        for seg in segments:
            start_ms = int(seg["start"] * 1000)
            end_ms = int(seg["end"] * 1000)

            l_seg = l_channel[start_ms:end_ms]
            r_seg = r_channel[start_ms:end_ms]

            speaker = self.speaker0 if l_seg.rms > r_seg.rms else self.speaker1

            diarized.append({
                "start": seg["start"],
                "end": seg["end"],
                "text": seg["text"].strip(),
                "speaker": speaker
            })

        # merge consecutive speaker turns if enabled
        if self.merge_turns:
            diarized = self._merge_turns(diarized)

        return diarized

    def _merge_turns(self, segments):
        """
        Merge consecutive segments from the same speaker into larger turns.
        """
        if not segments:
            return segments

        merged = []
        current = segments[0].copy()

        for seg in segments[1:]:
            if seg["speaker"] == current["speaker"]:
                current["end"] = seg["end"]
                current["text"] = current["text"].rstrip() + " " + seg["text"].lstrip()
            else:
                merged.append(current)
                current = seg.copy()

        merged.append(current)
        return merged
