#!/usr/bin/env python3
"""
gate.py - factual invariance gate for tailored incident briefs.

Given a claim ledger (factbase.json) and a draft brief, this checks that the
brief's facts are a subset of the ledger's facts. It does not judge style,
tone, length or framing: those are exactly the things a tailored brief is
supposed to change.

Six checks:
  1. UNSUPPORTED_NUMBER   a numeral appears that is not in the ledger
  2. UNSUPPORTED_DATE     a date appears that is not in the ledger
  3. UNSUPPORTED_ENTITY   a named entity appears that is not in the ledger
  4. BANNED_CLAIM         a known distortion pattern appears
  5. UNATTRIBUTED         a self-serving or contested claim appears with no attribution cue
  6. OVERSTATED_CERTAINTY a certainty marker is applied to a non-established claim

Exit code is 0 if the brief passes, 1 if any finding is raised.

Usage:
  python3 gate.py --factbase factbase.json brief.md
  python3 gate.py --factbase factbase.json --csv results.csv briefs/*.md
"""

import argparse
import csv
import glob
import json
import os
import re
import sys

# ---------------------------------------------------------------- normalising

MONTHS = {
    "january": 1, "february": 2, "march": 3, "april": 4, "may": 5, "june": 6,
    "july": 7, "august": 8, "september": 9, "october": 10, "november": 11,
    "december": 12,
    "jan": 1, "feb": 2, "mar": 3, "apr": 4, "jun": 6, "jul": 7, "aug": 8,
    "sep": 9, "sept": 9, "oct": 10, "nov": 11, "dec": 12,
}


def norm_number(tok):
    """'17,600' -> '17600'; '14.1%' -> '14.1%'; '100x' -> '100x'."""
    t = tok.strip().lower().replace(",", "").replace("~", "")
    t = t.replace("approximately", "").replace(" ", "")
    if t.endswith(".") :
        t = t[:-1]
    # drop a trailing decimal zero so 5.0 == 5
    m = re.fullmatch(r"(\d+)\.0+", t)
    if m:
        t = m.group(1)
    return t


WORD_NUM = {
    "zero": 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5, "six": 6,
    "seven": 7, "eight": 8, "nine": 9, "ten": 10, "eleven": 11, "twelve": 12,
    "thirteen": 13, "fourteen": 14, "fifteen": 15, "sixteen": 16,
    "seventeen": 17, "eighteen": 18, "nineteen": 19, "twenty": 20,
    "thirty": 30, "forty": 40, "fifty": 50, "sixty": 60, "seventy": 70,
    "eighty": 80, "ninety": 90, "hundred": 100, "thousand": 1000,
}
SCALES = {"hundred": 100, "thousand": 1000, "million": 1000000}

WORDNUM_RE = re.compile(
    r"\b((?:%s)(?:[\s-](?:%s|hundred|thousand|million))*)\b"
    % ("|".join(WORD_NUM), "|".join(WORD_NUM)),
    re.IGNORECASE,
)


def parse_word_number(phrase):
    """'twenty-three' -> 23, 'eight' -> 8, 'two thousand' -> 2000."""
    total, current = 0, 0
    for w in re.split(r"[\s-]+", phrase.lower()):
        if w in SCALES:
            scale = SCALES[w]
            if scale >= 1000:
                total += (current or 1) * scale
                current = 0
            else:
                current = (current or 1) * scale
        elif w in WORD_NUM:
            current += WORD_NUM[w]
        else:
            return None
    return total + current


def norm_date(tok):
    """Normalise the date spellings that appear in practice to MM-DD."""
    t = tok.strip().lower().rstrip(".,;:")
    # 2026-07-09  /  07-09
    m = re.fullmatch(r"(?:(\d{4})-)?(\d{1,2})-(\d{1,2})", t)
    if m:
        return "%02d-%02d" % (int(m.group(2)), int(m.group(3)))
    # 9 july / 9 july 2026 / 9th july
    m = re.fullmatch(r"(\d{1,2})(?:st|nd|rd|th)?\s+([a-z]+)(?:\s+\d{4})?", t)
    if m and m.group(2) in MONTHS:
        return "%02d-%02d" % (MONTHS[m.group(2)], int(m.group(1)))
    # july 9 / july 9th / july 9, 2026
    m = re.fullmatch(r"([a-z]+)\s+(\d{1,2})(?:st|nd|rd|th)?(?:,?\s+\d{4})?", t)
    if m and m.group(1) in MONTHS:
        return "%02d-%02d" % (MONTHS[m.group(1)], int(m.group(2)))
    # bare times, 02:28
    m = re.fullmatch(r"(\d{1,2}):(\d{2})", t)
    if m:
        return "%02d:%02d" % (int(m.group(1)), int(m.group(2)))
    return None


# ---------------------------------------------------------------- ledger load

class Ledger:
    def __init__(self, path):
        with open(path) as fh:
            d = json.load(fh)
        self.raw = d
        self.claims = {c["id"]: c for c in d["claims"]}

        self.numbers = set()
        self.dates = set()
        self.entities = set()

        all_claims = list(d["claims"])
        reg = d.get("regulatory_extension", {}).get("claims", [])
        all_claims.extend(reg)
        for c in reg:
            self.claims[c["id"]] = c

        for c in all_claims:
            a = c.get("atoms", {})
            for n in a.get("numbers", []):
                self.numbers.add(norm_number(n))
                wv = parse_word_number(n)
                if wv is not None:
                    self.numbers.add(str(wv))
            for dt in a.get("dates", []):
                nd = norm_date(dt)
                if nd:
                    self.dates.add(nd)
            for e in a.get("entities", []):
                self.entities.add(e.lower())
                # a shortened form of a known entity is still that entity:
                # "UK AI Security Institute" also licenses "AI Security Institute"
                parts = e.lower().split()
                for i in range(1, len(parts)):
                    self.entities.add(" ".join(parts[i:]))

        # quantities only: these take part in rounding comparisons and in
        # spelled-out numeral checks. Figures that come from a real claim atom
        # stay here even if they also appear in the generic small-number list.
        self.quantities = set(self.numbers)
        claim_numbers = set(self.numbers)

        # identifiers and generic small numbers are accepted but are not
        # meaningful rounding targets (a section number is not a magnitude)
        for n in (d.get("generic_number_allowlist", [])
                  + d.get("identifier_numbers", [])):
            nn = norm_number(n)
            self.numbers.add(nn)
            if nn not in claim_numbers:
                self.quantities.discard(nn)
        for e in d.get("generic_entity_allowlist", []) + d.get("technical_vocabulary", []):
            self.entities.add(e.lower())
            # allow individual words of multiword generic entities
            for w in e.split():
                self.entities.add(w.lower())

        # every year mentioned anywhere in the ledger is fine as a date
        for y in ("2025", "2026", "2027"):
            self.dates.add(y)

        self.banned = [(re.compile(b["pattern"]), b["why"])
                       for b in d.get("banned_claims", [])]
        self.attribution_required = [
            (re.compile(a["pattern"]), a["claim"], a["why"])
            for a in d.get("attribution_required", [])
        ]
        self.certainty = [m.lower() for m in d.get("certainty_markers", [])]
        self.cues = [c.lower() for c in d.get("attribution_cues", [])]
        self.counted_units = d.get("counted_units", [])
        self.closed_sets = [{
            "name": c["name"],
            "members": c["members"],
            "trigger": re.compile(c["trigger"], re.IGNORECASE),
            "candidate": re.compile(c["candidate"]),
            "why": c["why"],
        } for c in d.get("closed_sets", [])]

    # capitalised words that begin a clause but carry no attribution
    LEADING_NOISE = {
        "the", "their", "its", "our", "his", "her", "this", "that", "these",
        "those", "under", "strip", "over", "after", "before", "within", "when",
        "on", "at", "in", "from", "by", "per", "across", "into", "with",
        "to", "for", "against", "during", "through", "between", "via",
        "while", "since", "both", "each", "every", "one", "two", "no", "not",
        "an", "a", "and", "but", "so", "if", "then", "also", "only", "even",
        "most", "many", "some", "all", "here", "there", "what", "which", "who",
        "how", "why", "where", "yes", "yet", "still", "now", "next", "first",
        "second", "third", "last", "same", "such", "other", "another", "roughly",
        "around", "about", "nearly", "almost", "just", "very", "more", "less",
    }

    # A rounded figure is acceptable if it is hedged and stays within
    # ROUND_TOL of a ledger figure, in the direction the hedge claims.
    ROUND_TOL = 0.05
    HEDGE_DOWN = ("over", "more than", "at least", "upwards of", "north of", "above")
    HEDGE_UP = ("under", "fewer than", "less than", "up to", "below", "nearly", "almost")
    HEDGE_NEAR = ("about", "roughly", "around", "approximately", "some", "~", "circa")

    def rounding_ok(self, tok, sent, pos):
        """Return (accepted, reason_if_rejected)."""
        try:
            val = float(tok.rstrip("%x"))
        except ValueError:
            return False, ""
        prefix = sent[max(0, pos - 24):pos].lower()
        near = any(h in prefix for h in self.HEDGE_NEAR)
        down = any(h in prefix for h in self.HEDGE_DOWN)
        up = any(h in prefix for h in self.HEDGE_UP)
        if not (near or down or up):
            return False, ""
        for cand in self.quantities:
            try:
                cv = float(cand.rstrip("%x"))
            except ValueError:
                continue
            if cv == 0:
                continue
            err = (val - cv) / cv
            if near and abs(err) <= self.ROUND_TOL:
                return True, ""
            if down and -self.ROUND_TOL <= err <= 0:
                return True, ""
            if up and 0 <= err <= self.ROUND_TOL:
                return True, ""
        # hedged but outside tolerance: say by how much against the nearest atom
        best, berr = None, None
        for cand in self.quantities:
            try:
                cv = float(cand.rstrip("%x"))
            except ValueError:
                continue
            if cv == 0:
                continue
            e = abs(val - cv) / cv
            if berr is None or e < berr:
                best, berr = cand, e
        if best is not None and berr < 1.0:
            return False, ("hedged rounding of %s, but %.0f%% off (tolerance %.0f%%)"
                           % (best, berr * 100, self.ROUND_TOL * 100))
        return False, ""

    def _known_token(self, w):
        if w in self.entities:
            return True
        for suffix in ("'s", "s", "'"):
            if w.endswith(suffix) and w[: -len(suffix)] in self.entities:
                return True
        return False

    def entity_known(self, name):
        n = name.lower().strip().rstrip(".,;:'")
        if self._known_token(n):
            return True
        words = [w for w in re.split(r"[\s/&-]+", n) if w]
        if not words:
            return True
        # drop leading clause noise ("Their AI", "Under Article", "Strip the AI")
        while words and words[0] in self.LEADING_NOISE:
            words = words[1:]
        if not words:
            return True
        if self._known_token(" ".join(words)):
            return True
        # a hyphenated modifier of a known entity is still that entity:
        # "CyberGym-style", "OpenAI-registered", "German-language"
        if "-" in n:
            head = n.split("-")[0].strip()
            if head and self._known_token(head):
                return True
        # every remaining word independently known ("AI-assisted", "Hugging Face")
        if all(self._known_token(w) or w in self.LEADING_NOISE for w in words):
            return True
        return False


# ------------------------------------------------------------------ extraction

SENT_SPLIT = re.compile(r"(?<=[.!?])\s+")

NUM_RE = re.compile(r"(?<![\w.-])(~?\d[\d,]*(?:\.\d+)?\s?%?x?)(?![\w-])")

DATE_RE = re.compile(
    r"\b(?:"
    r"\d{4}-\d{1,2}-\d{1,2}"
    r"|\d{1,2}-\d{1,2}"
    r"|\d{1,2}(?:st|nd|rd|th)?\s+(?:%s)(?:\s+\d{4})?"
    r"|(?:%s)\s+\d{1,2}(?:st|nd|rd|th)?(?:,?\s+\d{4})?"
    r"|\d{1,2}:\d{2}"
    r")\b" % ("|".join(MONTHS), "|".join(MONTHS)),
    re.IGNORECASE,
)

# capitalised runs, allowing internal lowercase joiners and version-ish tokens
ENT_RE = re.compile(
    r"\b([A-Z][A-Za-z0-9.+-]*(?:\s+(?:of|the|for|and|de)\s+[A-Z][A-Za-z0-9.+-]*"
    r"|\s+[A-Z][A-Za-z0-9.+-]*)*)"
)

MD_STRIP = re.compile(r"(\*\*|[*_`>]|^\s*[-•]\s+|\[[^\]]*\]\([^)]*\))", re.MULTILINE)

HEADING = re.compile(r"^\s*#+\s")


def strip_front_matter(text):
    """Briefs carry a metadata block ending in a '---' rule. It is not prose."""
    lines = text.split("\n")
    for i, ln in enumerate(lines):
        if ln.strip() == "---":
            return "\n".join(lines[i + 1:])
    return text


def clean(text, drop_headings=True):
    lines = []
    for ln in strip_front_matter(text).split("\n"):
        if drop_headings and HEADING.match(ln):
            # headings are styling, not assertions, but their numerals still count
            ln = HEADING.sub("", ln)
            ln = "␟" + ln  # mark: entity check skips this line
        lines.append(ln)
    return MD_STRIP.sub(" ", "\n".join(lines))


def unwrap(text):
    """Join hard-wrapped lines inside a paragraph so a name is never split."""
    out, buf = [], []
    for ln in text.split("\n"):
        if not ln.strip() or ln.lstrip().startswith("\u241f") or ln.lstrip().startswith("|"):
            if buf:
                out.append(" ".join(buf))
                buf = []
            out.append(ln)
        else:
            buf.append(ln.strip())
    if buf:
        out.append(" ".join(buf))
    return "\n".join(out)


def sentences(text):
    for para in unwrap(text).split("\n"):
        para = para.strip()
        if not para:
            continue
        for s in SENT_SPLIT.split(para):
            s = s.strip()
            if s:
                yield s


# --------------------------------------------------------------------- checking

def check(text, ledger, path="<brief>", wordnum=True, closedset=True):
    findings = []
    body = clean(text)

    def add(kind, sent, detail, why=""):
        findings.append({
            "file": os.path.basename(path),
            "check": kind,
            "detail": detail,
            "why": why,
            "sentence": sent[:200],
        })

    for sent in sentences(body):
        is_heading = sent.startswith("␟")
        sent = sent.lstrip("␟").strip()
        if not sent:
            continue
        low = sent.lower()

        # spans covered by a recognised date, so we do not double-report their digits
        date_spans = []
        for m in DATE_RE.finditer(sent):
            date_spans.append((m.start(), m.end()))
            nd = norm_date(m.group(0))
            if nd is None:
                continue
            if nd not in ledger.dates:
                add("UNSUPPORTED_DATE", sent, m.group(0),
                    "no ledger claim carries this date")

        for m in NUM_RE.finditer(sent):
            if any(a <= m.start() < b for a, b in date_spans):
                continue
            tok = norm_number(m.group(1))
            if tok in ledger.numbers:
                continue
            # a percentage or multiplier of an allowed number
            if tok.endswith("%") and tok[:-1] in ledger.numbers:
                continue
            if tok.endswith("x") and tok[:-1] in ledger.numbers:
                continue
            ok, why = ledger.rounding_ok(tok, sent, m.start())
            if ok:
                continue
            add("UNSUPPORTED_NUMBER", sent, m.group(1),
                why or "no ledger claim carries this figure")

        # spelled-out numerals attached to a unit the ledger quantifies
        for m in (WORDNUM_RE.finditer(sent) if wordnum else ()):
            val = parse_word_number(m.group(1))
            if val is None or val <= 2:
                continue
            tail = sent[m.end():m.end() + 24].lower()
            unit = next((u for u in ledger.counted_units if tail.lstrip().startswith(u)), None)
            if unit is None:
                continue
            # measured against real quantities only: the generic small-number
            # allowlist must not launder "eight days" when the record says six
            if str(val) in ledger.quantities:
                continue
            add("UNSUPPORTED_NUMBER", sent, "%s %s" % (m.group(1), unit),
                "spelled-out figure not carried by any ledger claim")

        # closed-set roles: only these entities may fill this role
        for cs in (ledger.closed_sets if closedset else []):
            if not cs["trigger"].search(sent):
                continue
            for m in cs["candidate"].finditer(sent):
                found = m.group(0)
                if any(found.lower() == mem.lower() or found.lower() in mem.lower()
                       for mem in cs["members"]):
                    continue
                add("CLOSED_SET_VIOLATION", sent, "%s (role: %s)" % (found, cs["name"]),
                    cs["why"])

        for m in (() if is_heading else ENT_RE.finditer(sent)):
            name = m.group(1).strip().rstrip(".,;:'")
            # ignore a capitalised first word of a sentence when it is a common word
            if m.start() == 0 and " " not in name and len(name) < 12:
                if not ledger.entity_known(name):
                    continue
            if len(name) < 3:
                continue
            if ledger.entity_known(name):
                continue
            add("UNSUPPORTED_ENTITY", sent, name,
                "named entity not present in any ledger claim")

        for rx, why in ledger.banned:
            m = rx.search(sent)
            if m:
                add("BANNED_CLAIM", sent, m.group(0), why)

        for rx, cid, why in ledger.attribution_required:
            if rx.search(sent) and not any(c in low for c in ledger.cues):
                add("UNATTRIBUTED", sent, cid, why)

        if any(c in low for c in ledger.certainty):
            for rx, cid, why in ledger.attribution_required:
                if rx.search(sent):
                    add("OVERSTATED_CERTAINTY", sent, cid,
                        "certainty marker applied to a %s claim"
                        % ledger.claims[cid]["status"].lower())

    # de-duplicate identical (check, detail) pairs within a file
    seen, uniq = set(), []
    for f in findings:
        k = (f["check"], f["detail"].lower())
        if k in seen:
            continue
        seen.add(k)
        uniq.append(f)
    return uniq


def word_count(text):
    return len(re.findall(r"\b\w+\b", clean(text)))


# ------------------------------------------------------------------------- cli

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("briefs", nargs="+")
    ap.add_argument("--factbase", default="factbase.json")
    ap.add_argument("--csv", default=None)
    ap.add_argument("--quiet", action="store_true")
    ap.add_argument("--no-wordnum", action="store_true",
                    help="ablation: disable spelled-out numeral parsing (gate v0.1)")
    ap.add_argument("--no-closedset", action="store_true",
                    help="ablation: disable closed-set role checks (gate v0.1)")
    args = ap.parse_args()

    ledger = Ledger(args.factbase)

    paths = []
    for pattern in args.briefs:
        hits = sorted(glob.glob(pattern))
        paths.extend(hits if hits else [pattern])

    rows, failed = [], 0
    for p in paths:
        with open(p) as fh:
            text = fh.read()
        findings = check(text, ledger, p,
                         wordnum=not args.no_wordnum,
                         closedset=not args.no_closedset)
        wc = word_count(text)
        rows.append({
            "file": os.path.basename(p),
            "words": wc,
            "findings": len(findings),
            "per_1000_words": round(len(findings) * 1000.0 / wc, 1) if wc else 0.0,
            **{k: sum(1 for f in findings if f["check"] == k) for k in (
                "UNSUPPORTED_NUMBER", "UNSUPPORTED_DATE", "UNSUPPORTED_ENTITY",
                "BANNED_CLAIM", "UNATTRIBUTED", "OVERSTATED_CERTAINTY",
                "CLOSED_SET_VIOLATION")},
        })
        if findings:
            failed += 1
        if not args.quiet:
            print("\n=== %s  (%d words, %d findings)" % (os.path.basename(p), wc, len(findings)))
            for f in findings:
                print("  [%s] %s" % (f["check"], f["detail"]))
                if f["why"]:
                    print("      why: %s" % f["why"])
                print("      in : %s" % f["sentence"])

    print("\n--- summary")
    hdr = ["file", "words", "findings", "per_1000_words"]
    print("  " + "  ".join(h.ljust(34 if h == "file" else 14) for h in hdr))
    for r in rows:
        print("  " + "  ".join(str(r[h]).ljust(34 if h == "file" else 14) for h in hdr))
    print("  %d/%d briefs passed" % (len(rows) - failed, len(rows)))

    if args.csv:
        with open(args.csv, "w", newline="") as fh:
            w = csv.DictWriter(fh, fieldnames=list(rows[0].keys()))
            w.writeheader()
            w.writerows(rows)
        print("  wrote %s" % args.csv)

    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())