# Originally from https://github.com/google/speaker-id/blob/master/DiarizationLM/diarizationlm/utils.py
# SPDX-FileCopyrightText: 2025 Quan Wang and Yiling Huang and Guanlong Zhao and Evan Clark and Wei Xia and Hank Liao
#
# SPDX-License-Identifier: Apache-2.0

import copy
import dataclasses
import json
import sys
from collections.abc import Generator, Sequence
from typing import Any, List, Optional

import numpy as np
import word_levenshtein as levenshtein
from scipy import optimize

PUNCTUATIONS = [",", ".", "_", "?", "!", "-", '"', "'"]


@dataclasses.dataclass
class PromptOptions:
    """Options for generating prompts."""

    # For prompt segmentation.
    emit_input_length: int = 896
    emit_target_length: int = 896

    # Prefix and suffix for prompt and completion.
    # As a reference, OpenAI finetuning API usually suggests:
    # - No prompt prefix
    # - Prompt suffix: " -> "
    # - Completion suffix: " END"
    prompt_prefix: str = ""
    prompt_suffix: str = " --> "
    completion_suffix: str = ""

    # How do we represent the speaker token.
    # We may consider shorter prefix for token efficiency.
    speaker_prefix: str = "<speaker:"
    speaker_suffix: str = ">"


def normalize_text(text: str) -> str:
    """Normalize text."""
    # Convert to lower case.
    text_lower = text.lower().strip()

    # Remove punctuations.
    words = text_lower.split()
    new_words = []
    for word in words:
        new_word = word
        for punc in PUNCTUATIONS:
            replaced = new_word.replace(punc, "")
            if len(replaced.split()) != 1:
                continue
            new_word = replaced
        new_words.append(new_word)
    return " ".join(new_words)


def speakers_transform(speakers: Sequence[str]) -> list[str]:
    """Transform list of speakers to be order based."""
    spk_map = {}
    index = 0
    for spk in speakers:
        if spk not in spk_map:
            index += 1
            spk_map[spk] = index
    return [str(spk_map[spk]) for spk in speakers]


def get_aligned_hyp_speakers(
    hyp_text: str,
    ref_text: str,
    ref_spk: str,
    print_debug_info: bool = False,
) -> str:
    """Align ref_text to hyp_text, then apply the alignment to ref_spk."""
    # Counters for insertions and deletions in hyp and ref text alignment.
    num_insertions, num_deletions = 0, 0

    # Get the alignment.
    _, align = levenshtein.levenshtein_with_edits(normalize_text(ref_text), normalize_text(hyp_text))

    ref_spk_list = ref_spk.split()
    hyp_spk_align = []

    # Apply the alignment on ref speakers.
    for i, j in align:
        if i == -1:
            # hyp has insertion
            hyp_spk_align.append("-1")
            num_insertions += 1
        elif j == -1:
            # hyp has deletion
            num_deletions += 1
            continue
        else:
            hyp_spk_align.append(ref_spk_list[i])
    hyp_spk_align = " ".join(hyp_spk_align)

    if print_debug_info:
        print("Number of insertions: ", num_insertions)
        print("Number of deletions: ", num_deletions)
        # This is not the traditional denominator of WER. Instead, this is
        # len(hyp) + len(ref) - len(SUB).
        print("Length of align pairs: ", len(align))
    return hyp_spk_align


def get_oracle_speakers(hyp_spk: str, hyp_spk_align: str) -> Sequence[int]:
    """Get the oracle speakers for hypothesis."""
    hyp_spk_list = [int(x) for x in hyp_spk.split()]
    hyp_spk_align_list = [int(x) for x in hyp_spk_align.split()]

    # Build cost matrix.
    max_spk = max(*hyp_spk_list, *hyp_spk_align_list)
    cost_matrix = np.zeros((max_spk, max_spk))
    for aligned, original in zip(hyp_spk_align_list, hyp_spk_list):
        cost_matrix[aligned - 1, original - 1] += 1

    # Solve alignment.
    row_index, col_index = optimize.linear_sum_assignment(cost_matrix, maximize=True)

    # Build oracle.
    hyp_spk_oracle = hyp_spk_list.copy()
    for i, align_value in enumerate(hyp_spk_align_list):
        if align_value == -1:
            # There are some missing words. In such cases, we just use the original
            # speaker for these words if possible.
            if hyp_spk_list[i] == -1:
                # If we don't have original speaker for missing words, just use the
                # previous speaker if possible.
                # This is useful for the update_hyp_text_in_utt_dict() function.
                hyp_spk_oracle[i] = 1 if i == 0 else hyp_spk_oracle[i - 1]
            continue
        if row_index[align_value - 1] != align_value - 1:
            raise ValueError(f"Alignment mismatch at index {align_value - 1}: expected {align_value - 1}, got {row_index[align_value - 1]}")
        hyp_spk_oracle[i] = col_index[align_value - 1] + 1

    return hyp_spk_oracle


# Transcript-Preserving Speaker Transfer (TPST)
def transcript_preserving_speaker_transfer(src_text: str, src_spk: str, tgt_text: str, tgt_spk: str) -> str:
    """Apply source speakers to target."""
    if len(tgt_text.split()) != len(tgt_spk.split()):
        raise ValueError("tgt_text and tgt_spk must have the same length")
    if len(src_text.split()) != len(src_spk.split()):
        raise ValueError("src_text and src_spk must have the same length")
    tgt_spk_align = get_aligned_hyp_speakers(
        hyp_text=tgt_text,
        ref_text=src_text,
        ref_spk=src_spk,
    )
    oracle_speakers = get_oracle_speakers(hyp_spk=tgt_spk, hyp_spk_align=tgt_spk_align)
    return " ".join([str(x) for x in oracle_speakers])


# We can use this to finetune LLM.
# Inputs (prompts): hyp diarized text
# Targets: hyp diarized text with oracle speakers
def ref_to_oracle(json_dict: dict[str, str]) -> str:
    """Apply reference speakers to hypothesis."""
    return transcript_preserving_speaker_transfer(
        src_text=json_dict["ref_text"],
        src_spk=json_dict["ref_spk"],
        tgt_text=json_dict["hyp_text"],
        tgt_spk=json_dict["hyp_spk"],
    )


# Similar to ref_to_oracle, but the opposite direction.
# We can use this to finetune LLM.
# Inputs (prompts): ref diarized text with degraded speakers
# Targets: ref diarized text
def hyp_to_degraded(json_dict: dict[str, str]) -> str:
    """Apply hypothesis speakers to reference."""
    return transcript_preserving_speaker_transfer(
        src_text=json_dict["hyp_text"],
        src_spk=json_dict["hyp_spk"],
        tgt_text=json_dict["ref_text"],
        tgt_spk=json_dict["ref_spk"],
    )


def create_diarized_text(
    word_labels: Sequence[str],
    speaker_labels: Sequence[str],
    use_new_line: bool = False,
    po: PromptOptions = PromptOptions(),
) -> str:
    """Create diarized text from words and speaker labels."""
    output = []
    previous_speaker = None
    for word, speaker in zip(word_labels, speaker_labels):
        if speaker != previous_speaker:
            if previous_speaker and use_new_line:
                output.append("\n")
            output.append(po.speaker_prefix + speaker + po.speaker_suffix)
        output.append(word)
        previous_speaker = speaker
    return " ".join(output)


def extract_text_and_spk(completions: str, po: PromptOptions, skip_meaningless_speaker: bool = True) -> tuple[str, str]:
    """Extract the text and spk from the completions string."""
    spk = "1"
    previous_spk = "1"
    result_text = []
    result_spk = []
    for word in completions.split():
        if word.startswith(po.speaker_prefix):
            if not word.endswith(po.speaker_suffix):
                word += po.speaker_suffix
            spk = word[len(po.speaker_prefix) : -len(po.speaker_suffix)]
            # Handle undefined behaviors of non-recognizable spk with a placeholder.
            try:
                spk_int = int(spk)
                if not spk or spk_int < 1 or spk_int > 10:
                    raise ValueError("Seeing unexpected word: ", word)
                previous_spk = spk
            except ValueError as exc:
                if skip_meaningless_speaker:
                    print("Skipping meaningless speaker token:", word)
                    spk = previous_spk
                else:
                    raise exc
        else:
            result_text.append(word)
            result_spk.append(spk)
    return " ".join(result_text), " ".join(result_spk)


def discard_empty_str_and_remove_boundary_white_space(
    inputs: List[str],
) -> List[str]:
    return [x.strip() for x in inputs if x.strip()]


@dataclasses.dataclass
class JsonUtteranceReader:
    """Read the json files and generate prompts and targets."""

    json_files: str  # Ignored if utt is given.
    text_field: str
    input_speaker_field: str
    target_speaker_field: str  # If not given, will skip targets.
    po: PromptOptions
    utt: dict[str, str] = dataclasses.field(default_factory=dict)
    seg_id: int = 0

    def generate_utts(self) -> Generator[dict[str, str], None, None]:
        """Generate an utterance from all json files."""
        if self.utt:
            yield self.utt
            return

        for json_file in self.json_files.split(","):
            with open(json_file, encoding="utf-8") as f:
                data_dict = json.load(f)
                yield from data_dict["utterances"]

    def generate_data_tuple(self) -> Generator[tuple[str, str, str], None, None]:
        """Generate uttid-prompt-target tuples."""
        for utt in self.generate_utts():
            yield from self.generate_data_tuple_for_utt(utt)

    def generate_data_dict(self) -> Generator[dict[str, str], None, None]:
        """Generate a dict that can be used for datasets.Dataset.from_generator."""
        for uttid, prompt, target in self.generate_data_tuple():
            yield {"uttid": uttid, "prompt": prompt, "target": target}

    def generate_data_tuple_for_utt(self, utt: dict[str, str]) -> Generator[tuple[str, str, str], None, None]:
        """Generate uttid-prompt-target tuples from a single utterance."""
        self.seg_id = 0
        utt_id = utt["utterance_id"]

        # Get the fields from the utterance.
        words = discard_empty_str_and_remove_boundary_white_space(utt[self.text_field].split(" "))
        p_speakers = discard_empty_str_and_remove_boundary_white_space(utt[self.input_speaker_field].split(" "))

        # Check for length mismatch between words and p_speakers
        if len(words) != len(p_speakers):
            raise ValueError(
                f"Mismatch between words and input speakers: {len(words)} words, {len(p_speakers)} input speakers in utterance {utt_id}."
            )

        if self.target_speaker_field:
            t_speakers = discard_empty_str_and_remove_boundary_white_space(utt[self.target_speaker_field].split(" "))

            # Check for length mismatch between words and t_speakers
            if len(words) != len(t_speakers):
                raise ValueError(
                    f"Mismatch between words and target speakers: {len(words)} words, {len(t_speakers)} target speakers in utterance {utt_id}."
                )
        else:
            t_speakers = []

        yield from self.generate_data_tuple_from_range(
            utt_id=utt_id, words=words, p_speakers=p_speakers, t_speakers=t_speakers, start=0, end=len(words)
        )

    def generate_data_tuple_from_range(self, *, utt_id, words, p_speakers, t_speakers, start, end) -> Generator[tuple[str, str, str], None, None]:
        """Generate uttid-prompt-target tuples from a range of words."""
        # Decide whether to call recursively from the estimated length.
        estimated_prompt_length = len(self.po.prompt_prefix) + len(" ".join(words[start:end])) + len(self.po.prompt_suffix)
        if estimated_prompt_length > self.po.emit_input_length or estimated_prompt_length > self.po.emit_target_length:
            yield from self.generate_data_tuple_from_range(
                utt_id=utt_id, words=words, p_speakers=p_speakers, t_speakers=t_speakers, start=start, end=(start + end) // 2
            )
            yield from self.generate_data_tuple_from_range(
                utt_id=utt_id, words=words, p_speakers=p_speakers, t_speakers=t_speakers, start=(start + end) // 2, end=end
            )
            return

        prompt = self.po.prompt_prefix
        previous_p_spk = ""
        target = ""
        previous_t_spk = ""

        # Main loop.
        for i in range(start, end):
            word = words[i]
            p_spk = p_speakers[i]
            if p_spk != previous_p_spk:
                if previous_p_spk:
                    prompt += " "
                prompt += self.po.speaker_prefix + p_spk + self.po.speaker_suffix
            prompt += " " + word
            previous_p_spk = p_spk

            if self.target_speaker_field:
                t_spk = t_speakers[i]
                if t_spk != previous_t_spk:
                    if previous_t_spk:
                        target += " "
                    target += self.po.speaker_prefix + t_spk + self.po.speaker_suffix
                target += " " + word
                previous_t_spk = t_spk

        prompt_id = utt_id + "_seg" + str(self.seg_id)
        prompt += self.po.prompt_suffix
        target += self.po.completion_suffix
        if len(prompt) <= self.po.emit_input_length and len(target) <= self.po.emit_target_length:
            yield (prompt_id, prompt, target)
            self.seg_id += 1
        else:
            yield from self.generate_data_tuple_from_range(
                utt_id=utt_id, words=words, p_speakers=p_speakers, t_speakers=t_speakers, start=start, end=(start + end) // 2
            )
            yield from self.generate_data_tuple_from_range(
                utt_id=utt_id, words=words, p_speakers=p_speakers, t_speakers=t_speakers, start=(start + end) // 2, end=end
            )


def generate_prompts(
    utt: dict[str, str],
    po: PromptOptions,
    text_field: str = "hyp_text",
    input_speaker_field: str = "hyp_spk",
) -> list[str]:
    """Generate a list of prompts for a given utt."""
    po_modified = copy.deepcopy(po)
    po_modified.emit_target_length = sys.maxsize
    reader = JsonUtteranceReader(
        json_files="",
        text_field=text_field,
        input_speaker_field=input_speaker_field,
        target_speaker_field="",
        po=po_modified,
        utt=utt,
    )
    prompts = []
    for _, prompt, _ in reader.generate_data_tuple():
        prompts.append(prompt)
    if len(prompts) > 1:
        for prompt in prompts:
            if len(prompt) < po.emit_input_length / 3:
                raise RuntimeError("Prompt too short: ", prompt)
    return prompts


def find_utt_dict(utt_id: str, data_dict: dict[str, Any]) -> Optional[dict[str, str]]:
    """Find a utt_dict with a speicifc utterance_id from data_dict."""
    for utt_dict in data_dict["utterances"]:
        if utt_dict["utterance_id"] == utt_id:
            return utt_dict
    return None


def update_hyp_text_in_utt_dict(input_utt_dict: dict[str, str], new_hyp_text) -> dict[str, str]:
    """Update the hyp_text of a json utt_dict.

    We also transfer its original hyp_spk to the new hyp_text.

    This is useful if we want to use USM ASR transcripts to replace the
    turn-to-diarize transcripts, as the WER of turn-to-diarize transcripts is too
    high.

    Args:
      input_utt_dict: the input utt_dict
      new_hyp_text: the new hyp_text

    Returns:
      the new utt_dict
    """
    utt_dict = copy.deepcopy(input_utt_dict)
    # We don't know the speakers for new_hyp_text, so just use -1 as initial
    # speakers.
    new_hyp_spk = transcript_preserving_speaker_transfer(
        src_text=utt_dict["hyp_text"],
        src_spk=utt_dict["hyp_spk"],
        tgt_text=new_hyp_text,
        tgt_spk=" ".join(["-1" for _ in new_hyp_text.split()]),
    )
    # Update the utt_dict.
    utt_dict["hyp_text"] = new_hyp_text
    utt_dict["hyp_spk"] = new_hyp_spk
    utt_dict["hyp_diarized_text"] = create_diarized_text(new_hyp_text.split(), new_hyp_spk.split())
    return utt_dict


def truncate_suffix_and_tailing_text(text: str, suffix: str) -> str:
    """Tailing text after suffix should be removed as well."""
    if suffix and suffix in text:
        return text[: text.find(suffix)]
    return text


def postprocess_completions_for_utt(
    *,
    utt: dict[str, Any],
    llm_text_field: str = "llm_text",
    llm_speaker_field: str = "llm_spk",
    transfered_llm_speaker_field: str = "hyp_spk_llm",
    hyp_text_field: str = "hyp_text",
    hyp_spk_field: str = "hyp_spk",
    po: PromptOptions = PromptOptions(),
) -> None:
    """Postprocess the LLM completions of an utterance json dict."""
    # Remove completion suffix if it exists.
    completions_list = []
    for completion in utt["completions"]:
        if po.completion_suffix and po.completion_suffix in completion:
            completion = truncate_suffix_and_tailing_text(completion, po.completion_suffix)
        completions_list.append(completion)
    completions = " ".join(completions_list).strip()

    # Extract text and speaker.
    utt[llm_text_field], utt[llm_speaker_field] = extract_text_and_spk(completions, po=po)
    # Tha TPST alignment on LLM output against recognized hypothesis text can be
    # considered as a postprocessing step to ensure the hypothesis text does not
    # change too much from the diarization baseline.
    # Note: this step can arguably be skipped and we directly use LLM output
    # for evaluation. The assumption is LLM does not change original text too
    # much. `update_sstable_with_speakers` should be updated accordingly if so.
    utt[transfered_llm_speaker_field] = transcript_preserving_speaker_transfer(
        src_text=utt[llm_text_field],
        src_spk=utt[llm_speaker_field],
        tgt_text=utt[hyp_text_field],
        tgt_spk=utt[hyp_spk_field],
    )


def transfer_llm_completion(
    llm_completion: str,
    hyp: str,
    po: PromptOptions = PromptOptions(),
) -> str:
    """Transfer the LLM completion text to use text from hypothesis."""
    llm_text, llm_speaker = extract_text_and_spk(llm_completion, po=po)
    hyp_text, hyp_speaker = extract_text_and_spk(hyp, po=po)
    transfered_llm_speaker = transcript_preserving_speaker_transfer(
        src_text=llm_text,
        src_spk=llm_speaker,
        tgt_text=hyp_text,
        tgt_spk=hyp_speaker,
    )
    transferred = create_diarized_text(
        word_labels=hyp_text.split(),
        speaker_labels=transfered_llm_speaker.split(),
        po=po,
    )
    return transferred
