#!/usr/bin/env python3
"""Total-count gate for the TQ* rules in scripts/check_test_quality.py.

Sibling of scripts/type_discipline_gate.py, pointed at the test tree instead of
the package. Each rule listed in test-quality-budget.json has a hard ``limit``.
The gate counts each rule across the whole `tests` tree and fails when a rule is
both over its limit and higher than the base it merges into, so a change is
blamed for the violations it adds, never for drift that already exists in the
base.

Every rule is seeded at exactly its count on the day the gate landed, so the
suite's existing debt is grandfathered and any net-new violation trips the gate
immediately. ``--update`` ratchets a limit down by the violations this branch
fixed relative to its branch point (the merge-base), so the ceilings only ever
fall. A rule absent from the budget at the merge-base was seeded on this branch;
``--update`` leaves its limit untouched, because the base tree predates the rule
and its whole grandfathered count would otherwise be misread as "fixed".

The deliberate difference from its sibling: this gate has no headroom anywhere.
Type discipline seeded LIT010/LIT011 at 1.5x to leave room for an in-flight
sweep; a test-quality violation has no such transition to absorb, so the line is
today's count and the only legal direction is down.
"""

from __future__ import annotations

import argparse
import json
import re
import shutil
import subprocess
import sys
import tempfile
from collections import Counter
from collections.abc import Mapping, Sequence
from pathlib import Path
from types import MappingProxyType
from typing import Final, NamedTuple

REPO_ROOT: Final = Path(__file__).resolve().parent.parent
CHECKER: Final = REPO_ROOT / "scripts" / "check_test_quality.py"
BUDGET_PATH: Final = REPO_ROOT / "test-quality-budget.json"
TARGET: Final = "tests"
DEFAULT_BASE: Final = "origin/litellm_internal_staging"

_HUNK: Final = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@", re.MULTILINE)
_FILE_HEADER: Final = re.compile(r"^\+\+\+ b/(.+)$", re.MULTILINE)
_LINE: Final = re.compile(r"^(?P<file>.+?):(?P<line>\d+): (?P<code>TQ\d+) ")


class Violation(NamedTuple):
    file: str
    line: int
    code: str


class Breach(NamedTuple):
    rule: str
    total: int
    cap: int
    added: int


def _run(cmd: Sequence[str], cwd: Path = REPO_ROOT) -> str:
    proc: Final = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True)
    if proc.returncode not in (0, 1):
        sys.stderr.write(proc.stderr)
        raise SystemExit(f"{cmd[0]} exited {proc.returncode}")
    return proc.stdout


def resolve_base_point(base_ref: str, cwd: Path = REPO_ROOT) -> str:
    """The snapshot commit base counts are measured at: merge-base(base_ref, HEAD),
    made aware of an in-progress merge. Mid-merge, HEAD is still the pre-merge tip,
    so its merge-base is the old branch point and every violation the base gained
    since then would be blamed on this change."""
    head_point: Final = _run(["git", "merge-base", base_ref, "HEAD"], cwd=cwd).strip()
    if not head_point:
        return base_ref
    merge_head: Final = _run(["git", "rev-parse", "--verify", "--quiet", "MERGE_HEAD"], cwd=cwd).strip()
    if not merge_head:
        return head_point
    merge_point: Final = _run(["git", "merge-base", base_ref, merge_head], cwd=cwd).strip()
    if not merge_point:
        return head_point
    older: Final = _run(["git", "merge-base", head_point, merge_point], cwd=cwd).strip()
    return merge_point if older == head_point else head_point


def _check(root: Path, checker: Path) -> tuple[Violation, ...]:
    # macOS tempfile dirs (/var/...) resolve to /private/var/..., so relative_to needs both sides resolved.
    resolved: Final = root.resolve()
    out: Final = _run([sys.executable, str(checker), str(resolved / TARGET)], cwd=resolved)
    return tuple(
        Violation(
            (resolved / match.group("file")).resolve().relative_to(resolved).as_posix(),
            int(match.group("line")),
            match.group("code"),
        )
        for line in out.splitlines()
        if (match := _LINE.match(line)) is not None
    )


def head_violations() -> tuple[Violation, ...]:
    return _check(REPO_ROOT, CHECKER)


def count_by_rule(violations: Sequence[Violation]) -> Mapping[str, int]:
    return MappingProxyType(dict(Counter(v.code for v in violations)))


def base_counts(ref: str) -> Mapping[str, int]:
    """Rule counts at `ref`, measured with the *current* rule logic rather than
    whatever the checker looked like at that commit."""
    parent: Final = Path(tempfile.mkdtemp(prefix="tq_base_"))
    worktree: Final = parent / "wt"
    try:
        _run(["git", "worktree", "add", "--detach", str(worktree), ref])
        (worktree / "scripts").mkdir(parents=True, exist_ok=True)
        checker: Final = worktree / "scripts" / "check_test_quality.py"
        shutil.copy(CHECKER, checker)
        return count_by_rule(_check(worktree, checker))
    finally:
        # Teardown must never raise, or it masks the real error when the body failed.
        subprocess.run(
            ["git", "worktree", "remove", "--force", str(worktree)],
            cwd=REPO_ROOT, capture_output=True, text=True,
        )
        shutil.rmtree(parent, ignore_errors=True)


def over_ceiling(head: Mapping[str, int], budget: Mapping[str, Mapping[str, int]]) -> frozenset[str]:
    """Rules whose head count already exceeds their limit. When none are, the base
    comparison cannot change the verdict and the base worktree scan is skipped."""
    return frozenset(
        rule for rule, spec in budget.items() if head.get(rule, 0) > spec["limit"]
    )


def evaluate(
    head: Mapping[str, int],
    base: Mapping[str, int],
    budget: Mapping[str, Mapping[str, int]],
) -> tuple[Breach, ...]:
    return tuple(sorted(
        Breach(rule, head.get(rule, 0), spec["limit"], head.get(rule, 0) - base.get(rule, 0))
        for rule, spec in budget.items()
        if head.get(rule, 0) > spec["limit"] and head.get(rule, 0) > base.get(rule, 0)
    ))


def _hunk_lines(body: str) -> frozenset[int]:
    return frozenset(
        line
        for match in _HUNK.finditer(body)
        for start in (int(match.group(1)),)
        for line in range(start, start + (int(match.group(2)) if match.group(2) is not None else 1))
    )


def parse_changed_lines(diff_text: str) -> Mapping[str, frozenset[int]]:
    """Each file in the diff mapped to the line numbers it adds. Splitting on the
    `+++ b/` headers keeps this a pure expression: `split` hands back
    [preamble, path, body, path, body, ...], so each file's hunks are already
    grouped with it."""
    parts: Final = _FILE_HEADER.split(diff_text)
    return MappingProxyType({
        path: _hunk_lines(body)
        for path, body in zip(parts[1::2], parts[2::2])
    })


def introduced(
    violations: Sequence[Violation], changed: Mapping[str, frozenset[int]]
) -> tuple[Violation, ...]:
    return tuple(v for v in violations if v.line in changed.get(v.file, frozenset()))


def cmd_check(base: str) -> None:
    budget: Final = json.loads(BUDGET_PATH.read_text())
    head: Final = head_violations()
    head_counts: Final = count_by_rule(head)
    if not over_ceiling(head_counts, budget):
        print(f"OK: every TQ rule is within its test-suite ceiling (base {base})")
        return
    base_point: Final = resolve_base_point(base)
    breaches: Final = evaluate(head_counts, base_counts(base_point), budget)
    if not breaches:
        print(f"OK: every TQ rule is within its test-suite ceiling (base {base})")
        return
    new: Final = introduced(
        head,
        parse_changed_lines(
            _run(["git", "diff", base_point, "--unified=0", "--no-color", "--", TARGET])
        ),
    )
    print(f"FAIL: TQ-rule totals exceed their limit (base {base}):")
    for breach in breaches:
        print(
            f"  {breach.rule}: total {breach.total} over limit {breach.cap} "
            f"(this change added {breach.added})"
        )
        for violation in sorted(v for v in new if v.code == breach.rule):
            print(f"    {violation.file}:{violation.line}")
    print(
        "Fix the new violations, or give each one a reason "
        "(`# test-quality-ok: <reason>`), or remove an equal number elsewhere; "
        "the ceiling is the limit in test-quality-budget.json. "
        "Run `python scripts/check_test_quality.py tests/` to see every finding."
    )
    raise SystemExit(1)


def ratcheted_budget(
    budget: Mapping[str, Mapping[str, int]],
    current: Mapping[str, int],
    base: Mapping[str, int],
    seeded: frozenset[str] = frozenset(),
) -> Mapping[str, Mapping[str, int]]:
    """Each rule's limit lowered by the violations `current` fixed vs `base`. The drop
    is clamped to what was actually cleared, so a limit only ever falls. Rules in
    `seeded` were introduced on this branch and pass through untouched."""
    return MappingProxyType({
        rule: {
            "limit": spec["limit"] if rule in seeded
            else max(0, spec["limit"] - max(0, base.get(rule, 0) - current.get(rule, 0)))
        }
        for rule, spec in sorted(budget.items())
    })


def _base_budget_rules(base_point: str) -> frozenset[str]:
    proc: Final = subprocess.run(
        ["git", "show", f"{base_point}:{BUDGET_PATH.name}"],
        cwd=REPO_ROOT, capture_output=True, text=True,
    )
    if proc.returncode != 0:
        return frozenset()
    return frozenset(json.loads(proc.stdout))


def cmd_update(base_ref: str = DEFAULT_BASE) -> None:
    """Ratchet each rule's limit down by the violations this branch fixed."""
    budget: Final = json.loads(BUDGET_PATH.read_text())
    base_point: Final = resolve_base_point(base_ref)
    seeded: Final = frozenset(budget) - _base_budget_rules(base_point)
    updated: Final = ratcheted_budget(
        budget, count_by_rule(head_violations()), base_counts(base_point), seeded
    )
    BUDGET_PATH.write_text(json.dumps(dict(updated), indent=2, sort_keys=True) + "\n")
    cleared: Final = sum(budget[rule]["limit"] - updated[rule]["limit"] for rule in updated)
    print(f"Ratcheted TQ-rule limits down by {cleared} violations this branch fixed")
    if seeded:
        print(
            "Left untouched (seeded on this branch, absent from the base budget): "
            + ", ".join(sorted(seeded))
        )


def cmd_seed() -> None:
    """Write the budget from the working tree's current counts. Used once, to land
    the gate; afterwards `--update` is the only thing that may move a limit."""
    counts: Final = count_by_rule(head_violations())
    BUDGET_PATH.write_text(
        json.dumps({rule: {"limit": counts[rule]} for rule in sorted(counts)}, indent=2) + "\n"
    )
    print(f"Seeded {BUDGET_PATH.name} at " + ", ".join(f"{r}={counts[r]}" for r in sorted(counts)))


def main() -> None:
    parser: Final = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base", default=DEFAULT_BASE)
    parser.add_argument("--update", action="store_true")
    parser.add_argument("--seed", action="store_true")
    args: Final = parser.parse_args()
    from gate_slot_lock import held_slot

    with held_slot():
        if args.seed:
            cmd_seed()
        elif args.update:
            cmd_update(args.base)
        else:
            cmd_check(args.base)


if __name__ == "__main__":
    main()