"""aider-repo-map-edit-formats-explained.py — a miniature "repo map" for Python code.

What it does: walks a directory of .py files, collects top-level function and class
signatures with the ast module, builds an import graph between the files (file A imports
module B => edge A -> B), ranks the files with a small PageRank (power iteration, no
third-party packages), and prints a signature-only map trimmed to a token budget.

The token budget uses the chars/4 heuristic (the rough "1 token ~ 4 characters" figure
from the Claude pricing FAQ); it is NOT a tokenizer, so treat budget numbers as estimates.

Inputs:  a directory path, optional --map-tokens N (default 1000), optional --top K.
Run:     python aider-repo-map-edit-formats-explained.py <dir> [--map-tokens 1000] [--top 20]
Needs:   Python 3.10+ (tested on 3.13.12 on Windows). Standard library only.
"""
from __future__ import annotations

import argparse
import ast
import os
import sys


def module_name(root: str, path: str) -> str:
    rel = os.path.relpath(path, root).replace(os.sep, "/")
    return rel[:-3].replace("/", ".")


def collect(root: str) -> dict[str, dict]:
    """Return {module: {"path":..., "sigs": [...], "imports": set()}} for every .py under root."""
    files: dict[str, dict] = {}
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = [d for d in dirnames if not d.startswith((".", "__pycache__", "prod_v1", "node_modules"))]
        for name in filenames:
            if not name.endswith(".py"):
                continue
            path = os.path.join(dirpath, name)
            try:
                tree = ast.parse(open(path, encoding="utf-8").read(), filename=path)
            except (SyntaxError, UnicodeDecodeError) as e:
                print(f"skip {path}: {e}", file=sys.stderr)
                continue
            sigs: list[str] = []
            imports: set[str] = set()
            for node in tree.body:
                if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    sigs.append(f"def {node.name}({ast.unparse(node.args)})")
                elif isinstance(node, ast.ClassDef):
                    bases = ", ".join(ast.unparse(b) for b in node.bases)
                    sigs.append(f"class {node.name}({bases})" if bases else f"class {node.name}")
                    for sub in node.body:
                        if isinstance(sub, (ast.FunctionDef, ast.AsyncFunctionDef)):
                            sigs.append(f"    def {sub.name}({ast.unparse(sub.args)})")
            for node in ast.walk(tree):
                if isinstance(node, ast.Import):
                    imports.update(a.name.split(".")[0] for a in node.names)
                elif isinstance(node, ast.ImportFrom) and node.module:
                    imports.add(node.module.split(".")[0])
            files[module_name(root, path)] = {"path": path, "sigs": sigs, "imports": imports}
    return files


def pagerank(files: dict[str, dict], damping: float = 0.85, iters: int = 50) -> dict[str, float]:
    """PageRank over the import graph: an edge A -> B when A imports module B (B in the set)."""
    names = list(files)
    short = {n.split(".")[-1]: n for n in names}  # allow `import engine` to hit `engine`
    edges = {n: set() for n in names}
    for n, info in files.items():
        for imp in info["imports"]:
            target = files.get(imp) and imp or short.get(imp)
            if target and target != n:
                edges[n].add(target)
    n_nodes = len(names)
    rank = {n: 1.0 / n_nodes for n in names}
    for _ in range(iters):
        new = {n: (1.0 - damping) / n_nodes for n in names}
        for src, targets in edges.items():
            share = rank[src] / len(targets) if targets else rank[src] / n_nodes
            for t in (targets or names):
                new[t] += damping * share
        rank = new
    return rank


def render(files: dict[str, dict], rank: dict[str, float], map_tokens: int, top: int) -> str:
    """Emit files best-first, whole file blocks only while the chars/4 budget holds."""
    out: list[str] = []
    used = 0
    for name in sorted(rank, key=rank.get, reverse=True)[:top]:
        info = files[name]
        block = [f"{os.path.relpath(info['path'])}:  (rank {rank[name]:.3f})"] + [f"  {s}" for s in info["sigs"]] + [""]
        text = "\n".join(block)
        cost = len(text) // 4  # heuristic tokens
        if used + cost > map_tokens:
            out.append(f"... {os.path.relpath(info['path'])} omitted ({cost} est. tokens over budget)")
            continue
        used += cost
        out.append(text)
    out.append(f"[map: ~{used} est. tokens of {map_tokens} budget, chars/4 heuristic]")
    return "\n".join(out)


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("root")
    ap.add_argument("--map-tokens", type=int, default=1000)
    ap.add_argument("--top", type=int, default=20)
    a = ap.parse_args()
    files = collect(a.root)
    if not files:
        sys.exit("no .py files found")
    rank = pagerank(files)
    print(render(files, rank, a.map_tokens, a.top))


if __name__ == "__main__":
    main()