"""
Decode helpers for Pidgin Whisper.

Two distinct text treatments:
  - postprocess()      : eval normalizer. Lowercases, strips punctuation,
                         merges digits — used to compute WER against the
                         lowercase/unpunctuated training references.
  - format_output()    : user-facing formatter. Restores punctuation
                         (commas, full stops, question marks) and applies
                         capitalization (sentence starts, "I", proper nouns)
                         on top of the model's raw lowercase output.

The model itself emits lowercase, unpunctuated text because that's how the
v1 training labels were written. format_output() is the formatting layer
(the same role the formatting step plays behind production ASR products):
the recognizer produces raw words, a separate model adds the punctuation
and casing.
"""
import re
from functools import lru_cache

# Domain priming context. Whisper's initial_prompt is read as prior speech
# context — natural-sounding, lowercase, no punctuation (matches v1 labels).
INITIAL_PROMPT = (
    "dis na bbc news pidgin tori about buhari tinubu atiku saraki "
    "tony nwoye femi otedola akinwunmi ambode oseloka henry obaze "
    "zainab balogun jimoh moshood and aisha for nigeria politics "
    "for states like lagos anambra delta kogi niger abuja kano rivers "
    "edo ogun salford and offa with organizations like apc pdp nema "
    "jamb frsc brt jp morgan and wikipedia pipo dey tok say di dey na "
    "wey pikin tori sabi hapun sometin anytin becos redi alredi neva "
    "dem una abi oga chillax snakebite "
    # everyday conversational / slang register (texting friends, gist)
    "how far my guy oyibo orobo omo abeg wahala gist sabi comot chop "
    "belle package sharp sharp fine girl fine boy no wahala wetin dey "
    "happen body dey inside cloth shey you dey alright sha gon make sense"
)

_DIGIT_PAIR = re.compile(r"(\d) (\d)")
_PUNCT = re.compile(r"[.,!?;:\"]")
_INTRA_NUM_COMMA = re.compile(r"(\d),(\d)")


# ----- proper-noun casing -----
# Multi-word entries are applied first (longest match wins).
_MULTIWORD_PROPER = {
    "femi fani kayode": "Femi Fani-Kayode",
    "oseloka henry obaze": "Oseloka Henry Obaze",
    "akinwunmi ambode": "Akinwunmi Ambode",
    "femi otedola": "Femi Otedola",
    "tony nwoye": "Tony Nwoye",
    "zainab balogun": "Zainab Balogun",
    "jimoh moshood": "Jimoh Moshood",
    "jp morgan": "JP Morgan",
}
_PROPER_NOUNS = {
    "nigeria": "Nigeria", "nigerian": "Nigerian", "nigerians": "Nigerians",
    "naija": "Naija", "africa": "Africa", "african": "African",
    "lagos": "Lagos", "abuja": "Abuja", "anambra": "Anambra", "delta": "Delta",
    "kogi": "Kogi", "kano": "Kano", "rivers": "Rivers", "edo": "Edo",
    "ogun": "Ogun", "oyo": "Oyo", "enugu": "Enugu", "imo": "Imo", "niger": "Niger",
    "yoruba": "Yoruba", "igbo": "Igbo", "hausa": "Hausa", "ghana": "Ghana",
    "buhari": "Buhari", "tinubu": "Tinubu", "atiku": "Atiku", "saraki": "Saraki",
    "obasanjo": "Obasanjo", "jonathan": "Jonathan", "aisha": "Aisha",
    "apc": "APC", "pdp": "PDP", "nema": "NEMA", "jamb": "JAMB",
    "frsc": "FRSC", "brt": "BRT", "bbc": "BBC", "efcc": "EFCC", "inec": "INEC",
    "wikipedia": "Wikipedia", "whatsapp": "WhatsApp", "facebook": "Facebook",
    "youtube": "YouTube", "google": "Google", "twitter": "Twitter",
    "salford": "Salford", "offa": "Offa",
}

# Sentence-initial words that usually signal a Pidgin/English question.
_Q_STARTERS = (
    "wetin", "wettin", "who", "wia", "where", "why", "how", "when",
    "which", "wich", "wheda", "shey", "abi",
)
_Q_STARTER_RE = re.compile(r"^(?:" + "|".join(_Q_STARTERS) + r")\b", re.IGNORECASE)

# Pidgin clause connectors that almost never legitimately *start* a sentence —
# the multilingual punctuator wrongly breaks before them, so we re-join.
_RELATIVIZERS = ("wey", "wia", "weh", "wae", "wey")
_FALSE_BREAK_RE = re.compile(
    r"\s*[.!?]\s+(" + "|".join(set(_RELATIVIZERS)) + r")\b", re.IGNORECASE
)


def _merge_digits(text: str) -> str:
    while True:
        new = _INTRA_NUM_COMMA.sub(r"\1\2", text)
        if new == text:
            break
        text = new
    while True:
        new = _DIGIT_PAIR.sub(r"\1\2", text)
        if new == text:
            break
        text = new
    return text


def postprocess(text: str) -> str:
    """Eval normalizer: lowercase-style, punctuation stripped, digits merged.
    Used to score WER against the lowercase/unpunctuated references."""
    text = text.lower()
    while True:
        new = _INTRA_NUM_COMMA.sub(r"\1\2", text)
        if new == text:
            break
        text = new
    text = _PUNCT.sub("", text)
    text = _merge_digits(text)
    return re.sub(r" +", " ", text).strip()


@lru_cache(maxsize=1)
def _punctuator():
    """Lazily load the punctuation-restoration model (downloaded once)."""
    from deepmultilingualpunctuation import PunctuationModel
    # xlm-roberta-base multilingual; lighter than the -large variant, fine for
    # short utterances on CPU. Restores . , ? - : ; on lowercased input.
    return PunctuationModel(model="kredor/punctuate-all")


def _truecase(text: str) -> str:
    # proper nouns (multi-word first, then single tokens)
    for k, v in _MULTIWORD_PROPER.items():
        text = re.sub(r"\b" + re.escape(k) + r"\b", v, text, flags=re.IGNORECASE)
    for k, v in _PROPER_NOUNS.items():
        text = re.sub(r"\b" + re.escape(k) + r"\b", v, text, flags=re.IGNORECASE)
    # standalone "i" -> "I"
    text = re.sub(r"\bi\b", "I", text)
    # capitalize the first alphabetic character of the whole string
    text = re.sub(r"^(\s*)([a-z])", lambda m: m.group(1) + m.group(2).upper(), text)
    # capitalize the first letter after sentence-ending punctuation
    text = re.sub(r"([.!?]\s+)([a-z])", lambda m: m.group(1) + m.group(2).upper(), text)
    return text


def _fix_questions(text: str) -> str:
    """Turn a trailing '.' into '?' for sentences that clearly start as questions."""
    pieces = re.split(r"([.!?])", text)
    out = []
    # pieces = [sentence, delim, sentence, delim, ...]
    for i in range(0, len(pieces) - 1, 2):
        sentence, delim = pieces[i], pieces[i + 1]
        if delim == "." and _Q_STARTER_RE.match(sentence.strip()):
            delim = "?"
        out.append(sentence + delim)
    if len(pieces) % 2 == 1:
        out.append(pieces[-1])
    return "".join(out)


def format_output(text: str) -> str:
    """User-facing formatter: restore punctuation + capitalization."""
    text = text.strip().lower()
    if not text:
        return ""
    text = _merge_digits(text)
    try:
        text = _punctuator().restore_punctuation(text)
    except Exception:
        # If the punctuation model is unavailable, at least end with a full stop.
        if text and text[-1] not in ".!?":
            text += "."
    # Drop colons/semicolons the model likes to insert after Pidgin "say"/"tok".
    text = re.sub(r"\s*[:;]\s*", " ", text)
    # Re-join sentences the model wrongly split before a Pidgin relativizer.
    text = _FALSE_BREAK_RE.sub(lambda m: " " + m.group(1), text)
    text = _fix_questions(text)
    text = _truecase(text)
    # ensure the utterance ends with terminal punctuation
    if text and text[-1] not in ".!?":
        text += "."
    return re.sub(r" +", " ", text).strip()


def transcribe(model, audio, *, use_hotwords: bool = True,
               use_postprocess: bool = True, pretty: bool = False,
               beam_size: int = 1) -> str:
    """Run faster-whisper with the decode pipeline.

    pretty=True  -> user-facing formatted text (punctuation + capitalization)
    pretty=False -> eval-normalized text (lowercase, no punctuation) when
                    use_postprocess is True; raw otherwise.
    beam_size    -> 1 (greedy, fastest) … 5 (more accurate, ~slower).
    """
    kwargs = {"language": "en", "task": "transcribe", "beam_size": beam_size}
    if use_hotwords:
        kwargs["initial_prompt"] = INITIAL_PROMPT
    segments, _ = model.transcribe(audio, **kwargs)
    text = "".join(s.text for s in segments).strip().lower()
    if pretty:
        return format_output(text)
    if use_postprocess:
        return postprocess(text)
    return text