"""LLM answering: prompts, generation, and output parsing (scaffolded and lean)."""

from __future__ import annotations

import re
from typing import List, Optional, Sequence, Tuple

from .llm import LLMClient
from .preprocess import Puzzle

SYSTEM = (
    "You solve International Linguistics Olympiad problems. Every problem is "
    "about a language you have never seen; ALL the evidence you need is in the "
    "given data. Derive the rules ONLY from that data — do not assume the "
    "language works like English or any language you know. "
    "A 'Mechanical analysis' section may provide segmentation, alignment, and "
    "candidate answers computed by exact algorithms: use them as hypotheses, "
    "adopt candidates that fit the data, and correct them when the data "
    "disagrees. "
    "First, reason briefly about the linguistic patterns — keep it to a few "
    "lines, not an essay: line up the given examples, segment the words into "
    "morphemes, note what each recurring morpheme means and the order they "
    "combine in, and any sound changes; check that the pattern holds across the "
    "examples, then apply it to each query item. "
    "Answer format by task type -- translation: the translated form only, in "
    "the language asked for; fill_blanks: only the missing form for each "
    "blank; match_letters: only the option letter (for example A, B, C); "
    "text_to_num: the number in digits; num_to_text: the number written out "
    "in the puzzle language; anything else: exactly what the instruction "
    "asks, nothing more. Match the punctuation style of the given examples. "
    "Then write a line that says exactly "
    "FINAL ANSWERS: and, below it, one answer per line in the order the items "
    "are asked -- the bare answer only, no numbering, no quotes, no markdown, "
    "no extra text. If you are unsure, still give your single best guess for "
    "every item -- never leave an item blank. After the answers, write a line "
    "that says exactly EXPLANATION: followed by 2-4 short bullet points "
    "stating the rules you found (word order, affixes and their functions, "
    "sound changes, numeral bases) and the key evidence for them. Do not "
    "repeat your reasoning.\n\n"
    "Output shape (follow it exactly; do not bold or number the markers):\n"
    "<a few lines of pattern reasoning>\n"
    "FINAL ANSWERS:\n"
    "first answer\n"
    "second answer\n"
    "EXPLANATION:\n"
    "- rule and evidence\n"
    "- rule and evidence"
)

# Minimal prompt: no scaffold, no chain-of-thought. Selected by LEAN_MODE.
LEAN_SYSTEM = (
    "You solve International Linguistics Olympiad problems. Answer every "
    "numbered item. Put each answer on its own line, in order, with no "
    "numbering and no extra text. Give your best guess for every item; never "
    "leave one blank."
)

MAX_NEW_TOKENS = 1536

# Markers tolerate leading markdown/quote decoration (#, *, >, -, spaces) and
# trailing decoration/colon, and capture any inline content after the marker
# ("FINAL ANSWERS: foo" -> "foo" is the first answer). Qwen habitually bolds
# these headers; the strict "^marker$" form silently dropped every such output.
# trailing class excludes newline ([^\S\n]) so the marker never swallows the
# line break and misreads the next line as inline content
_FINAL_RX = re.compile(r"(?im)^[ \t#>*_`-]*final[ \t]*answers?\b[^\S\n]*[:*_`]?[^\S\n]*(.*)$")
_EXPL_RX = re.compile(r"(?im)^[ \t#>*_`-]*explanation\b[^\S\n]*[:*_`]?[^\S\n]*(.*)$")
_NUMBERING_RX = re.compile(r"^\s*\(?\d{1,3}[.)]\s*")
# leading markdown decoration on an answer line: bullets, bold, backticks
_ANSWER_DECOR_RX = re.compile(r"^[\s>*_`•·–-]+")


def build_prompt(puzzle: Puzzle, scaffold: str = "") -> str:
    base = f"{puzzle.context.strip()}\n\n{puzzle.query.strip()}"
    if scaffold:
        base += f"\n\n{scaffold}"
    return base


def _clean_answer_line(line: str) -> str:
    line = _NUMBERING_RX.sub("", line.strip())
    line = _ANSWER_DECOR_RX.sub("", line)
    # strip trailing markdown emphasis but keep sentence punctuation (EM-significant)
    line = re.sub(r"[*_`]+$", "", line)
    return line.strip().strip("'\"“”").strip()


def parse_output(text: str, n_items: Optional[int] = None
                 ) -> Tuple[List[str], str, bool]:
    """Returns (answers, explanation, found_marker).

    `found_marker` is True only when an explicit FINAL ANSWERS marker was
    present. On a parse failure (no marker) we return ([], expl_if_any, False)
    rather than treating the reasoning prose as answers -- the caller then
    keeps its symbolic answers instead of overwriting them with garbage.
    When `n_items` is given it is used only as a sanity cap on how many
    answer lines to accept (guards against a runaway list)."""
    markers = list(_FINAL_RX.finditer(text))
    found = bool(markers)
    if markers:
        m = markers[-1]
        # a marker may carry its first answer inline: "FINAL ANSWERS: foo"
        inline = m.group(1).strip() if m.lastindex else ""
        tail = (inline + "\n" if inline else "") + text[m.end():]
    else:
        tail = ""

    expl = ""
    em = _EXPL_RX.search(tail) if found else None
    if em:
        inline_e = em.group(1).strip() if em.lastindex else ""
        rest_e = tail[em.end():].strip()
        expl = (inline_e + ("\n" + rest_e if rest_e else "")) if inline_e else rest_e
        tail = tail[: em.start()]

    answers: List[str] = []
    cap = (2 * n_items + 4) if n_items else None
    for line in tail.splitlines():
        cleaned = _clean_answer_line(line)
        if cleaned:
            answers.append(cleaned)
            if cap and len(answers) >= cap:
                break

    # models sometimes put EXPLANATION before FINAL ANSWERS despite the prompt
    if not expl:
        em2 = _EXPL_RX.search(text)
        if em2 and (not markers or em2.start() < markers[-1].start()):
            seg = text[em2.end():]
            stop = _FINAL_RX.search(seg)
            expl = (seg[: stop.start()] if stop else seg).strip()
    return answers, expl, found


# lines the model may prepend/append around a bare answer list
_CHATTY_RX = re.compile(
    r"^(?:here (?:are|is)\b|answers?\s*:?\s*$|the answers?\b|translations?\s*:?\s*$|"
    r"note\b|okay\b|sure\b|solution\b|let me\b)", re.IGNORECASE)


def parse_output_lean(text: str, n_items: Optional[int] = None
                      ) -> Tuple[List[str], str, bool]:
    """Lean parse: the minimal prompt asks for bare answers, one per line, with
    no FINAL ANSWERS marker — so every non-empty line IS an answer. We only
    strip numbering/markdown decoration and drop obvious chatty preamble lines.
    When more than n_items lines survive, keep the LAST n_items (any stray
    preamble sits at the top). No marker requirement, no salvage."""
    if not text:
        return [], "", False
    answers: List[str] = []
    for line in text.splitlines():
        cleaned = _clean_answer_line(line)
        if cleaned and not _CHATTY_RX.match(cleaned):
            answers.append(cleaned)
    if n_items and len(answers) > n_items:
        answers = answers[-n_items:]
    return answers, "", bool(answers)


_LEAN_LABEL_RX = re.compile(r"^\s*\(?(\d{1,3})\)?[.):\]]\s")


def align_lean(text: str, puzzle: Puzzle) -> List[Optional[str]]:
    """Turn a lean (bare-lines) model output into exactly len(puzzle.items)
    answers, blending two placement methods that stack:

    1. LABEL-AWARE placement — if the model numbered its answer lines and those
       numbers cover most of the item labels, place each answer under its own
       label. This is robust to the model reordering items or skipping one (a
       single skipped line would otherwise shift every later answer and zero the
       whole block on both metrics).
    2. POSITIONAL last-N fallback — when the output isn't reliably numbered, take
       the cleaned, non-chatty lines in order (dropping any preamble at the top).

    Both share our line hygiene (numbering/markdown stripping via
    `_clean_answer_line`, preamble drop via `_CHATTY_RX`). Returns a length-N
    list; None marks items the model did not answer, which the caller leaves to
    the fallback."""
    items = puzzle.items
    n = len(items)
    if n == 0:
        return []
    labeled: dict = {}
    for line in (text or "").splitlines():
        s = line.strip()
        if not s:
            continue
        m = _LEAN_LABEL_RX.match(s)
        cleaned = _clean_answer_line(s)
        if not cleaned or _CHATTY_RX.match(cleaned):
            continue
        if m:
            labeled[m.group(1)] = cleaned  # last write wins (models restate)
    item_labels = [it.number for it in items]
    if labeled and all(lbl for lbl in item_labels):
        covered = sum(1 for lbl in item_labels if lbl in labeled)
        if covered >= max(1, (2 * n + 2) // 3):  # ~2/3 labelled -> trust labels
            return [labeled.get(lbl) for lbl in item_labels]
    # positional fallback: our bare-line parse, aligned by position
    bare, _e, _f = parse_output_lean(text, n)
    out: List[Optional[str]] = [None] * n
    for i in range(min(n, len(bare))):
        out[i] = bare[i]
    return out


def build_salvage_prompt(base_prompt: str, raw_reasoning: str,
                         max_chars: int = 2400) -> str:
    """A short follow-up prompt that reuses reasoning the model already
    produced (but never terminated with a FINAL ANSWERS block, e.g. it hit the
    token cap). We feed the reasoning back as context and ask ONLY for the
    answer block -- cheaper and more reliable than a bare retry."""
    reasoning = raw_reasoning.strip()
    if len(reasoning) > max_chars:
        reasoning = reasoning[-max_chars:]  # keep the most recent reasoning
    return (
        f"{base_prompt}\n\n"
        "You already worked through this problem:\n"
        "-----\n"
        f"{reasoning}\n"
        "-----\n"
        "Now output ONLY the answer block, nothing else. Write the line "
        "FINAL ANSWERS: then one bare answer per line in item order (your best "
        "guess for every item, never blank), then a line EXPLANATION: with 2-4 "
        "short bullets."
    )


def solve_direct(puzzles: Sequence[Puzzle], client: LLMClient,
                 scaffolds: Optional[Sequence[str]] = None,
                 max_new_tokens: int = MAX_NEW_TOKENS,
                 salvage: bool = True,
                 system: str = SYSTEM,
                 lean: bool = False,
                 sample: bool = False,
                 temperature: float = 0.5
                 ) -> List[Tuple[List[str], str, str, bool]]:
    """Batched single-shot answering. Returns per puzzle
    (answers, explanation, raw_text, found_marker). `sample`/`temperature` drive
    the sampled passes used by self-consistency voting; salvage is skipped when
    sampling. In `lean` mode every non-empty line is an answer (no marker)."""
    if not client.available or not puzzles:
        return [([], "", "", False) for _ in puzzles]
    prompts = [
        build_prompt(p, scaffolds[i] if scaffolds else "")
        for i, p in enumerate(puzzles)
    ]
    raws = client.generate(prompts, max_new_tokens=max_new_tokens, system=system,
                           sample=sample, temperature=temperature)
    out: List[Tuple[List[str], str, str, bool]] = []
    salvage_idx: List[int] = []
    for i, r in enumerate(raws):
        if lean:
            aligned = align_lean(r, puzzles[i])
            answers, expl, found = aligned, "", any(a for a in aligned)
        else:
            answers, expl, found = parse_output(r, len(puzzles[i].items))
        out.append((answers, expl, r, found))
        if salvage and not lean and not sample and not found and r.strip():
            salvage_idx.append(i)

    if salvage_idx:
        sp = [build_salvage_prompt(prompts[i], raws[i]) for i in salvage_idx]
        # answers only — a small cap keeps the salvage pass cheap
        cap = min(max_new_tokens, 384)
        sraws = client.generate(sp, max_new_tokens=cap, system=system)
        for i, sr in zip(salvage_idx, sraws):
            a2, e2, f2 = parse_output(sr, len(puzzles[i].items))
            if f2 and a2:
                prev = out[i]
                out[i] = (a2, e2 or prev[1], prev[2], True)
    return out


def align_answers(direct: List[str], n_items: int) -> List[Optional[str]]:
    """Position-align a FINAL ANSWERS block to the expected item count."""
    out: List[Optional[str]] = [None] * n_items
    for i in range(min(n_items, len(direct))):
        out[i] = direct[i]
    return out