import argparse
import json
import re
from dataclasses import dataclass
from itertools import groupby
from pathlib import Path


@dataclass(frozen=True)
class Turn:
    start: float
    end: float
    speaker: str


@dataclass
class Token:
    start: float
    end: float
    text: str


def load_tokens(path: Path) -> list[Token]:
    with path.open(encoding="utf-8") as file:
        transcript = json.load(file)

    tokens: list[Token] = []
    pending = ""
    # Whisper reserves all IDs from end-of-text onward for control tokens.
    end_of_text = 50256 + int(transcript["model"]["multilingual"])

    for segment in transcript["transcription"]:
        for raw_token in segment["tokens"]:
            if raw_token["id"] >= end_of_text:
                continue
            text = raw_token["text"]
            offsets = raw_token.get("offsets")

            if not text:
                continue
            if offsets is None:
                if tokens:
                    tokens[-1].text += text
                else:
                    pending += text
                continue

            tokens.append(
                Token(offsets["from"] / 1000, offsets["to"] / 1000, pending + text)
            )
            pending = ""

    if not tokens:
        raise ValueError("Whisper produced no timed text")

    return tokens


def load_turns(audio: Path, model: Path) -> list[Turn]:
    from pyannote.audio import Pipeline

    pipeline = Pipeline.from_pretrained(model)
    result = pipeline(str(audio))
    turns = [
        Turn(segment.start, segment.end, speaker)
        for segment, _, speaker in result.exclusive_speaker_diarization.itertracks(
            yield_label=True
        )
    ]
    if not turns:
        raise RuntimeError("Pyannote detected no speech")

    return turns


def token_speaker(token: Token, turns: list[Turn]) -> str:
    midpoint = (token.start + token.end) / 2

    def rank(turn: Turn) -> tuple[float, float]:
        overlap = max(0.0, min(token.end, turn.end) - max(token.start, turn.start))
        distance = max(turn.start - midpoint, midpoint - turn.end, 0.0)
        return overlap, -distance

    return max(turns, key=rank).speaker


def render(tokens: list[Token], turns: list[Turn]) -> str:
    speaker_numbers: dict[str, int] = {}
    rendered = []
    for speaker, group in groupby(
        tokens, key=lambda token: token_speaker(token, turns)
    ):
        number = speaker_numbers.setdefault(speaker, len(speaker_numbers) + 1)
        text = "".join(token.text for token in group)
        text = re.sub(r"\s+", " ", text).strip()
        if text:
            rendered.append(f"**Speaker {number}:**\n\n{text}")

    if not rendered:
        raise ValueError("Whisper produced no transcript text")
    return "\n\n".join(rendered) + "\n"


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--audio", required=True, type=Path)
    parser.add_argument("--model", required=True, type=Path)
    parser.add_argument("--transcript", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args()

    tokens = load_tokens(args.transcript)
    turns = load_turns(args.audio, args.model)
    args.output.write_text(render(tokens, turns), encoding="utf-8")


if __name__ == "__main__":
    main()