"""Query ranking and budgeted context assembly."""

from __future__ import annotations

import math
import re
from collections import Counter
from pathlib import Path

from . import graph as graph_module
from . import tokens, weights

K1 = 1.2
B = 0.75
PATH_WEIGHT = 3
NAME_WEIGHT = 4
PRESELECT_FILES = 40
FILE_WINDOW_LINES = 60
SEED_FILES = 6
ONE_HOP_BONUS = 0.45
TWO_HOP_BONUS = 0.18
CALL_HOP_BONUS = 0.30
EXACT_NAME_BONUS = 1.6
PARTIAL_NAME_BONUS = 0.5
PATH_STEM_BONUS = 0.8
BODY_HIT_BONUS = 0.2
BODY_HIT_CAP = 0.6
SESSION_SCALE = 0.5
SESSION_CAP = 0.9
LOW_COVERAGE = 0.5
HIGH_COVERAGE = 0.6
FLAT_SEPARATION = 0.12
HIGH_SEPARATION = 0.25
# Adaptive packing: the budget is a ceiling, spend scales with UNCERTAINTY, not
# with rank strength. High confidence means the ranking already nailed it - the
# winner plus a couple of neighbours answers the question and everything past
# them is padding the caller pays for. Medium keeps the most room because that
# is where extra context genuinely helps the model resolve ambiguity; low gets
# the least because the ranking itself is suspect. An exact-name hit on top of
# high confidence shrinks the ceiling further still. Measured on this repo's
# BM25 distributions: tails are flat (rarely under 40% of the top score), so
# the mass cut carries the early-stop weight and the floor only trims outliers.
CONFIDENCE_BUDGET_SCALE = {"high": 0.55, "medium": 0.65, "low": 0.4}
EXACT_HIGH_BUDGET_SCALE = 0.35
ADAPTIVE_FLOOR_RATIO = 0.35
ADAPTIVE_MASS = 0.70
ADAPTIVE_MIN_ITEMS = 2
# Skeleton packing: supporting items keep only the lines that earn their
# place - query hits, signatures, small constants, one line of context - and
# blank the rest, one marker line per gap. Line count is preserved so the
# reported ranges stay true; the chars are what the caller pays for. The
# first-packed item and exact hits always ship whole.
SKELETON_MIN_LINES = 16
SKELETON_CONTEXT = 1
SKELETON_MARKER = "⋯"
SKELETON_CONST_SPAN_CAP = 12
ENUMERATION_CUES = re.compile(r"\b(all|every|each|list|enumerate|complete|entire)\b")


class Query:
    def __init__(self, text: str) -> None:
        self.text = text if isinstance(text, str) else ""
        self.terms = Counter(tokens.tokenize(self.text))
        self.identifiers = [
            identifier
            for identifier in tokens.identifiers(self.text)
            if len(identifier) >= 3 and identifier.lower() not in tokens.STOPWORDS
        ]
        self.lowered = {identifier.lower() for identifier in self.identifiers}

    def __bool__(self) -> bool:
        return bool(self.terms)


def _bm25(counts: Counter, length: int, average: float, idf: dict[str, float], query: Counter) -> float:
    total = 0.0
    for term, frequency in query.items():
        occurrences = counts.get(term, 0)
        if not occurrences:
            continue
        denominator = occurrences + K1 * (1 - B + B * (length / average if average else 1.0))
        total += idf.get(term, 0.0) * frequency * (occurrences * (K1 + 1)) / denominator
    return total


def _documents(graph: graph_module.Graph) -> tuple[dict[str, Counter], float]:
    cached = getattr(graph, "_rank_documents", None)
    if cached is not None:
        return cached
    documents: dict[str, Counter] = {}
    total_length = 0
    for rel, record in graph.files.items():
        counts: Counter = Counter()
        for token in tokens.path_tokens(rel):
            counts[token] += PATH_WEIGHT
        for symbol in record["symbols"]:
            name = symbol.get("name")
            if isinstance(name, str):
                for token in tokens.tokenize(name):
                    counts[token] += NAME_WEIGHT
        for token in record["terms"]:
            if isinstance(token, str):
                counts[token] += 1
        documents[rel] = counts
        total_length += sum(counts.values())
    average = total_length / len(documents) if documents else 1.0
    result = (documents, average or 1.0)
    graph._rank_documents = result  # noqa: SLF001 - cache lives with the cached graph
    return result


def _idf(documents: dict[str, Counter], terms: Counter) -> dict[str, float]:
    total = max(1, len(documents))
    table: dict[str, float] = {}
    for term in terms:
        holders = sum(1 for counts in documents.values() if term in counts)
        table[term] = math.log(1 + (total - holders + 0.5) / (holders + 0.5))
    return table


def _session_bonus(session: dict[tuple[str, str | None], float], rel: str, symbol: str | None) -> float:
    value = session.get((rel, symbol), 0.0) if symbol else 0.0
    value += session.get((rel, None), 0.0)
    if symbol:
        value += 0.5 * sum(
            weight for (path, name), weight in session.items() if path == rel and name and name != symbol
        )
    return min(SESSION_CAP, SESSION_SCALE * value)


def _proximity(graph: graph_module.Graph, seeds: set[str]) -> dict[str, float]:
    if not seeds:
        return {}
    first: set[str] = set()
    for seed in seeds:
        first |= graph.imports.get(seed, set())
        first |= graph.imported_by.get(seed, set())
    first -= seeds
    second: set[str] = set()
    for rel in first:
        second |= graph.imports.get(rel, set())
        second |= graph.imported_by.get(rel, set())
    second -= seeds | first
    # Call edges are one hop only and weaker than an import: a resolved call is
    # a strong hint but its name matching is looser than import resolution.
    call_hop: set[str] = set()
    for seed in seeds:
        call_hop |= graph.calls.get(seed, set())
        call_hop |= graph.called_by.get(seed, set())
    call_hop -= seeds | first
    bonus = {rel: ONE_HOP_BONUS for rel in first}
    for rel in call_hop:
        bonus.setdefault(rel, CALL_HOP_BONUS)
    for rel in second:
        bonus.setdefault(rel, TWO_HOP_BONUS)
    return bonus


def _exact_bonus(query: Query, rel: str, symbol: str | None, body: str) -> tuple[float, bool]:
    """Return the identifier bonus and whether a query identifier matched a name outright."""
    bonus = 0.0
    exact = False
    if symbol:
        leaf = symbol.split(".")[-1].lower()
        if leaf in query.lowered:
            bonus += EXACT_NAME_BONUS
            exact = True
        elif any(len(identifier) >= 4 and identifier in leaf for identifier in query.lowered):
            bonus += PARTIAL_NAME_BONUS
    stem = rel.rsplit("/", 1)[-1].split(".")[0].lower()
    if stem in query.lowered:
        bonus += PATH_STEM_BONUS
        exact = True
    if body:
        hits = sum(1 for identifier in query.identifiers if identifier in body)
        bonus += min(BODY_HIT_CAP, BODY_HIT_BONUS * hits)
    return bonus, exact


def _best_window(lines: list[str], query: Query, span: int) -> tuple[int, int]:
    if not lines:
        return 1, 1
    if not query.terms:
        return 1, min(len(lines), span)
    hits = []
    for line in lines:
        lowered = line.lower()
        hits.append(sum(1 for term in query.terms if term in lowered))
    span = min(span, len(lines))
    running = sum(hits[:span])
    best_score, best_start = running, 0
    for index in range(span, len(lines)):
        running += hits[index] - hits[index - span]
        if running > best_score:
            best_score, best_start = running, index - span + 1
    if best_score == 0:
        return 1, span
    return best_start + 1, best_start + span


def _candidates(
    graph: graph_module.Graph,
    query: Query,
    idf: dict[str, float],
    average: float,
    preselected: list[str],
    proximity: dict[str, float],
    session: dict[tuple[str, str | None], float],
) -> list[dict]:
    out: list[dict] = []
    normalizer = max(1, len(query.terms))
    for rel in preselected:
        lines = graph.source(rel)
        if lines is None:
            continue
        record = graph.files[rel]
        path_counts: Counter = Counter()
        for token in tokens.path_tokens(rel):
            path_counts[token] += PATH_WEIGHT

        for symbol in record["symbols"]:
            name = symbol.get("name")
            try:
                start = int(symbol["line"])
                end = int(symbol["end_line"])
            except (KeyError, TypeError, ValueError):
                continue
            if not isinstance(name, str) or start < 1 or end < start or start > len(lines):
                continue
            body = "\n".join(lines[start - 1 : min(end, len(lines))])
            counts = Counter(path_counts)
            for token in tokens.tokenize(name):
                counts[token] += NAME_WEIGHT
            for token in tokens.tokenize(body):
                counts[token] += 1
            bonus, exact = _exact_bonus(query, rel, name, body)
            score = _bm25(counts, sum(counts.values()), average, idf, query.terms) / normalizer
            score += bonus + proximity.get(rel, 0.0) + _session_bonus(session, rel, name)
            out.append(
                {
                    "path": rel,
                    "symbol": name,
                    "score": score,
                    "start": start,
                    "end": min(end, len(lines)),
                    "lines": lines,
                    "matched": sum(1 for term in query.terms if counts.get(term)),
                    "exact": exact,
                }
            )

        start, end = _best_window(lines, query, FILE_WINDOW_LINES)
        body = "\n".join(lines[start - 1 : end])
        counts = Counter(path_counts)
        for token in tokens.tokenize(body):
            counts[token] += 1
        bonus, exact = _exact_bonus(query, rel, None, body)
        score = _bm25(counts, sum(counts.values()), average, idf, query.terms) / normalizer
        score += bonus + proximity.get(rel, 0.0) + _session_bonus(session, rel, None)
        out.append(
            {
                "path": rel,
                "symbol": None,
                "score": score,
                "start": start,
                "end": end,
                "lines": lines,
                "matched": sum(1 for term in query.terms if counts.get(term)),
                "exact": exact,
            }
        )
    out.sort(key=lambda item: item["score"], reverse=True)
    return out


def _const_keep_rows(graph: graph_module.Graph, rel: str, start: int, end: int) -> set[int]:
    """Window-relative rows of small const symbols fully inside [start, end].

    Multi-line constant values carry their content on continuation lines that
    match no query term; blanking those would delete the very answer the item
    was packed for.
    """
    record = graph.files.get(rel)
    if record is None:
        return set()
    rows: set[int] = set()
    for symbol in record["symbols"]:
        if symbol.get("kind") != "const":
            continue
        try:
            first, last = int(symbol["line"]), int(symbol["end_line"])
        except (KeyError, TypeError, ValueError):
            continue
        if first < start or last > end or last - first + 1 > SKELETON_CONST_SPAN_CAP:
            continue
        rows.update(range(first - start, last - start + 1))
    return rows


def _skeleton_window(window: list[str], query: Query, const_rows: set[int]) -> list[str]:
    """Blank the rows that earn no place, preserving the row count exactly."""
    total = len(window)
    keep = [False] * total
    for row, line in enumerate(window):
        if row <= 1 or row == total - 1 or row in const_rows:
            keep[row] = True
            continue
        lowered = line.lower()
        if any(term in lowered for term in query.terms):
            keep[row] = True
            continue
        if any(identifier in line for identifier in query.identifiers):
            keep[row] = True
    padded = list(keep)
    for row, kept in enumerate(keep):
        if kept:
            if row > 0:
                padded[row - 1] = True
            if row + 1 < total:
                padded[row + 1] = True
    out: list[str] = []
    gap = 0
    for row, line in enumerate(window):
        if padded[row]:
            out.append(line)
            gap = 0
        else:
            gap += 1
            out.append(SKELETON_MARKER if gap == 1 and _gap_length(padded, row) >= 2 else "")
    return out


def _gap_length(padded: list[bool], row: int) -> int:
    length = 0
    index = row
    while index < len(padded) and not padded[index]:
        length += 1
        index += 1
    return length


def _confidence(selected: list[dict], query: Query) -> str:
    """Judge the answer by how much of the query the winner covers and how far it leads.

    Absolute BM25 values are not comparable between queries, so confidence rests
    on scale-free signals: term coverage of the top hit, whether a query
    identifier named it outright, and its lead over the best hit from a
    *different* file. Several hits from one file is agreement, not ambiguity.

    Coverage is measured over the whole pack, not over selected[0]. BM25 favours
    short symbols, so a bare `Session` can outrank `validateSessionToken` on a
    question about validating session tokens; judging the pack by its first row
    then throws away an answer it is already carrying.
    """
    if not selected or not query.terms:
        return "low"
    top = selected[0]
    coverage = max(item["matched"] for item in selected) / len(query.terms)
    rivals = [item["score"] for item in selected[1:] if item["path"] != top["path"]]
    best = top["score"]
    separation = 1.0 if not rivals or best <= 0 else (best - max(rivals)) / best
    if coverage < LOW_COVERAGE:
        return "low"
    if coverage < HIGH_COVERAGE and not top["exact"] and separation < FLAT_SEPARATION:
        return "low"
    if coverage >= HIGH_COVERAGE and (top["exact"] or separation >= HIGH_SEPARATION):
        return "high"
    return "medium"


def retrieve(
    root: Path,
    query_text: str,
    budget_chars: int = 18000,
    max_read_chars: int = 4000,
    limit: int = 12,
    adaptive: bool = True,
    skeleton: bool = True,
) -> dict:
    graph = graph_module.ensure(root)
    if graph is None or not graph.files:
        return {"items": [], "confidence": "low", "used_chars": 0, "stale": False, "budget_chars": 0}

    stale = graph.stale()
    query = Query(query_text)
    documents, average = _documents(graph)
    idf = _idf(documents, query.terms)
    try:
        session = weights.scores(root)
    except Exception:
        session = {}

    normalizer = max(1, len(query.terms))
    base: dict[str, float] = {}
    for rel, counts in documents.items():
        score = _bm25(counts, sum(counts.values()), average, idf, query.terms) / normalizer
        score += _session_bonus(session, rel, None)
        if score > 0:
            base[rel] = score

    seeds = {rel for rel, _ in sorted(base.items(), key=lambda item: item[1], reverse=True)[:SEED_FILES]}
    proximity = _proximity(graph, seeds)
    preliminary = dict(base)
    for rel, bonus in proximity.items():
        if rel in graph.files:
            preliminary[rel] = preliminary.get(rel, 0.0) + bonus

    preselected = [
        rel for rel, _ in sorted(preliminary.items(), key=lambda item: item[1], reverse=True)[:PRESELECT_FILES]
    ]
    if not preselected:
        return {"items": [], "confidence": "low", "used_chars": 0, "stale": stale, "budget_chars": 0}

    ranked = _candidates(graph, query, idf, average, preselected, proximity, session)

    effective_budget = budget_chars
    top_score = ranked[0]["score"] if ranked else 0.0
    total_mass = sum(max(0.0, candidate["score"]) for candidate in ranked[: max(0, int(limit))])
    if adaptive:
        provisional = _confidence(ranked[: max(0, int(limit))], query)
        scale = CONFIDENCE_BUDGET_SCALE.get(provisional, 1.0)
        if provisional == "high" and ranked and ranked[0]["exact"]:
            scale = EXACT_HIGH_BUDGET_SCALE
        effective_budget = int(budget_chars * scale)

    enumerating = bool(ENUMERATION_CUES.search(query.text.lower()))
    items: list[dict] = []
    chosen: list[dict] = []
    used = 0
    packed_mass = 0.0
    per_file: dict[str, int] = {}
    covered: dict[str, list[tuple[int, int]]] = {}
    for candidate in ranked:
        if len(items) >= max(0, int(limit)) or used >= effective_budget:
            break
        if adaptive and len(items) >= ADAPTIVE_MIN_ITEMS:
            if top_score > 0 and candidate["score"] < ADAPTIVE_FLOOR_RATIO * top_score:
                break
            if total_mass > 0 and packed_mass >= ADAPTIVE_MASS * total_mass:
                break
        rel = candidate["path"]
        if any(start <= candidate["start"] and candidate["end"] <= end for start, end in covered.get(rel, [])):
            continue
        allowance = min(max_read_chars - per_file.get(rel, 0), effective_budget - used)
        if allowance <= 0:
            continue
        window = candidate["lines"][candidate["start"] - 1 : candidate["end"]]
        compressed = False
        if (
            skeleton
            and not enumerating
            and items
            and not candidate["exact"]
            and (candidate["symbol"] is None or len(window) > SKELETON_MIN_LINES)
        ):
            slim = _skeleton_window(
                window, query, _const_keep_rows(graph, rel, candidate["start"], candidate["end"])
            )
            compressed = slim != window
            window = slim
        text = "\n".join(window)
        if not text.strip():
            continue
        text = text[:allowance]
        end_line = candidate["start"] + text.count("\n")
        item = {
            "path": rel,
            "symbol": candidate["symbol"],
            "score": round(candidate["score"], 4),
            "lines": [candidate["start"], end_line],
            "text": text,
        }
        if compressed:
            item["skeleton"] = True
        items.append(item)
        chosen.append(candidate)
        used += len(text)
        packed_mass += max(0.0, candidate["score"])
        per_file[rel] = per_file.get(rel, 0) + len(text)
        covered.setdefault(rel, []).append((candidate["start"], candidate["end"]))

    return {
        "items": items,
        "confidence": _confidence(chosen, query),
        "used_chars": used,
        "stale": stale,
        "budget_chars": effective_budget,
    }