#!/usr/bin/env python3
"""Replace fully-constant wire-array ROMs with Vivado dist_mem_gen instances.

- Copies every non-tb .v from <hdlsrc> into <hdlsrc>/rtl_ip/ (first run only, or --force).
- For each fully-constant ROM in the manifest:
    * extracts the table into a .coe (dedup by content hash: identical tables share one IP config)
    * deletes the wire-array declaration and the constant assigns from the copy
    * per READ SITE inserts:  wire [AW-1:0] <name>_a<k>; wire [W-1:0] <name>_q<k>;
                              assign <name>_a<k> = <read_expr>;
                              <ip> u_<name>_r<k> (.a(<name>_a<k>), .spo(<name>_q<k>));
      and rewrites `<name>[<read_expr>]` -> `$signed(<name>_q<k>)` (or unsigned).
- Appends dist_mem_gen entries to ip_manifest.json / ip_manifest.tcl.

dist_mem_gen is ASYNC read => bit- and cycle-true against the original combinational read.
Run AFTER rtl_ip_scan.py, BEFORE namespace_ip.py.
"""
import argparse, hashlib, json, os, re, shutil, sys

def to_bits(radix, lit, width):
    lit = lit.replace('_', '')
    base = {'b': 2, 'd': 10, 'h': 16}[radix]
    v = int(lit, base) & ((1 << width) - 1)
    return format(v, '0%db' % width)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('hdlsrc')
    ap.add_argument('outdir', help='FIXED output dir (e.g. <repo>/vivado/rtl_ip/<model>)')
    ap.add_argument('--manifest', default=None)
    ap.add_argument('--force', action='store_true')
    a = ap.parse_args()
    from ipflow_common import copy_sources
    hdlsrc = os.path.abspath(a.hdlsrc)
    rtl_ip = os.path.abspath(a.outdir)
    os.makedirs(rtl_ip, exist_ok=True)
    man = json.load(open(a.manifest or os.path.join(rtl_ip, 'manifest.json')))

    # 1. copy sources + PROVENANCE.log
    copy_sources(hdlsrc, rtl_ip, force=a.force)

    ipman_p = os.path.join(rtl_ip, 'ip_manifest.json')
    ipman = json.load(open(ipman_p)) if os.path.exists(ipman_p) else {'ips': []}
    coe_of_hash, roms_done = {}, 0

    CONST_ASSIGN = re.compile(r"^\s*assign\s+(\w+)\[(\d+)\]\s*=\s*(\d+)'s?([bdh])([0-9a-fA-F_]+)\s*;\s*$", re.M)

    by_file = {}
    for rom in man['const_rom']:
        if not rom.get('fully_const'):
            print('SKIP (partial const):', rom['name']); continue
        by_file.setdefault(rom['file'], []).append(rom)

    for fn, roms in by_file.items():
        path = os.path.join(rtl_ip, fn)
        txt = open(path, encoding='utf-8', errors='replace').read()
        for rom in roms:
            name, W, D, signed = rom['name'], rom['width'], rom['depth'], rom['signed']
            AW = max(1, (D - 1).bit_length())
            # collect table
            bits = {}
            for m in CONST_ASSIGN.finditer(txt):
                if m.group(1) == name:
                    bits[int(m.group(2))] = to_bits(m.group(4), m.group(5), W)
            if len(bits) != D:
                print(f'ERROR {name}: found {len(bits)}/{D} const assigns -- skipped'); continue
            content = '\n'.join(bits[i] for i in range(D))
            h = hashlib.sha1((f'{W}x{D}\n' + content).encode()).hexdigest()[:10]
            if h not in coe_of_hash:
                ipname = f'rom_{W}x{D}_{h}'
                coe = os.path.join(rtl_ip, ipname + '.coe')
                with open(coe, 'w') as f:
                    f.write('memory_initialization_radix=2;\nmemory_initialization_vector=\n')
                    f.write(',\n'.join(bits[i] for i in range(D)) + ';\n')
                coe_of_hash[h] = ipname
                ipman['ips'].append({'type': 'dist_mem_gen', 'name': ipname, 'width': W,
                                     'depth': D, 'coe': os.path.abspath(coe)})
            ipname = coe_of_hash[h]

            # remove decl + assigns
            txt = re.sub(r'^\s*wire\s+(signed\s+)?\[%d:0\]\s+%s\s+\[0:%d\]\s*;.*\n'
                         % (W - 1, re.escape(name), D - 1), '', txt, flags=re.M)
            txt = re.sub(r"^\s*assign\s+%s\[\d+\]\s*=\s*[^;]+;\s*\n" % re.escape(name), '', txt, flags=re.M)

            # rewrite read sites (unique exprs; every occurrence of same expr shares one instance)
            inst_lines, k = [], 0
            for expr in dict.fromkeys(rom['read_exprs']):     # preserve order, dedup
                aq, qq = f'{name}_a{k}', f'{name}_q{k}'
                inst_lines.append(f'  wire [{AW-1}:0] {aq};\n  wire [{W-1}:0] {qq};\n'
                                  f'  assign {aq} = {expr};\n'
                                  f'  {ipname} u_{name}_r{k} (.a({aq}), .spo({qq}));\n')
                rd = f'$signed({qq})' if signed else qq
                txt = txt.replace(f'{name}[{expr}]', rd)
                k += 1
            # place instances right before `endmodule`
            txt = txt.replace('endmodule', '  // ---- rtl-ip-replace: %s -> %s (%d read site(s)) ----\n%s\nendmodule'
                              % (name, ipname, k, ''.join(inst_lines)), 1)
            roms_done += 1
            print(f'ROM {name} ({W}x{D}, {k} reads) -> {ipname}  [{fn}]')
        open(path, 'w', encoding='utf-8').write(txt)

    json.dump(ipman, open(ipman_p, 'w'), indent=1)
    with open(os.path.join(rtl_ip, 'ip_manifest.tcl'), 'w') as f:
        for ip in ipman['ips']:
            f.write('lappend ips {%s}\n' % ' '.join(f'{k} {{{v}}}' for k, v in ip.items()))
    print(f'done: {roms_done} ROMs, {len(coe_of_hash)} unique tables; ip_manifest updated')

if __name__ == '__main__':
    main()