#!/usr/bin/env python3
"""Count the shape functions the `declare!` form has to cover.

The population is the denominator of phase 1 of the declaration transition
(GoingsOn `65d99281`, wiki `quasi-declare-form`). Five counts of it existed
before this script -- 289, 375, 402, 425, 487 -- and none reproduced, because
none had written its predicate down. This is the predicate, executable.

THE PREDICATE. A Rust `fn` item with a body, in one of the three shape
directories, outside `tests.rs` and `parity.rs`, outside any `tests/` path
segment, outside every inline `#[cfg(test)]` block, whose return type -- after
stripping `Vec`, `Option`, `Result`, `Box`, `Arc`, `Rc` and `Cow` -- names a
`quasi_router` description type, with `use ... as` aliases resolved and
file-local types of the same name excluded.

Comments, string literals, raw strings and char literals are blanked in place
before anything is matched, so a type named in a doc comment is not a
construction. That is the defect that produced two of the five old numbers.

    python3 scripts/population.py            # the count, per tree
    python3 scripts/population.py --list     # one `path:line name -> type` per line
    python3 scripts/population.py --json     # the same, machine-readable
    python3 scripts/population.py --selftest # the predicate against known answers

Run `--selftest` before quoting a number from this. Every case in it is a defect
that actually happened: a `#[cfg(test)]` stripper that brace-matched from the
next `{` blanked four real functions in goingson, a doc-comment mention of a
vocabulary type was counted as a construction, and a tuple return was missed.

Paths are read relative to `~/Code` by default; pass `--root` for a tree
somewhere else.
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path

# The three shape directories, relative to the root. A fourth app joining the
# transition is one line here and a re-run.
SHAPE_DIRS = [
    "MNW/server/src/quasi",
    "Apps/goingson/src-tauri/src/quasi",
    "Apps/audiofiles/crates/audiofiles-browser/src/quasi",
]

# The vocabulary. Every public description type `quasi_router` exports that a
# shape function can return, plus the two enums that are returned as leaves.
# `Node` alone is a fifth of the population.
VOCABULARY = {
    "Node", "Screen", "Slot", "Run", "Field", "Row", "Cell", "Cells",
    "Act", "Action", "Choice", "Column", "Tag", "Part", "Figure",
    "Placed", "Consult", "Meter", "Image", "Lexeme", "Document", "Feed",
}

# Returned by a handful of leaf suppliers that section 6 of `quasi-declare-form`
# defers, so they sit outside the ratified population and inside `--wide`.
# Adding all five moves the count by 10 on the trees as they stand.
DEFERRED_LEAVES = {"Rest", "Candidate", "Accepted", "RegionKind", "Chrome"}

# Wrappers stripped before the name is looked up. A `-> Result<Vec<Row>, E>` is
# a shape function returning rows; the wrapper says how it fails, not what it is.
WRAPPERS = {"Vec", "Option", "Result", "Box", "Arc", "Rc", "Cow"}


def blank_noncode(src: str) -> str:
    """Replace every comment, string, raw string and char literal with spaces.

    Newlines are kept so line numbers still hold. Nothing is deleted, so every
    offset in the returned text is the offset in the original.
    """
    out = list(src)
    i, n = 0, len(src)
    while i < n:
        c = src[i]
        # Raw string, with any hash count: r"..", r#".."#, br##".."##
        m = re.match(r'(?:b)?r(#*)"', src[i:])
        if m and (i == 0 or not (src[i - 1].isalnum() or src[i - 1] == "_")):
            hashes = m.group(1)
            close = '"' + hashes
            end = src.find(close, i + m.end())
            end = n if end < 0 else end + len(close)
            for j in range(i, end):
                if out[j] != "\n":
                    out[j] = " "
            i = end
            continue
        if c == "/" and i + 1 < n and src[i + 1] == "/":
            end = src.find("\n", i)
            end = n if end < 0 else end
            for j in range(i, end):
                out[j] = " "
            i = end
            continue
        if c == "/" and i + 1 < n and src[i + 1] == "*":
            depth, j = 1, i + 2
            while j < n and depth:
                if src.startswith("/*", j):
                    depth += 1
                    j += 2
                elif src.startswith("*/", j):
                    depth -= 1
                    j += 2
                else:
                    j += 1
            for k in range(i, j):
                if out[k] != "\n":
                    out[k] = " "
            i = j
            continue
        if c == '"':
            j = i + 1
            while j < n:
                if src[j] == "\\":
                    j += 2
                    continue
                if src[j] == '"':
                    j += 1
                    break
                j += 1
            for k in range(i, min(j, n)):
                if out[k] != "\n":
                    out[k] = " "
            i = j
            continue
        if c == "'":
            # A char literal, not a lifetime: `'a` is a lifetime, `'a'` is a char.
            m = re.match(r"'(?:\\.|[^\\'])'", src[i:])
            if m:
                for k in range(i, i + m.end()):
                    out[k] = " "
                i += m.end()
                continue
        i += 1
    return "".join(out)


def strip_cfg_test(src: str) -> str:
    """Blank every `#[cfg(test)]` item that is followed by an inline block.

    Only a block: `#[cfg(test)] mod tests;` names a file that is excluded by
    path anyway, and blanking to the next `}` there would eat the rest of the
    file.
    """
    out = list(src)
    for m in re.finditer(r"#\[cfg\(test\)\]", src):
        j = m.end()
        while j < len(src) and src[j] not in "{;":
            j += 1
        if j >= len(src) or src[j] == ";":
            continue
        depth, k = 0, j
        while k < len(src):
            if src[k] == "{":
                depth += 1
            elif src[k] == "}":
                depth -= 1
                if depth == 0:
                    k += 1
                    break
            k += 1
        for p in range(m.start(), k):
            if out[p] != "\n":
                out[p] = " "
    return "".join(out)


def local_types(src: str) -> set[str]:
    """Types the file declares itself, which shadow the vocabulary name."""
    return set(
        re.findall(r"^\s*(?:pub(?:\([^)]*\))?\s+)?(?:struct|enum|type|union)\s+(\w+)",
                   src, re.M)
    )


def imported(src: str) -> tuple[set[str], dict[str, str]]:
    """Names brought in from `quasi_router`, and every `as` alias in the file.

    Returns (names in scope from quasi_router, alias -> real name). Braced and
    single-item `use` forms both, at any nesting depth: the only thing that
    matters is whether the path starts at `quasi_router`.
    """
    names: set[str] = set()
    aliases: dict[str, str] = {}
    for m in re.finditer(r"use\s+([^;]+);", src):
        body = m.group(1)
        if "quasi_router" not in body and "quasi_webview" not in body:
            continue
        for leaf in re.findall(r"(\w+)(?:\s+as\s+(\w+))?", body):
            name, alias = leaf
            if name in VOCABULARY:
                names.add(alias or name)
                if alias:
                    aliases[alias] = name
    return names, aliases


def split_top(s: str) -> list[str]:
    """Split on commas at bracket depth zero."""
    out, depth, start = [], 0, 0
    for i, c in enumerate(s):
        if c in "<([":
            depth += 1
        elif c in ">)]":
            depth -= 1
        elif c == "," and depth == 0:
            out.append(s[start:i])
            start = i + 1
    out.append(s[start:])
    return [p for p in (x.strip() for x in out) if p]


def head_type(ret: str) -> str | None:
    """The vocabulary name a return type resolves to, or None.

    Strips the wrappers outside-in, takes the last `::` segment, and stops at
    the first name that is not a wrapper. `Result<Vec<Node>, RouteError>` is
    `Node`; `Result<(), E>` is nothing.
    """
    ret = ret.strip()
    for _ in range(6):
        ret = ret.strip().lstrip("&").strip()
        ret = re.sub(r"^'\w+\s+", "", ret)
        m = re.match(r"([\w:]+)\s*<(.*)>$", ret, re.S)
        if not m:
            break
        head = m.group(1).split("::")[-1]
        if head not in WRAPPERS:
            return head
        inner = m.group(2)
        # `Result<T, E>`: the first argument at depth zero is the payload.
        depth, cut = 0, len(inner)
        for i, c in enumerate(inner):
            if c in "<([":
                depth += 1
            elif c in ">)]":
                depth -= 1
            elif c == "," and depth == 0:
                cut = i
                break
        ret = inner[:cut]
    ret = ret.strip().lstrip("&").strip()
    ret = re.sub(r"^'\w+\s+", "", ret)
    # A tuple return. One site in the trees as they stand (goingson
    # `data.rs:380 csv_rows`), and `shaped` names one type, so the form
    # defers it -- but it is a shape function and it counts.
    if ret.startswith("(") and ret.endswith(")"):
        for part in split_top(ret[1:-1]):
            head = head_type(part)
            if head in VOCABULARY:
                return head
        return None
    name = ret.split("::")[-1].strip()
    return name or None


def match_angle(src: str, start: int) -> int:
    """Index just past the `>` closing the `<` at `start`."""
    depth, i = 0, start
    while i < len(src):
        if src[i] == "<":
            depth += 1
        elif src[i] == ">":
            depth -= 1
            if depth == 0:
                return i + 1
        elif src[i] == ";":
            return -1
        i += 1
    return -1


def match_paren(src: str, start: int) -> int:
    depth, i = 0, start
    while i < len(src):
        if src[i] == "(":
            depth += 1
        elif src[i] == ")":
            depth -= 1
            if depth == 0:
                return i + 1
        i += 1
    return -1


def match_brace(src: str, start: int) -> int:
    """Index just past the `}` closing the `{` at `start`, or -1.

    `src` must already be through `blank_noncode`, so a brace inside a string
    or a comment cannot unbalance the count.
    """
    depth = 0
    i = start
    while i < len(src):
        if src[i] == "{":
            depth += 1
        elif src[i] == "}":
            depth -= 1
            if depth == 0:
                return i + 1
        i += 1
    return -1


def shapes_in(path: Path) -> list[dict]:
    raw = path.read_text(encoding="utf-8", errors="replace")
    src = strip_cfg_test(blank_noncode(raw))
    scope, aliases = imported(raw)
    shadowed = local_types(src)
    found = []
    for m in re.finditer(r"\bfn\s+(\w+)\s*(?=[<(])", src):
        name = m.group(1)
        i = m.end()
        if src[i] == "<":
            i = match_angle(src, i)
            if i < 0:
                continue
            while i < len(src) and src[i].isspace():
                i += 1
        if i >= len(src) or src[i] != "(":
            continue
        i = match_paren(src, i)
        if i < 0:
            continue
        # Return type runs to the body's `{`, the `where` clause, or a `;`.
        j = i
        depth = 0
        while j < len(src):
            c = src[j]
            if c in "<([":
                depth += 1
            elif c in ">)]":
                depth -= 1
            elif c == "{" and depth <= 0:
                break
            elif c == ";" and depth <= 0:
                j = -1
                break
            j += 1
        if j < 0 or j >= len(src):
            continue  # no body: a trait signature or an extern declaration
        sig = src[i:j]
        if re.search(r"\bwhere\b", sig):
            sig = sig[: sig.index("where")]
        sig = sig.strip()
        if not sig.startswith("->"):
            continue  # returns unit
        ret = sig[2:].strip()
        head = head_type(ret)
        if head is None:
            continue
        real = aliases.get(head, head)
        if real not in VOCABULARY:
            continue
        if head in shadowed:
            continue  # a file-local type wearing a vocabulary name
        if head not in scope and real not in scope:
            # Reached by path (`quasi_router::Node`) rather than by import.
            if "quasi_router" not in ret and "quasi_webview" not in ret:
                continue
        end = match_brace(src, j)
        start_line = src.count("\n", 0, m.start()) + 1
        # Body length in lines, signature included. Used to stratify a probe
        # sample (`scripts/probe-sample.py`); it does not enter the predicate,
        # so the count is unaffected.
        lines = (src.count("\n", m.start(), end) + 1) if end > 0 else 1
        found.append(
            {"file": str(path), "line": start_line,
             "fn": name, "returns": " ".join(ret.split()), "lines": lines}
        )
    return found


# The cases the implementation has been wrong on, as executable assertions.
# Every one is a defect that actually happened: a naive `#[cfg(test)]` stripper
# ate four real functions, a doc-comment mention was counted as a construction,
# and a tuple return was missed. `--selftest` is how those stay fixed.
SELFTEST = [
    # (source, expected [(fn, head)])
    # A plain shape.
    ("use quasi_router::Node;\nfn body(cx: &Cx) -> Node { Node::page(\"x\") }",
     [("body", "Node")]),
    # Wrappers are stripped outside-in; `Result`'s payload is its first argument.
    ("use quasi_router::{Node, Row};\n"
     "fn a(c: &C) -> Result<Vec<Node>, E> { todo!() }\n"
     "fn b(c: &C) -> Option<Row> { todo!() }",
     [("a", "Node"), ("b", "Row")]),
    # `Result<(), E>` names no vocabulary type and is not a shape.
    ("use quasi_router::Node;\nfn a(c: &C) -> Result<(), E> { todo!() }", []),
    # A doc comment naming a vocabulary type is not a construction. This is the
    # defect behind two of the five irreproducible counts.
    ("use quasi_router::Node;\n/// Returns a Node, eventually.\nfn a(c: &C) -> u32 { 0 }",
     []),
    # Nor is a string literal.
    ("use quasi_router::Node;\nfn a(c: &C) -> u32 { let s = \"fn x() -> Node\"; 0 }",
     []),
    # An inline `#[cfg(test)]` block is stripped ...
    ("use quasi_router::Node;\n"
     "fn real(c: &C) -> Node { todo!() }\n"
     "#[cfg(test)]\nmod tests {\n  fn fake(c: &C) -> Node { todo!() }\n}\n",
     [("real", "Node")]),
    # ... but `#[cfg(test)] mod tests;` has no block, and brace-matching from
    # the next `{` would blank an arbitrary later region. Four real goingson
    # functions were lost to exactly this.
    ("use quasi_router::Node;\n"
     "#[cfg(test)]\nmod tests;\n"
     "fn survives(c: &C) -> Node { todo!() }\n",
     [("survives", "Node")]),
    # A signature with no body is a trait method or an extern, not a shape.
    ("use quasi_router::Node;\ntrait T { fn a(&self) -> Node; }", []),
    # A returned unit is not a shape.
    ("use quasi_router::Node;\nfn a(c: &C) {}", []),
    # A generic parameter list before the arguments must not break the parse.
    ("use quasi_router::Node;\nfn a<T: Into<String>>(t: T) -> Node { todo!() }",
     [("a", "Node")]),
    # A `where` clause sits between the return type and the body.
    ("use quasi_router::Node;\nfn a<T>(t: T) -> Node where T: Into<String> { todo!() }",
     [("a", "Node")]),
    # A tuple return: one site in the trees, and `shaped` names one type, but it
    # is a shape function and it counts.
    ("use quasi_router::Cells;\n"
     "fn a(c: &C) -> (Vec<&'static str>, Vec<Cells>) { todo!() }",
     [("a", "Cells")]),
    # A file-local type wearing a vocabulary name shadows it.
    ("use quasi_router::Node;\nstruct Row;\nfn a(c: &C) -> Row { Row }", []),
    # Reached by path rather than by import.
    ("fn a(c: &C) -> quasi_router::Node { todo!() }", [("a", "Node")]),
    # A vocabulary name that was never imported and is not path-qualified is
    # somebody else's type.
    ("fn a(c: &C) -> Node { todo!() }", []),
    # An alias is resolved for MEMBERSHIP and reported AS WRITTEN. Both halves
    # are deliberate: the shape counts, and `--list` says what the source says,
    # which is why the record's own return-type histogram lists "Described
    # (Screen alias) 18" as its own row rather than folding it into Screen.
    ("use quasi_router::Screen as Described;\nfn a(c: &C) -> Described { todo!() }",
     [("a", "Described")]),
    # `Response` is quasi's transport, not its description vocabulary. 566
    # further functions in these directories return one; admitting them roughly
    # doubles the count.
    ("use quasi_router::Response;\nfn a(c: &C) -> Response { todo!() }", []),
]


def selftest() -> int:
    """Run the predicate over synthetic sources with known answers."""
    import tempfile

    failures = 0
    for n, (src, expected) in enumerate(SELFTEST, 1):
        with tempfile.TemporaryDirectory() as d:
            f = Path(d) / "case.rs"
            f.write_text(src, encoding="utf-8")
            got = [(h["fn"], head_type(h["returns"]) or "?") for h in shapes_in(f)]
            # head_type is re-applied for the label only; membership already
            # filtered. Compare as sets, since order is not the claim.
            want = {(a, b) for a, b in expected}
            have = {(a, b) for a, b in got}
            if have != want:
                failures += 1
                print(f"case {n} FAILED")
                print(f"  source:   {src!r}")
                print(f"  expected: {sorted(want)}")
                print(f"  got:      {sorted(have)}")
    total = len(SELFTEST)
    if failures:
        print(f"\n{failures} of {total} cases failed.")
        return 1
    print(f"{total} of {total} cases pass.")
    return 0


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--root", default=str(Path.home() / "Code"))
    ap.add_argument("--list", action="store_true")
    ap.add_argument("--json", action="store_true")
    ap.add_argument("--selftest", action="store_true",
                    help="run the predicate over synthetic sources with known answers")
    ap.add_argument("--wide", action="store_true",
                    help="also count the deferred leaf suppliers (Rest, Candidate, "
                         "Accepted, RegionKind, Chrome)")
    args = ap.parse_args()

    if args.selftest:
        return selftest()

    if args.wide:
        VOCABULARY.update(DEFERRED_LEAVES)

    root = Path(args.root)
    per_tree: dict[str, list[dict]] = {}
    missing = []
    for rel in SHAPE_DIRS:
        d = root / rel
        if not d.is_dir():
            missing.append(rel)
            continue
        hits: list[dict] = []
        for f in sorted(d.rglob("*.rs")):
            parts = f.relative_to(root).parts
            if f.name in ("tests.rs", "parity.rs") or "tests" in parts:
                continue
            hits.extend(shapes_in(f))
        for h in hits:
            h["file"] = str(Path(h["file"]).relative_to(root))
        per_tree[rel] = hits

    total = sum(len(v) for v in per_tree.values())

    if args.json:
        json.dump({"total": total,
                   "per_tree": {k: len(v) for k, v in per_tree.items()},
                   "shapes": [h for v in per_tree.values() for h in v]},
                  sys.stdout, indent=1)
        print()
        return 0

    if args.list:
        for rel, hits in per_tree.items():
            for h in hits:
                print(f"{h['file']}:{h['line']} {h['fn']} -> {h['returns']}")

    for rel, hits in per_tree.items():
        print(f"{len(hits):5}  {rel}")
    print(f"{total:5}  TOTAL")
    for rel in missing:
        print(f"    ?  {rel} (not present under {root})", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())