#!/usr/bin/env python3
"""Scan HDL Coder generated Verilog for IP-replaceable structures.

Emits a JSON manifest of:
  const_rom : `wire [W:0] name [0:D];` fully populated by constant assigns -> dist/blk ROM IP
  reg_table : `reg` arrays (runtime-written tables) -> review by hand (RAM or leave)
  ram       : SimpleDualPortRAM_generic instantiations (AddrWidth/DataWidth combos) -> blk_mem_gen
  fifo      : modules/instances whose name contains 'fifo' -> fifo_generator (guided)
  cordic    : files with the cx/cy/cz + atanT iterative kernel -> cordic IP (guided)
  filter    : FIR filters (Task12 §六) -> fir_compiler / cic_compiler (guided)
                fir_coef_bank        : HDL Coder <Name>Coef module, constant CoefOut_i taps
                fir_block            : module whose header says "Discrete FIR Filter"
                adaptive_fir_candidate: hand-rolled tap regs + delay line + MAC (LMS etc.)

Usage: python rtl_ip_scan.py <hdlsrc_dir> [<outdir>] [-o manifest.json]
"""
import argparse, json, os, re, sys

ROM_DECL = re.compile(r'^\s*wire\s+(signed\s+)?\[(\d+):0\]\s+(\w+)\s+\[0:(\d+)\]\s*;', re.M)
REG_DECL = re.compile(r'^\s*reg\s+(signed\s+)?\[(\d+):0\]\s+(\w+)\s+\[0:(\d+)\]\s*;', re.M)
CONST_ASSIGN = re.compile(r"^\s*assign\s+(\w+)\[(\d+)\]\s*=\s*(\d+)'s?([bdh])([0-9a-fA-F_xzXZ]+)\s*;", re.M)
RAM_INST = re.compile(r'SimpleDualPortRAM_\w+\s*#\s*\(((?:[^()]|\([^()]*\))*)\)\s*(\w+)\s*\(', re.S)
PARAM = re.compile(r'\.(\w+)\s*\(\s*(\d+)\s*\)')

# ---- filter (Task12 §六) ----------------------------------------------------
MODULE = re.compile(r'^[ \t]*module\s+(\w+)\b(.*?)^\s*endmodule', re.M | re.S)
DECL = re.compile(r'^\s*(input|output|inout|wire|reg)\s+(signed\s+)?(?:\[(\d+):0\]\s+)?(\w+)\s*;'
                  r'(?:[ \t]*//[ \t]*([usq]fix\d+(?:_En-?\d+)?|\w+))?', re.M)
CONST_SIG = re.compile(r"^\s*assign\s+(\w+)\s*=\s*(\d+)'s?([bdh])([0-9a-fA-F_]+)\s*;", re.M)
BASE_IDX = re.compile(r'^(.*?)_?(\d+)$')
DLY_CHAIN = re.compile(r'^\s*(\w*?)(\d+)_next(_re|_im)?\s*=\s*\1(\d+)(?:_re|_im)?\s*;', re.M)
FIXPT = re.compile(r'([us])fix(\d+)(?:_En(-?\d+))?')


def to_signed(width, radix, digits):
    v = int(digits.replace('_', ''), {'b': 2, 'd': 10, 'h': 16}[radix])
    return v - (1 << width) if v >> (width - 1) else v


def split_modules(txt):
    """[(name, body, header_comment_block)] -- HDL Coder writes the Simulink block
    kind ('Discrete FIR Filter') and Source Path in the comment above the module."""
    out = []
    for m in MODULE.finditer(txt):
        out.append((m.group(1), m.group(2), txt[max(0, m.start() - 1200):m.start()]))
    return out


def decls_of(body):
    d = {}
    for m in DECL.finditer(body):
        d[m.group(4)] = {'dir': m.group(1), 'signed': m.group(2) is not None,
                         'width': int(m.group(3)) + 1 if m.group(3) else 1,
                         'type': m.group(5) or ''}
    return d


def fixpt_of(type_comment, width):
    m = FIXPT.search(type_comment or '')
    if not m:
        return {'width': width, 'frac': None, 'signed': None}
    return {'width': int(m.group(2)), 'frac': int(m.group(3) or 0), 'signed': m.group(1) == 's'}


def const_groups(body, decls):
    """constant `assign base_<i> = W'sh...;` families, index-complete from 0."""
    vals = {}
    for m in CONST_SIG.finditer(body):
        name, w, radix, digits = m.group(1), int(m.group(2)), m.group(3), m.group(4)
        vals[name] = (w, to_signed(w, radix, digits))
    groups = {}
    for name, (w, v) in vals.items():
        bm = BASE_IDX.match(name)
        if bm:
            groups.setdefault((bm.group(1), w), {})[int(bm.group(2))] = v
    # HDL Coder names the first element without an index (CoefData, CoefData_1, ...)
    for (base, w), members in groups.items():
        if 0 not in members and base in vals and vals[base][0] == w:
            members[0] = vals[base][1]
    out = []
    for (base, w), members in groups.items():
        n = len(members)
        if n < 4 or sorted(members) != list(range(n)):
            continue
        coeffs = [members[i] for i in range(n)]
        ex = decls.get(base + '_0') or decls.get(base + '0') or decls.get(base) or {}
        out.append({'base': base, 'n': n, 'width': w, 'coeffs': coeffs,
                    'frac': fixpt_of(ex.get('type', ''), w)['frac'],
                    'symmetric': coeffs == coeffs[::-1]})
    return out


def scan_filters(path, txt):
    """fir_coef_bank / fir_block / adaptive_fir_candidate (guided replacement)."""
    found = []
    fn = os.path.basename(path)
    for name, body, hdr in split_modules(txt):
        decls = decls_of(body)
        is_fir_hdr = 'Discrete FIR Filter' in hdr
        srcpath = (re.search(r'//\s*Source Path:\s*(\S+)', hdr) or [None, ''])[1] \
            if re.search(r'//\s*Source Path:\s*(\S+)', hdr) else ''
        ports = {k: v for k, v in decls.items() if v['dir'] in ('input', 'output')}
        coef_in = [p for p in ports if re.match(r'coef', p, re.I) and ports[p]['dir'] == 'input']

        for g in const_groups(body, decls):
            # A constant family is only a TAP BANK if it is named/exported as one, or
            # if it is fed to multipliers inside a FIR module.  Without this guard every
            # HDL Coder state/parameter constant family (state_1, mu_2, ...) looks like
            # a filter -- 5 false positives on RxAlignCore alone.
            exports = [p for p in ports if ports[p]['dir'] == 'output'
                       and re.search(r'coef|tap', p, re.I)]
            named = re.search(r'coef|tap', name, re.I) or re.search(r'coef|tap', g['base'], re.I)
            multiplied = any(re.search(r'\b%s(?:_\d+)?\s*\*|\*\s*%s(?:_\d+)?\b'
                                       % (re.escape(g['base']), re.escape(g['base'])), body)
                             for _ in (0,)) and (is_fir_hdr or re.search(r'fir|filt', name, re.I))
            if not (exports or named or multiplied):
                continue
            found.append({'file': fn, 'module': name, 'kind': 'fir_coef_bank',
                          'source_path': srcpath, 'coef_signal': g['base'],
                          'taps': g['n'], 'coef_width': g['width'], 'coef_frac': g['frac'],
                          'symmetric': g['symmetric'], 'coeffs': g['coeffs'],
                          'exports_taps': bool(exports),
                          'note': 'fixed coefficients -> fir_compiler (.coe), guided (Task12 §6.2)'})

        din = [p for p in ports if re.match(r'(data|d)?in', p, re.I) and ports[p]['dir'] == 'input'
               and ports[p]['width'] > 1]
        dout = [p for p in ports if re.search(r'out', p, re.I) and ports[p]['dir'] == 'output'
                and ports[p]['width'] > 1]
        if is_fir_hdr or (re.search(r'fir|filt', name, re.I) and din and dout):
            e = {'file': fn, 'module': name,
                 'kind': 'fir_block' if is_fir_hdr else 'fir_block_candidate',
                 'source_path': srcpath,
                 'complex': any(p.endswith('_re') for p in din),
                 'data_in': {p: fixpt_of(ports[p]['type'], ports[p]['width']) for p in sorted(din)},
                 'data_out': {p: fixpt_of(ports[p]['type'], ports[p]['width']) for p in sorted(dout)},
                 'valid_in': any(re.match(r'valid|.*vld', p, re.I) and ports[p]['dir'] == 'input'
                                 for p in ports),
                 'valid_out': any(re.match(r'valid|.*vld', p, re.I) and ports[p]['dir'] == 'output'
                                  for p in ports),
                 'coef_ports': sorted(coef_in),
                 'adaptive': bool(coef_in),
                 'rate_hint': ('interpolation' if re.search(r'interp|up', name, re.I) else
                               'decimation' if re.search(r'decim|down', name, re.I) else 'single_rate?'),
                 'note': 'rate/phase must be confirmed against the Simulink block before IP config'}
            found.append(e)

        # hand-rolled adaptive FIR: tap regs + delay-line chain + multiplies
        chains = {}
        for m in DLY_CHAIN.finditer(body):
            if int(m.group(2)) == int(m.group(4)) + 1:
                chains.setdefault(m.group(1), set()).add(int(m.group(2)))
        chain = max(((b, len(s)) for b, s in chains.items()), key=lambda x: x[1], default=None)
        # tap registers: indexed reg group named like weights/taps.  Do NOT require the
        # reg itself to sit next to a `*` -- HDL Coder feeds multipliers through operand
        # temporaries (da0_re * db0_re), so that test misses every real case (EqCore).
        tapregs = {}
        for rn, d in decls.items():
            bm = BASE_IDX.match(re.sub(r'_(re|im)$', '', rn))
            if d['dir'] == 'reg' and d['width'] > 1 and bm \
                    and re.fullmatch(r'w\w*|tap\w*|coef\w*|h', bm.group(1), re.I):
                tapregs.setdefault(bm.group(1), set()).add(int(bm.group(2)))
        taps = max(((b, len(s)) for b, s in tapregs.items()), key=lambda x: x[1], default=None)
        if chain and chain[1] >= 3 and taps and taps[1] >= 3 and '*' in body:
            found.append({'file': fn, 'module': name, 'kind': 'adaptive_fir_candidate',
                          'source_path': srcpath, 'tap_regs': taps[0], 'taps': taps[1],
                          'delay_line': chain[0], 'delay_depth': chain[1] + 1,
                          'note': 'runtime-updated taps: fir_compiler Reloadable, or KEEP the RTL '
                                  'with a written reason (Task12 §6.1)'})
    return found

KINDS = ('const_rom', 'reg_table', 'ram', 'fifo', 'cordic', 'filter')


def scan_file(path):
    txt = open(path, encoding='utf-8', errors='replace').read()
    out = {k: [] for k in KINDS}

    # constant assigns grouped by array name
    assigns = {}
    for m in CONST_ASSIGN.finditer(txt):
        assigns.setdefault(m.group(1), {})[int(m.group(2))] = (m.group(3), m.group(4), m.group(5))

    for m in ROM_DECL.finditer(txt):
        signed, msb, name, dmax = m.group(1) is not None, int(m.group(2)), m.group(3), int(m.group(4))
        depth = dmax + 1
        vals = assigns.get(name, {})
        # read sites: name[expr] not on its own assign-constant lines
        reads = []
        for rm in re.finditer(re.escape(name) + r'\[([^\]\n]+)\]', txt):
            idx = rm.group(1)
            if idx.isdigit():          # constant index = the table-fill assigns
                continue
            reads.append(idx.strip())
        entry = {'file': os.path.basename(path), 'name': name, 'width': msb + 1,
                 'depth': depth, 'signed': signed, 'n_const': len(vals),
                 'fully_const': len(vals) == depth, 'read_exprs': reads}
        if len(vals) == depth and reads:
            out['const_rom'].append(entry)
        elif len(vals) > 0:
            entry['note'] = 'PARTIAL const fill -- inspect by hand'
            out['const_rom'].append(entry)

    for m in REG_DECL.finditer(txt):
        out['reg_table'].append({'file': os.path.basename(path), 'name': m.group(3),
                                 'width': int(m.group(2)) + 1, 'depth': int(m.group(4)) + 1,
                                 'note': 'runtime-written reg array: treat as RAM or leave as-is'})

    for m in RAM_INST.finditer(txt):
        params = dict((k, int(v)) for k, v in PARAM.findall(m.group(1)))
        out['ram'].append({'file': os.path.basename(path), 'inst': m.group(2),
                           'AddrWidth': params.get('AddrWidth'), 'DataWidth': params.get('DataWidth')})

    for m in re.finditer(r'^\s*module\s+(\w*[fF][iI][fF][oO]\w*)', txt, re.M):
        out['fifo'].append({'file': os.path.basename(path), 'module': m.group(1)})

    if re.search(r'\bcx\b', txt) and re.search(r'\bcz\b', txt) and 'atanT' in txt:
        out['cordic'].append({'file': os.path.basename(path),
                              'note': 'iterative CORDIC kernel (cx/cy/cz + atanT) -- guided replacement'})

    out['filter'] = scan_filters(path, txt)
    return out

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('hdlsrc')
    ap.add_argument('outdir', nargs='?', default=None,
                    help='FIXED output dir; manifest goes to <outdir>/manifest.json')
    ap.add_argument('-o', '--out', default=None)
    a = ap.parse_args()
    if a.outdir and not a.out:
        a.out = os.path.join(a.outdir, 'manifest.json')
    manifest = dict({'hdlsrc': os.path.abspath(a.hdlsrc), 'files': {}}, **{k: [] for k in KINDS})
    for fn in sorted(os.listdir(a.hdlsrc)):
        if not fn.endswith('.v') or fn.endswith('_tb.v'):
            continue
        r = scan_file(os.path.join(a.hdlsrc, fn))
        for k in KINDS:
            manifest[k].extend(r[k])

    # a Discrete FIR Filter block expands into a tree of helper modules (Filter,
    # subFilter, FilterTapSystolic*): keep the block, drop its own guts.
    tops = [f['source_path'] for f in manifest['filter']
            if f['kind'] == 'fir_block' and f.get('source_path')]
    kept = []
    for f in manifest['filter']:
        inner = any(f.get('source_path', '').startswith(t + '/') for t in tops)
        if inner and f['kind'] == 'fir_block_candidate':
            continue                      # internal plumbing of a detected FIR block
        if inner and f['kind'] == 'fir_coef_bank':
            f['belongs_to'] = next(t for t in tops if f['source_path'].startswith(t + '/'))
        kept.append(f)
    manifest['filter'] = kept
    outp = a.out or os.path.join(a.hdlsrc, 'rtl_ip', 'manifest.json')
    os.makedirs(os.path.dirname(outp), exist_ok=True)
    json.dump(manifest, open(outp, 'w'), indent=1)
    print('manifest ->', outp)
    for k in KINDS:
        print(f'  {k}: {len(manifest[k])}')
        for e in manifest[k]:
            brief = {kk: (f'[{len(vv)} values]' if kk == 'coeffs' else vv)
                     for kk, vv in e.items() if kk != 'read_exprs'}
            print('    ', brief,
                  f"reads={len(e.get('read_exprs', []))}" if 'read_exprs' in e else '')

if __name__ == '__main__':
    main()