import re, os, sys, glob, json
ROOT = r"C:\Users\decid\Documents\projects\spt-core\.worktrees\hertz-census\crates"
files = glob.glob(ROOT + r"\spt\src\**\*.rs", recursive=True) + glob.glob(ROOT + r"\spt-daemon\src\**\*.rs", recursive=True)

def strip(src):
    # replace string/char/comment contents with spaces, keep newlines
    out = []; i = 0; n = len(src)
    while i < n:
        c = src[i]
        if src.startswith("//", i):
            j = src.find("\n", i);  j = n if j < 0 else j
            out.append(" " * (j - i)); i = j; continue
        if src.startswith("/*", i):
            depth = 0; j = i
            while j < n:
                if src.startswith("/*", j): depth += 1; j += 2; continue
                if src.startswith("*/", j):
                    depth -= 1; j += 2
                    if depth == 0: break
                    continue
                j += 1
            out.append(re.sub(r"[^\n]", " ", src[i:j])); i = j; continue
        m = re.match(r'b?r(#*)"', src[i:i+10])
        if m and (i == 0 or not (src[i-1].isalnum() or src[i-1] == '_')):
            hashes = m.group(1); end = '"' + hashes
            j = src.find(end, i + len(m.group(0)))
            j = n if j < 0 else j + len(end)
            out.append('"' + re.sub(r"[^\n]", " ", src[i+1:j-1]) + '"'); i = j; continue
        if c == '"':
            j = i + 1
            while j < n and src[j] != '"':
                j += 2 if src[j] == '\\' else 1
            j += 1
            out.append('"' + re.sub(r"[^\n]", " ", src[i+1:j-1]) + '"'); i = j; continue
        if c == "'":
            m = re.match(r"'(\\.|\\u\{[0-9a-fA-F]+\}|\\x..|[^\\'])'", src[i:i+12])
            if m:
                out.append(" " * len(m.group(0))); i += len(m.group(0)); continue
        out.append(c); i += 1
    return "".join(out)

FN = re.compile(r"\bfn\s+([A-Za-z_][A-Za-z0-9_]*)")
funcs = []  # (file, name, start, end, is_test)
texts = {}
for f in files:
    raw = open(f, encoding="utf-8").read()
    s = strip(raw)
    texts[f] = (raw, s)
    lines_off = [0]
    for i, ch in enumerate(s):
        if ch == "\n": lines_off.append(i + 1)
    import bisect
    def lineno(pos): return bisect.bisect_right(lines_off, pos)
    # test regions: #[cfg(test)] followed by mod X {  -> span
    test_spans = []
    for m in re.finditer(r"#\[cfg\(test\)\]\s*(?:#\[[^\]]*\]\s*)*(?:pub(?:\([a-z]+\))?\s+)?(mod|fn|impl|use|struct|const|static)\b", s):
        k = m.end()
        # find first { or ;
        j = k
        while j < len(s) and s[j] not in "{;": j += 1
        if j < len(s) and s[j] == "{":
            d = 0; e = j
            while e < len(s):
                if s[e] == "{": d += 1
                elif s[e] == "}":
                    d -= 1
                    if d == 0: break
                e += 1
            test_spans.append((m.start(), e))
    for m in FN.finditer(s):
        k = m.end(); j = k; depth = 0
        while j < len(s):
            ch = s[j]
            if ch in "(<[": depth += 1
            elif ch in ")>]": depth -= 1
            if ch == "{" and depth <= 0: break
            if ch == ";" and depth <= 0: j = -1; break
            j += 1
        if j < 0 or j >= len(s): continue
        d = 0; e = j
        while e < len(s):
            if s[e] == "{": d += 1
            elif s[e] == "}":
                d -= 1
                if d == 0: break
            e += 1
        is_test = any(a <= m.start() <= b for a, b in test_spans)
        # also a #[test] attribute directly before
        pre = s[max(0, m.start()-200):m.start()]
        if re.search(r"#\[(tokio::)?test\]\s*(#\[[^\]]*\]\s*)*$", pre): is_test = True
        funcs.append(dict(file=os.path.relpath(f, ROOT), name=m.group(1), s=m.start(), e=e, line=lineno(m.start()), eline=lineno(e), test=is_test))
    for fd in funcs:
        if fd["file"] == os.path.relpath(f, ROOT): fd["_ln"] = None
    texts[f] = (raw, s, lines_off)

def enclosing(file, pos):
    best = None
    for fd in funcs:
        if fd["file"] == file and fd["s"] <= pos <= fd["e"]:
            if best is None or fd["s"] > best["s"]: best = fd
    return best

json.dump([{k: v for k, v in fd.items() if not k.startswith("_")} for fd in funcs], open(os.path.join(os.path.dirname(__file__), "funcs.json"), "w"))

import bisect
def callers(name, restrict_file=None, qual=None):
    """return set of enclosing fns that reference `name` followed by ( or as path."""
    res = []
    pat = re.compile(r"(?<![A-Za-z0-9_])" + (re.escape(qual) + r"\s*::\s*" if qual else "") + re.escape(name) + r"(?![A-Za-z0-9_])")
    for f in files:
        rel = os.path.relpath(f, ROOT)
        if restrict_file and rel != restrict_file: continue
        raw, s, lo = texts[f]
        for m in pat.finditer(s):
            # skip definition
            if re.search(r"\bfn\s+$", s[max(0, m.start()-20):m.start()]) and not qual: continue
            if qual is None and re.search(r"\bfn\s*$", s[max(0,m.start()-10):m.start()]): continue
            enc = enclosing(rel, m.start())
            ln = bisect.bisect_right(lo, m.start())
            res.append((rel, ln, enc))
    return res

if __name__ == "__main__":
    for arg in sys.argv[1:]:
        parts = arg.split("@")
        name = parts[0]; rf = parts[1] if len(parts) > 1 and parts[1] else None; q = parts[2] if len(parts) > 2 else None
        print("==", arg)
        for rel, ln, enc in callers(name, rf, q):
            print(f"  {rel}:{ln}  in {enc['name'] if enc else '<top>'}{' [TEST]' if enc and enc['test'] else ''} (fn@{enc['line'] if enc else '-'})")