"""Dictation scoring (M2) — pure, network-free, reused by M5 pronunciation.

Compares what the learner typed against the reference sentence and produces a
word-level diagnosis. Design choices (ADR 0004 spirit):

- Tolerance is calibrated to what dictation *tests*: hearing and spelling the
  right word. Case, punctuation and spacing are ignored; contractions
  ("don't" = "do not") and US/UK spelling ("colour" = "color") are accepted
  (legitimate transcription variants, not listening errors); but the word
  itself counts ("ship" vs "sheep", "there" vs "their") — that is exactly what
  dictation must catch.
- Feedback is word-aligned (Levenshtein backtrace via jiwer): the reference is
  shown with substitutions, omissions and insertions marked, plus a WER and an
  error breakdown. Aligning two sequences of different lengths is the core bit.
"""

import re
from typing import Literal

from pydantic import BaseModel

# Minimal US/UK normalisation: fold common -our/-or, -ise/-ize, -re/-er, -ll-.
# Not exhaustive — a pragmatic set that covers the bulk of dictation cases.
_UK_US_SUFFIXES = (
    ("our", "or"),  # colour -> color
    ("ise", "ize"),  # realise -> realize
    ("isation", "ization"),
    ("yse", "yze"),  # analyse -> analyze
    ("tre", "ter"),  # centre -> center
)
_CONTRACTIONS = {
    "don't": "do not",
    "doesn't": "does not",
    "didn't": "did not",
    "won't": "will not",
    "can't": "cannot",
    "couldn't": "could not",
    "wouldn't": "would not",
    "shouldn't": "should not",
    "isn't": "is not",
    "aren't": "are not",
    "wasn't": "was not",
    "weren't": "were not",
    "haven't": "have not",
    "hasn't": "has not",
    "hadn't": "had not",
    "i'm": "i am",
    "you're": "you are",
    "we're": "we are",
    "they're": "they are",
    "it's": "it is",
    "that's": "that is",
    "i've": "i have",
    "i'll": "i will",
    "let's": "let us",
}

OpType = Literal["equal", "substitute", "delete", "insert"]


class WordOp(BaseModel):
    """One alignment operation. ref_word/hyp_word are None where not applicable."""

    op: OpType
    ref_word: str | None = None
    hyp_word: str | None = None


class DictationResult(BaseModel):
    wer: float
    hits: int
    substitutions: int
    deletions: int
    insertions: int
    reference_words: list[str]
    ops: list[WordOp]

    @property
    def is_perfect(self) -> bool:
        return self.substitutions == 0 and self.deletions == 0 and self.insertions == 0


def _expand_contractions(text: str) -> str:
    return " ".join(_CONTRACTIONS.get(token, token) for token in text.split())


def _fold_spelling(word: str) -> str:
    for uk, us in _UK_US_SUFFIXES:
        if word.endswith(uk):
            return word[: -len(uk)] + us
    return word


def normalize_words(text: str) -> list[str]:
    """Lowercase, drop punctuation, expand contractions, fold UK->US, split."""
    text = text.lower()
    text = re.sub(r"[^a-z0-9'\s]", " ", text)  # keep the apostrophe for contractions
    text = _expand_contractions(text)
    text = re.sub(r"[^a-z0-9\s]", " ", text)  # now drop stray apostrophes
    return [_fold_spelling(word) for word in text.split()]


def score_dictation(reference: str, hypothesis: str) -> DictationResult:
    """Word-level scoring with a tolerant normalisation. Pure, no I/O."""
    import jiwer

    ref_words = normalize_words(reference)
    hyp_words = normalize_words(hypothesis)

    # jiwer needs non-empty input; handle the degenerate cases explicitly.
    if not ref_words:
        msg = "reference is empty after normalisation"
        raise ValueError(msg)
    if not hyp_words:
        return DictationResult(
            wer=1.0,
            hits=0,
            substitutions=0,
            deletions=len(ref_words),
            insertions=0,
            reference_words=ref_words,
            ops=[WordOp(op="delete", ref_word=word) for word in ref_words],
        )

    output = jiwer.process_words([" ".join(ref_words)], [" ".join(hyp_words)])
    ops: list[WordOp] = []
    for chunk in output.alignments[0]:
        if chunk.type == "equal":
            for offset in range(chunk.ref_end_idx - chunk.ref_start_idx):
                word = ref_words[chunk.ref_start_idx + offset]
                ops.append(WordOp(op="equal", ref_word=word, hyp_word=word))
        elif chunk.type == "substitute":
            for offset in range(chunk.ref_end_idx - chunk.ref_start_idx):
                ops.append(
                    WordOp(
                        op="substitute",
                        ref_word=ref_words[chunk.ref_start_idx + offset],
                        hyp_word=hyp_words[chunk.hyp_start_idx + offset],
                    )
                )
        elif chunk.type == "delete":
            for offset in range(chunk.ref_end_idx - chunk.ref_start_idx):
                ops.append(WordOp(op="delete", ref_word=ref_words[chunk.ref_start_idx + offset]))
        elif chunk.type == "insert":
            for offset in range(chunk.hyp_end_idx - chunk.hyp_start_idx):
                ops.append(WordOp(op="insert", hyp_word=hyp_words[chunk.hyp_start_idx + offset]))

    return DictationResult(
        wer=output.wer,
        hits=output.hits,
        substitutions=output.substitutions,
        deletions=output.deletions,
        insertions=output.insertions,
        reference_words=ref_words,
        ops=ops,
    )