"""Map a canonical lyric sheet onto what the recording actually sings.

A lyric sheet is *canonical*: the chorus is written once, repeats are collapsed,
and unsung extra verses sometimes ride along. Karaoke needs the *performance*
sequence — the real order, with the chorus appearing as many times as it is sung.
Forced alignment cannot invent that: it consumes the reference in order, so a
chorus written once but sung three times leaves two thirds of the vocal to be
absorbed by whatever line happens to be adjacent (measured: 1.37 s mean line
error, worst 3.7 s, when a reference carried lines the recording never sang).

The two inputs have exactly complementary strengths:

    sheet       right words, wrong structure
    transcript  right structure, wrong words

So we use the transcript only to decide *which sheet line is being sung when*,
never for its words. That works even when the transcript is poor — measured CER
on real songs is 0.64, but token-overlap similarity still identifies the correct
sheet line, because picking one line out of ~30 needs far less signal than
reading it. This is why the mapper is worth more than a better ASR.

No network, no API key, no LLM: it is a similarity matrix plus a Viterbi pass
with a continuation bonus. See `resolve_with_llm` for where a model genuinely
helps (ambiguous sheets), which is a much smaller job than this one.
"""
from __future__ import annotations

import re
from typing import List, Tuple

# A sheet line only counts as "sung here" above this token-overlap score. Below
# it the transcript segment is an ad-lib, an instrumental mis-fire, or a line the
# sheet simply does not contain.
MIN_MATCH = 0.34
# Reward for continuing to the next sheet line, which disambiguates the common
# case of near-identical lines (a chorus whose lines differ by one word) without
# forbidding the backward jump that a chorus repeat *is*.
CONTINUE_BONUS = 0.22
# Cost of jumping backwards in the sheet, i.e. claiming a line is sung again.
# A *penalty*, not a reward: a repeat has to be earned by the similarity, because
# sheets legitimately contain the same chorus text twice and a second chorus
# reads as a backward jump otherwise. Swept against ground truth — at 0.0 two
# fixtures gained phantom repeats; at -0.10 both are exact and the real repeat is
# still found.
REPEAT_BONUS = -0.10


def _norm(s: str) -> str:
    s = s.lower().replace("ё", "е")
    s = re.sub(r"[^\w\s]|_", " ", s, flags=re.UNICODE)
    return re.sub(r"\s+", " ", s).strip()


def _tokens(s: str) -> List[str]:
    return _norm(s).split()


def _bigrams(word: str) -> set:
    w = f" {word} "
    return {w[i:i + 2] for i in range(len(w) - 1)}


def word_similarity(a: str, b: str) -> float:
    """Dice coefficient over character bigrams — tolerant of the one- or
    two-character errors that dominate sung ASR output."""
    if a == b:
        return 1.0
    ga, gb = _bigrams(a), _bigrams(b)
    if not ga or not gb:
        return 0.0
    return 2 * len(ga & gb) / (len(ga) + len(gb))


def line_similarity(hyp: str, ref: str) -> float:
    """Greedy token matching between two lines, 0…1.

    Token-level rather than character-level so that a transcript which gets a
    word wrong still scores the line it belongs to. Length-normalized against
    the *reference* so a long transcript run doesn't out-score a short line.
    """
    ht, rt = _tokens(hyp), _tokens(ref)
    if not ht or not rt:
        return 0.0
    used = [False] * len(ht)
    score = 0.0
    for rw in rt:
        best, bi = 0.0, -1
        for i, hw in enumerate(ht):
            if used[i]:
                continue
            s = word_similarity(rw, hw)
            if s > best:
                best, bi = s, i
        if bi >= 0 and best >= 0.5:
            used[bi] = True
            score += best
    return score / len(rt)


def map_performance(sheet: List[str], hyp_lines: List[dict],
                    min_match: float = MIN_MATCH) -> List[dict]:
    """Decide which sheet line each transcript segment is singing.

    `hyp_lines` are the transcript's timed lines ({startMs, endMs, text}).
    Returns one entry per transcript segment: the matched sheet index (or None),
    its score, and the segment's timing.

    Viterbi over sheet index, so the choice is made for the sequence as a whole
    rather than greedily per line — that is what lets a repeated chorus win over
    a locally-similar verse line.
    """
    n, m = len(hyp_lines), len(sheet)
    if not n or not m:
        return []

    sim = [[line_similarity(h["text"], s) for s in sheet] for h in hyp_lines]

    NONE = m                                  # an extra state: "matches nothing"
    best = [[float("-inf")] * (m + 1) for _ in range(n)]
    back = [[-1] * (m + 1) for _ in range(n)]
    for j in range(m):
        best[0][j] = sim[0][j]
    best[0][NONE] = min_match * 0.999          # ...just under any real match

    for i in range(1, n):
        for j in range(m + 1):
            emit = min_match * 0.999 if j == NONE else sim[i][j]
            for pj in range(m + 1):
                if best[i - 1][pj] == float("-inf"):
                    continue
                bonus = 0.0
                if j != NONE and pj != NONE:
                    if j == pj + 1:
                        bonus = CONTINUE_BONUS      # running through a section
                    elif j < pj:
                        bonus = REPEAT_BONUS        # jumped back: a repeat
                v = best[i - 1][pj] + emit + bonus
                if v > best[i][j]:
                    best[i][j] = v
                    back[i][j] = pj

    j = max(range(m + 1), key=lambda k: best[n - 1][k])
    path = [j]
    for i in range(n - 1, 0, -1):
        j = back[i][j]
        path.append(j)
    path.reverse()

    out = []
    for i, j in enumerate(path):
        matched = j != NONE and sim[i][j] >= min_match
        out.append({
            "startMs": hyp_lines[i]["startMs"],
            "endMs": hyp_lines[i]["endMs"],
            "sheetIdx": j if matched else None,
            "score": round(sim[i][j], 3) if j != NONE else 0.0,
            "hyp": hyp_lines[i]["text"],
        })
    return out


def expand_reference(sheet: List[str], hyp_lines: List[dict],
                     min_match: float = MIN_MATCH) -> Tuple[List[str], List[dict]]:
    """Build the reference the aligner should actually be given.

    Returns `(lines, plan)` where `lines` is the sheet rewritten in performance
    order — a chorus sung twice appears twice — and `plan` is the mapping detail.

    **Strictly additive: no sheet line is ever dropped.** The mapper's recall is
    bounded by the transcript's, and the transcript is poor — on a fixture where
    all 16 sheet lines are sung, the ASR produced 12 usable segments, so a
    "drop what wasn't matched" rule deleted 8 lines that really were sung. Adding
    a repeat that isn't there costs a little alignment drift; deleting a line the
    singer sings loses it from the karaoke entirely. So the sheet is the backbone
    and the transcript may only *insert* into it.

    Consecutive transcript segments matching the *same* sheet line collapse into
    one: the transcript often splits a sung line in two, which is an artefact
    rather than a repeat.
    """
    plan = map_performance(sheet, hyp_lines, min_match)

    # Collapse ASR-split duplicates, keeping the matched entries in time order.
    matched: List[dict] = []
    for p in plan:
        j = p["sheetIdx"]
        if j is None:
            continue
        if matched and j == matched[-1]["sheetIdx"] and \
                p["startMs"] - matched[-1]["endMs"] < 1500:
            matched[-1]["endMs"] = p["endMs"]
            continue
        matched.append({"sheetIdx": j, "startMs": p["startMs"],
                        "endMs": p["endMs"], "score": p["score"]})

    lines: List[str] = []
    order: List[dict] = []

    def emit(j: int, repeat: bool, hit: dict = None) -> None:
        lines.append(sheet[j])
        order.append({
            "sheetIdx": j, "repeat": repeat,
            "startMs": (hit or {}).get("startMs"),
            "endMs": (hit or {}).get("endMs"),
            "score": (hit or {}).get("score", 0.0),
        })

    # Walk the matched entries one at a time against a high-water mark. Grouping
    # them into runs first was wrong twice over: a run that began with a repeat
    # but then ran forward got classified as a repeat *whole*, and the high-water
    # mark wasn't advanced on that branch, so the tail re-emitted the entire
    # sheet — 16 lines came out as 28.
    emitted = -1
    for e in matched:
        j = e["sheetIdx"]
        if j > emitted:
            # Forward progress. Emit any sheet lines the transcript skipped over
            # (it has poor recall) so they are never lost, then this one.
            for k in range(emitted + 1, j):
                emit(k, False)
            emit(j, False, e)
            emitted = j
        else:
            # Already past this line, so the recording is singing it again.
            emit(j, True, e)
    for j in range(emitted + 1, len(sheet)):   # tail the transcript never reached
        emit(j, False)
    return lines, order


def coverage(sheet: List[str], order: List[dict]) -> dict:
    """How much of the sheet the performance used, and how much it repeated."""
    return {
        "sheetLines": len(sheet),
        "performanceLines": len(order),
        "repeatsInserted": sum(1 for o in order if o.get("repeat")),
        "linesWithEvidence": sum(1 for o in order if o.get("startMs") is not None),
    }


def resolve_with_llm(sheet: List[str], hyp_lines: List[dict], call) -> List[str]:
    """Optional escape hatch for sheets the matcher can't resolve.

    `call(prompt) -> str` is supplied by the caller so this module stays free of
    any SDK or API key. Only worth reaching for when `coverage()` looks wrong —
    a sheet in the wrong order, interleaved with a translation, or carrying a
    second song. For the ordinary "chorus written once, sung twice" case the
    deterministic path above is cheaper, faster and does not invent lines.

    The model is asked to *reorder and repeat the given lines only*; any line it
    returns that is not in the sheet is dropped, because an LLM inventing lyrics
    is the one failure this whole pipeline exists to avoid.
    """
    numbered = "\n".join(f"{i}: {l}" for i, l in enumerate(sheet))
    heard = "\n".join(f"{h['startMs']/1000:.1f}s: {h['text']}" for h in hyp_lines)
    prompt = (
        "A lyric sheet is written in canonical form (chorus once). A rough "
        "machine transcript shows what the recording actually sings, in order, "
        "with timings. The transcript has many wrong words — trust it only for "
        "ORDER and REPETITION.\n\n"
        f"SHEET (numbered):\n{numbered}\n\nTRANSCRIPT:\n{heard}\n\n"
        "Output the sheet line numbers in the order they are actually sung, one "
        "per line, repeating a number when its line is sung again. Output "
        "nothing but numbers."
    )
    raw = call(prompt)
    out = []
    for tok in re.findall(r"\d+", raw or ""):
        i = int(tok)
        if 0 <= i < len(sheet):          # never accept a line not in the sheet
            out.append(sheet[i])
    return out