#!/usr/bin/env python3
"""rtl_index.py — RTL signal-flow index for cross-module tracing.

Builds a module-fact database from Verilog/SystemVerilog sources:
  ports (ANSI & non-ANSI), nets, continuous assigns, always-block drivers,
  instantiations with named port connections, and cross-boundary edges.

pyslang 11 API. Index build is O(source); queries are in-memory graph walks.
"""
import os, re, json, glob
import pyslang
from pyslang import syntax as sl

IDENT = re.compile(r"^[A-Za-z_][A-Za-z0-9_$]*$")
ASSIGN_KINDS = ("AssignmentExpression", "NonblockingAssignmentExpression")


def _txt(node):
    t = getattr(node, "valueText", "")
    return t if t else (str(node).strip() if node is not None else "")


def _syntax_of(val_node):
    s = getattr(val_node, "syntax", None)
    return s if s is not None else val_node


def ids_from(expr_syntax):
    out = []
    def walk(n):
        if n is None:
            return
        if n.kind in (sl.SyntaxKind.IdentifierName, sl.SyntaxKind.IdentifierSelectName):
            ident = getattr(n, "identifier", None)
            if ident is not None:
                t = _txt(ident)
                if IDENT.match(t):
                    out.append(t)
            # IdentifierSelectName has .target identifier inside
            tgt = getattr(n, "target", None)
            if tgt is not None:
                walk(tgt)
            return
        for c in n:
            if hasattr(c, "__iter__"):
                walk(c)
    walk(expr_syntax)
    return out


def _lhs_names(expr_syntax):
    """names written by an assignment LHS (plain idents only, strip selects)"""
    out = []
    def walk(n):
        if n is None:
            return
        if n.kind == sl.SyntaxKind.IdentifierName:
            t = _txt(n.identifier)
            if IDENT.match(t):
                out.append(t)
            return
        for c in n:
            if hasattr(c, "__iter__"):
                walk(c)
    walk(expr_syntax)
    return out


class RtlIndex:
    def __init__(self):
        self.sm = pyslang.SourceManager()
        self.db = {}        # module -> facts
        self.edges = []     # (parent_mod, parent_net, child_mod, child_port, child_dir, inst_name)
        self.loc = {}       # (module, signal) -> "file:line"

    # ---------- build ----------
    def add_tree(self, path):
        tree = sl.SyntaxTree.fromFile(path, self.sm)
        root = tree.root
        for m in root.members:
            if m.kind == sl.SyntaxKind.ModuleDeclaration:
                info = self._scan_module(m)
                self.db[info["name"]] = info

    def _record_loc(self, module, name, node):
        try:
            if type(node).__name__ == "Token":
                tok = node
            else:
                tok = node.getFirstToken()
            if tok is not None:
                self.loc[(module, name)] = (
                    os.path.basename(self.sm.getFileName(tok.location)),
                    self.sm.getLineNumber(tok.location))
        except Exception:
            pass

    def _scan_module(self, mod):
        name = _txt(mod.header.name)
        info = dict(name=name, ports=[], nets=[], assigns=[], always=[], insts=[])
        self._record_loc(name, name, mod.header.name)

        # ---- ports (ANSI) ----
        pl = mod.header.ports
        if pl is not None:
            for p in pl:
                if type(p).__name__ == "Token":
                    continue
                try:
                    d = getattr(p, "declarator", None)
                    if d is None or d.name is None:
                        # explicit port like .name(expr) or non-ansi header decl — skip here
                        continue
                    pname = _txt(d.name)
                    h = p.header
                    dtok = getattr(h, "direction", None)
                    dtext = (_txt(dtok).capitalize() if dtok is not None else "In")
                    info["ports"].append((dtext, pname))
                    self._record_loc(name, pname, d.name)
                except Exception:
                    continue

        # ---- members ----
        for m in mod.members:
            k = m.kind
            if k in (sl.SyntaxKind.DataDeclaration, sl.SyntaxKind.NetDeclaration):
                for d in m.declarators:
                    if type(d).__name__ == "Token":
                        continue
                    dn = getattr(d, "name", None)
                    if dn is not None:
                        n = _txt(dn)
                        info["nets"].append(n)
                        self._record_loc(name, n, dn)
            elif k == sl.SyntaxKind.ContinuousAssign:
                for a in m.assignments:
                    if type(a).__name__ == "Token":
                        continue
                    lhs_ids = _lhs_names(a.left) if getattr(a, "left", None) is not None else []
                    rhs_ids = ids_from(a.right) if a.right is not None else []
                    if lhs_ids:
                        info["assigns"].append((lhs_ids[0], rhs_ids))
                        self._record_loc(name, lhs_ids[0], m.getFirstToken())
            elif k in (sl.SyntaxKind.AlwaysBlock, sl.SyntaxKind.AlwaysFFBlock,
                       sl.SyntaxKind.AlwaysCombBlock, sl.SyntaxKind.AlwaysLatchBlock):
                self._scan_always(name, m, info)
            elif k == sl.SyntaxKind.HierarchyInstantiation:
                self._scan_inst(name, m, info)
        info["nets"] = sorted(set(info["nets"]))
        return info

    def _scan_always(self, mod_name, m, info):
        if getattr(m, "statement", None) is None:
            return
        def walk_stmt(s, depth=0):
            if s is None or depth > 40:
                return
            k = s.kind
            if k == sl.SyntaxKind.ExpressionStatement:
                e = s.expr
                ek = str(e.kind)
                if any(a in ek for a in ASSIGN_KINDS):
                    lt = _syntax_of(e.left) if e.left is not None else None
                    rt = _syntax_of(e.right) if e.right is not None else None
                    lhs = _lhs_names(lt) if lt is not None else []
                    rhs = ids_from(rt) if rt is not None else []
                    if lhs:
                        info["always"].append((lhs[0], rhs))
                        self._record_loc(mod_name, lhs[0], s.getFirstToken())
                return
            # descend through timing/if/for/case wrappers
            inner = getattr(s, "statement", None)
            if inner is not None and hasattr(inner, "kind"):
                walk_stmt(inner, depth + 1)
            if k == sl.SyntaxKind.ConditionalStatement:
                cons = getattr(s, "statement", None)
                alt = getattr(s, "elseStatement", None)
                for part in (cons, alt):
                    if part is not None and hasattr(part, "kind"):
                        walk_stmt(part, depth + 1)
                return
            if hasattr(s, "__iter__"):
                for c in s:
                    tn = type(c).__name__
                    if hasattr(c, "__iter__") and ("Statement" in tn or "Case" in tn
                                                   or "Generate" in tn or "Block" in tn):
                        walk_stmt(c, depth + 1)
        walk_stmt(m.statement)

    def _scan_inst(self, mod_name, m, info):
        tnode = m.type
        itype = _txt(tnode.valueText if hasattr(tnode, "valueText") else tnode)
        itype = itype.split()[-1] if itype.split() else "?"
        for inst in m.instances:
            conns = []
            try:
                for pc in inst.connections:
                    if type(pc).__name__ == "Token":
                        continue
                    nm = getattr(pc, "name", None)
                    if nm is None:
                        continue
                    pname = _txt(nm)
                    expr_ids = ids_from(pc.expr) if getattr(pc, "expr", None) is not None else []
                    conns.append((pname, expr_ids))
            except Exception:
                pass
            iname = _txt(inst.decl.name) if inst.decl is not None and inst.decl.name else "?"
            info["insts"].append((itype, iname, conns))

    # ---------- finalize ----------
    def finalize(self):
        self.edges = []
        for mod in self.db.values():
            for (itype, iname, conns) in mod["insts"]:
                child = self.db.get(itype)
                if not child:
                    continue
                cports = {pn: d for d, pn in child["ports"]}
                for (pname, pids) in conns:
                    cdir = cports.get(pname)
                    if cdir is None or not pids:
                        continue
                    for pid in pids:
                        self.edges.append((mod["name"], pid, itype, pname, cdir, iname))
        # adjacency for speed
        self.in_by_child = {}
        for e in self.edges:
            if e[4].lower().startswith("in"):
                self.in_by_child.setdefault((e[2], e[3]), []).append(e)
        self.out_by_child = {}
        for e in self.edges:
            if e[4].lower().startswith("out"):
                self.out_by_child.setdefault((e[2], e[3]), []).append(e)

    # ---------- queries ----------
    def drivers(self, module, sig, max_depth=4, max_nodes=60):
        seen, out, rowseen = set(), [], set()
        def rec(mod_name, s, depth):
            if len(out) >= max_nodes or depth > max_depth or (mod_name, s) in seen:
                return
            seen.add((mod_name, s))
            m = self.db.get(mod_name)
            if not m:
                return
            for lhs, rhs in m["assigns"]:
                if lhs == s:
                    loc = self.loc.get((mod_name, lhs), ("?", "?"))
                    out.append(dict(kind="assign", where=f"{mod_name}  ({loc[0]}:{loc[1]})",
                                    sig=lhs, rhs=rhs[:8]))
                    for r in rhs[:3]:
                        rec(mod_name, r, depth + 1)
            for lhs, rhs in m["always"]:
                if lhs == s:
                    loc = self.loc.get((mod_name, lhs), ("?", "?"))
                    out.append(dict(kind="always", where=f"{mod_name}  ({loc[0]}:{loc[1]})",
                                    sig=lhs, rhs=rhs[:8]))
                    for r in rhs[:3]:
                        rec(mod_name, r, depth + 1)
            pdirs = {pn: d for d, pn in m["ports"]}
            if str(pdirs.get(s, "")).lower().startswith("in"):
                for (pm, pid, cm, cp, cdir, iname) in self.in_by_child.get((mod_name, s), []):
                    out.append(dict(kind="port-in", where=f"from {pm}.{iname}",
                                    sig=f".{cp}", rhs=[pid]))
                    rec(pm, pid, depth + 1)
            for (itype, iname, conns) in m["insts"]:
                for (pname, pids) in conns:
                    if s in pids:
                        child = self.db.get(itype)
                        if child:
                            cdirs = {pn: d for d, pn in child["ports"]}
                            if str(cdirs.get(pname, "")).lower().startswith("out"):
                                out.append(dict(kind="inst-out",
                                                where=f"{mod_name}.{iname} ({itype}.{pname})",
                                                sig=s, rhs=[f"{itype}.{pname}"]))
                                rec(itype, pname, depth + 1)
        rec(module, sig, 0)
        return [r for r in out
                if not (key := (r['kind'], r['where'], r['sig'], tuple(r['rhs']))) in rowseen and not rowseen.add(key)]

    def loads(self, module, sig, max_depth=3, max_nodes=60):
        """who reads sig (same module) and where does it flow (output port / down into insts)"""
        seen, out, rowseen = set(), [], set()
        def rec(mod_name, s, depth):
            if len(out) >= max_nodes or depth > max_depth or (mod_name, s) in seen:
                return
            seen.add((mod_name, s))
            m = self.db.get(mod_name)
            if not m:
                return
            for lhs, rhs in m["assigns"] + m["always"]:
                if s in rhs:
                    loc = self.loc.get((mod_name, lhs), ("?", "?"))
                    out.append(dict(kind="load", where=f"{mod_name}  ({loc[0]}:{loc[1]})",
                                    sig=lhs, rhs=[s]))
                    rec(mod_name, lhs, depth + 1)
            pdirs = {pn: d for d, pn in m["ports"]}
            if str(pdirs.get(s, "")).lower().startswith("out"):
                for (pm, pid, cm, cp, cdir, iname) in self.out_by_child.get((mod_name, s), []):
                    out.append(dict(kind="port-out", where=f"into {pm} (as {pid}, via {iname})",
                                    sig=f".{cp}", rhs=[pid]))
                    rec(pm, pid, depth + 1)
            for (itype, iname, conns) in m["insts"]:
                for (pname, pids) in conns:
                    if s in pids:
                        child = self.db.get(itype)
                        cdirs = {pn: d for d, pn in child["ports"]} if child else {}
                        kind = cdirs.get(pname, "?")
                        out.append(dict(kind=f"inst-load-{kind.lower() or 'unk'}",
                                        where=f"{mod_name}.{iname} -> {itype}.{pname}",
                                        sig=s, rhs=[pname]))
        rec(module, sig, 0)
        return [r for r in out
                if not (key := (r['kind'], r['where'], r['sig'], tuple(r['rhs']))) in rowseen and not rowseen.add(key)]

    def module_summary(self, name):
        m = self.db.get(name)
        if not m:
            return None
        return dict(name=m["name"],
                    ports=[f"{d} {p}" for d, p in m["ports"]],
                    nets=m["nets"][:120],
                    n_assign=len(m["assigns"]), n_always=len(m["always"]),
                    insts=[dict(type=t, name=i, ports=[f".{p}({','.join(ids[:3])})" for p, ids in c][:16])
                           for t, i, c in m["insts"]])

    def search(self, pat, limit=20):
        rx = re.compile(pat, re.I)
        hits = []
        for (mod, sig), (f, l) in self.loc.items():
            if rx.search(sig) and len(hits) < limit:
                hits.append(dict(module=mod, signal=sig, file=f, line=l))
        return hits

    def save(self, path):
        with open(path, "w", encoding="utf-8") as f:
            json.dump(dict(modules=list(self.db), n_edges=len(self.edges)), f)

    @classmethod
    def build(cls, roots):
        idx = cls()
        files = []
        for r in roots:
            if os.path.isfile(r):
                files.append(r)
            else:
                files += glob.glob(os.path.join(r, "**", "*.v"), recursive=True)
                files += glob.glob(os.path.join(r, "**", "*.sv"), recursive=True)
        idx._roots = [os.path.abspath(r) for r in roots]
        for f in sorted(set(files)):
            try:
                idx.add_tree(f)
            except Exception:
                pass
        idx.finalize()
        return idx