"""
Align Whisper words with speaker diarization turns.

Strategy: give every word the speaker whose turns overlap it the most, tidy
the speaker changes so they fall on sentence boundaries, then join consecutive
words from the same speaker into readable paragraphs.
"""
from __future__ import annotations

from typing import List

from modules.models import (
    AlignedSegment,
    DiarizationSegment,
    TranscriptionSegment,
)

# A pause this long starts a new paragraph even if the speaker is the same.
MAX_PAUSE_SECONDS = 1.5
# After this long, a paragraph ends at the next sentence boundary.
PARAGRAPH_SECONDS = 20.0
SENTENCE_ENDINGS = (".", "?", "!", "।", "。")
# Word timestamps and turn edges are each off by a few tenths of a second, so
# up to this many words on either side of a speaker change may be mislabelled.
BOUNDARY_WORDS = 3


def align_transcript_with_speakers(
    transcription_segments: List[TranscriptionSegment],
    diarization_segments: List[DiarizationSegment],
) -> List[AlignedSegment]:
    """
    Merge word timestamps with speaker labels.

    *transcription_segments* are words as produced by TranscriberService
    (leading space kept), so their text is concatenated as-is.
    """
    words = transcription_segments
    speakers = [_find_dominant_speaker(w, diarization_segments) for w in words]
    _snap_to_sentences(words, speakers)

    aligned: list[AlignedSegment] = []
    for word, speaker in zip(words, speakers):
        current = aligned[-1] if aligned else None
        if (
            current
            and current.speaker == speaker
            and word.start - current.end < MAX_PAUSE_SECONDS
            and not (
                current.end - current.start > PARAGRAPH_SECONDS
                and _ends_sentence(current.text)
            )
        ):
            current.text += word.text
            current.end = word.end
        else:
            aligned.append(
                AlignedSegment(
                    start=word.start,
                    end=word.end,
                    speaker=speaker,
                    text=word.text,
                )
            )

    for seg in aligned:
        seg.text = seg.text.strip()
    return aligned


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------
def _ends_sentence(text: str) -> bool:
    return text.rstrip().endswith(SENTENCE_ENDINGS)


def _snap_to_sentences(words: List[TranscriptionSegment], speakers: List[str]) -> None:
    """
    Move speaker changes that land mid-sentence onto the nearest sentence
    boundary, editing *speakers* in place. People mostly change at sentence
    ends; a change one or two words off is timing error, not an interruption.

    ponytail: relies on Whisper's punctuation, so it does nothing for
    unpunctuated output, and a genuine one-word interruption without
    punctuation is folded into the surrounding speaker.
    """
    n = len(words)

    # 1. A few words of one speaker inside another speaker's sentence
    #    ("... prohibition, [such] as Andhra Pradesh").
    i = 0
    while i < n:
        j = i
        while j + 1 < n and speakers[j + 1] == speakers[i]:
            j += 1
        if (
            0 < i and j < n - 1
            and j - i + 1 < BOUNDARY_WORDS
            and speakers[i - 1] == speakers[j + 1]
            and not _ends_sentence(words[i - 1].text)
            and not _ends_sentence(words[j].text)
        ):
            speakers[i: j + 1] = [speakers[i - 1]] * (j - i + 1)
        i = j + 1

    # 2. A change just before or just after a sentence end.
    for i in range(1, n):
        previous, current = speakers[i - 1], speakers[i]
        if previous == current or _ends_sentence(words[i - 1].text):
            continue
        # The sentence finishes within the new speaker's first words: they
        # belong to the previous speaker ("... still is | banned. When was").
        run_end = i
        while run_end + 1 < n and speakers[run_end + 1] == current:
            run_end += 1
        head = next(
            (j for j in range(i, min(i + BOUNDARY_WORDS, run_end + 1)) if _ends_sentence(words[j].text)),
            None,
        )
        if head is not None and head < run_end:
            speakers[i: head + 1] = [previous] * (head + 1 - i)
            continue
        # The previous speaker's last words start a new sentence: they belong
        # to the new speaker ("... New Jersey. And I'm | Sheila in Texas").
        tail = i - 1
        while tail > 0 and speakers[tail - 1] == previous and not _ends_sentence(words[tail - 1].text):
            tail -= 1
        if (
            i - tail <= BOUNDARY_WORDS
            and tail > 0
            and speakers[tail - 1] == previous
        ):
            speakers[tail:i] = [current] * (i - tail)


def _find_dominant_speaker(
    t_seg: TranscriptionSegment,
    d_segments: List[DiarizationSegment],
) -> str:
    """
    Return the speaker with the most overlap for *t_seg*. A word that falls
    in a gap between turns goes to the nearest turn.
    """
    if not d_segments:
        return "Speaker 1"

    # ponytail: scans every turn for every word, O(words x turns). Fine up to
    # a few hours of audio; switch to a two-pointer sweep if it ever shows up.
    speaker_overlap: dict[str, float] = {}
    for d_seg in d_segments:
        overlap = _calculate_overlap(
            t_seg.start, t_seg.end, d_seg.start, d_seg.end
        )
        if overlap > 0:
            speaker_overlap[d_seg.speaker] = (
                speaker_overlap.get(d_seg.speaker, 0.0) + overlap
            )

    if speaker_overlap:
        return max(speaker_overlap, key=speaker_overlap.get)

    nearest = min(
        d_segments,
        key=lambda d: max(d.start - t_seg.end, t_seg.start - d.end),
    )
    return nearest.speaker


def _calculate_overlap(
    start1: float, end1: float, start2: float, end2: float
) -> float:
    """Return the overlap duration between two intervals (≥ 0)."""
    return max(0.0, min(end1, end2) - max(start1, start2))