#!/usr/bin/env python3
"""gen_cm_dmem_direct.py
Bypass the EFX_MAP 2026.1 $readmemb→VDB bug by replacing the inferred
`dmem` array in church_ti60_f225.v with a module that explicitly
instantiates EFX_RAM10 primitives with inline INIT values.

EFX_MAP passes explicit primitive #(.INIT_N(...)) parameters straight into
the VDB binary — no $readmemb inference, no bug.

Confirmed INIT mapping (READ_WIDTH=1, WRITE_WIDTH=1):
  For bit-plane P (= lane*8 + bit, 0..31), address range R (0 or 1),
  INIT param N (0..3), address offset B (0..255):
    INIT_N bit B = (dmem[R*1024 + N*256 + B] >> P) & 1

Usage:
    python3 gen_cm_dmem_direct.py <SoC_dir>

  SoC_dir — directory containing church_ti60_f225.v (and where
             cm_dmem_bram.v will be written).
             Defaults to the directory containing this script's parent
             (i.e. the SoC root when run from work_syn/).

Outputs:
  <SoC_dir>/cm_dmem_bram.v       — EFX_RAM10 instantiation module
  <SoC_dir>/church_ti60_f225.v   — patched top-level (backup created)

Then re-run synthesis + P&R:
    bash work_syn/run_efx_map.sh
    bash work_pnr/run_efx_pnr.sh
    sudo openFPGALoader -b titanium_ti60_f225_jtag \\
         -f ~/church_project/SoC/outflow/church_soc_cm.hex
"""

import re
import sys
import os
import shutil

# ── EFX_RAM10 instance template parameters ──────────────────────────────────

EFX_PARAMS = """\
        .READ_WIDTH(1), .WRITE_WIDTH(1),
        .WCLK_POLARITY(1'b1), .WCLKE_POLARITY(1'b1),
        .WE_POLARITY(2'b11), .WADDREN_POLARITY(1'b1),
        .RADDREN_POLARITY(1'b1), .RST_POLARITY(1'b1),
        .RCLK_POLARITY(1'b1), .RE_POLARITY(1'b1),
        .OUTPUT_REG(1'b0),
        .WRITE_MODE("READ_FIRST"),
        .RESET_RAM("ASYNC"), .RESET_OUTREG("ASYNC")"""

# ── Patterns for patching church_ti60_f225.v ────────────────────────────────

DECL_PAT = re.compile(
    r'(?:\(\* src[^*]*\*\)\n)?'
    r'  reg \[31:0\] dmem \[(\d+):0\];\n',
    re.MULTILINE,
)

INIT_PAT = re.compile(
    r'  initial begin\n'
    r'((?:    dmem\[\d+\] = 32\'d\d+;\n)*)'
    r'  end\n',
    re.MULTILINE,
)

WRITE_PAT = re.compile(
    r'  always @\(posedge clk\) begin\n'
    r'    if \(dmem_wr__en\)\n'
    r'      dmem\[dmem_wr__addr\] <= dmem_wr__data;\n'
    r'  end\n',
    re.MULTILINE,
)

READ_PAT = re.compile(
    r'  reg \[31:0\] (_\d+_);\n'
    r'  always @\(posedge clk\) begin\n'
    r'    (_\d+_) <= dmem\[mem_addr\];\n'
    r'  end\n'
    r'  assign mem_rd_data = (_\d+_);\n',
    re.MULTILINE,
)

# $readmemb form: declarations (separate from write/read blocks below).
# Use [ ]+ (flexible indentation) and an optional (* src ... *) attribute
# line — the Chromebook-generated Verilog can have 2- or 4-space indent.
DECL_BRAM_PAT = re.compile(
    r'(?:[ ]+\(\*[^*]*\*\)\n)?'           # optional (* src ... *) attribute
    r'[ ]+reg \[7:0\] dmem_b0 \[0:\d+\];\n'
    r'[ ]+reg \[7:0\] dmem_b1 \[0:\d+\];\n'
    r'[ ]+reg \[7:0\] dmem_b2 \[0:\d+\];\n'
    r'[ ]+reg \[7:0\] dmem_b3 \[0:\d+\];\n'
    r'[ ]+initial \$readmemb\("[^"]*", dmem_b0\);\n'
    r'[ ]+initial \$readmemb\("[^"]*", dmem_b1\);\n'
    r'[ ]+initial \$readmemb\("[^"]*", dmem_b2\);\n'
    r'[ ]+initial \$readmemb\("[^"]*", dmem_b3\);\n',
    re.MULTILINE,
)

# $readmemb form: write always block (byte-lane version)
WRITE_BRAM_PAT = re.compile(
    r'  always @\(posedge clk\) begin\n'
    r'    if \(dmem_wr__en\) begin\n'
    r'      dmem_b0\[dmem_wr__addr\] <= dmem_wr__data\[7:0\];\n'
    r'      dmem_b1\[dmem_wr__addr\] <= dmem_wr__data\[15:8\];\n'
    r'      dmem_b2\[dmem_wr__addr\] <= dmem_wr__data\[23:16\];\n'
    r'      dmem_b3\[dmem_wr__addr\] <= dmem_wr__data\[31:24\];\n'
    r'    end\n'
    r'  end\n',
    re.MULTILINE,
)

# $readmemb form: read always block (byte-lane version)
READ_BRAM_PAT = re.compile(
    r'  reg \[31:0\] _0_;\n'
    r'  always @\(posedge clk\) begin\n'
    r'    _0_ <= \{dmem_b3\[mem_addr\], dmem_b2\[mem_addr\],\n'
    r'            dmem_b1\[mem_addr\], dmem_b0\[mem_addr\]\};\n'
    r'  end\n'
    r'  assign mem_rd_data = _0_;\n',
    re.MULTILINE,
)


# ── INIT computation ─────────────────────────────────────────────────────────

def compute_init(dmem_data: dict, plane: int, range_idx: int, param: int) -> int:
    """Return 256-bit integer for INIT_{param} of EFX_RAM10 instance.

    READ_WIDTH=1 mapping:
      INIT_{param} bit B = (dmem[range_idx*1024 + param*256 + B] >> plane) & 1
    """
    base = range_idx * 1024 + param * 256
    result = 0
    for b in range(256):
        addr = base + b
        word = dmem_data.get(addr, 0)
        if (word >> plane) & 1:
            result |= (1 << b)
    return result


# ── Verilog generator ─────────────────────────────────────────────────────────

def gen_cm_dmem_bram(dmem_data: dict, depth: int) -> str:
    """Generate complete cm_dmem_bram.v content."""

    lines = []
    a = lines.append

    a('// cm_dmem_bram.v — AUTO-GENERATED by gen_cm_dmem_direct.py')
    a('// Direct EFX_RAM10 instantiation with inline INIT values.')
    a('// Bypasses EFX_MAP 2026.1 $readmemb->VDB INIT encoding bug.')
    a('//')
    a('// Addresses 0-2047   : 64 explicit EFX_RAM10 (READ_WIDTH=1)')
    a('// Addresses 2048+    : reads return 32\'h0 (no FF array — saves 458K FFs)')
    a('')
    a('module cm_dmem_bram (')
    a('    input  wire        clk,')
    a('    input  wire [13:0] wr_addr,')
    a('    input  wire [31:0] wr_data,')
    a('    input  wire        wr_en,')
    a('    input  wire [13:0] rd_addr,')
    a('    output reg  [31:0] rd_data')
    a(');')
    a('')
    a('// ── Range-select signals ───────────────────────────────────────────────────')
    a('wire wr_in_efx  = wr_en && (wr_addr[13:11] == 3\'b000);')
    a('wire wr_range0  = wr_in_efx && !wr_addr[10];')
    a('wire wr_range1  = wr_in_efx &&  wr_addr[10];')
    a('wire rd_in_efx  = (rd_addr[13:11] == 3\'b000);')
    a('wire rd_range0  = rd_in_efx && !rd_addr[10];')
    a('wire rd_range1  = rd_in_efx &&  rd_addr[10];')
    a('')
    a('// NOTE: addresses 2048+ are write-ignored and read as 0.')
    a('')
    a('// ── EFX_RAM10 read-data wires ──────────────────────────────────────────────')
    for r in range(2):
        for lane in range(4):
            for bit in range(8):
                a(f'wire r{r}_l{lane}b{bit}_rdata;')
    a('')

    # Generate 64 EFX_RAM10 instances
    a('// ── 64 × EFX_RAM10 instances (READ_WIDTH=1) ─────────────────────────────────')

    for r in range(2):
        a(f'// Range {r} (addresses {r*1024}-{r*1024+1023})')
        for lane in range(4):
            for bit in range(8):
                plane = lane * 8 + bit
                iname = f'r{r}_l{lane}b{bit}'

                # Compute INIT_0..INIT_3
                init_vals = []
                for param in range(4):
                    v = compute_init(dmem_data, plane, r, param)
                    init_vals.append(v)

                # Only include non-zero INIT params in the parameter list
                # (always include at least INIT_0 and INIT_1 for clarity)
                init_strs = []
                for n, v in enumerate(init_vals):
                    init_strs.append(f'        .INIT_{n}(256\'h{v:064X})')

                data_bit = lane * 8 + bit
                nz = sum(1 for v in init_vals if v)
                a(f'// Plane {plane} (lane {lane}, bit {bit}), range {r}'
                  f' — {nz} non-zero INIT params')
                a(f'EFX_RAM10 #(')
                a(EFX_PARAMS + ',')
                for s in init_strs:
                    a(s + ',')
                # Remove trailing comma from last INIT
                lines[-1] = lines[-1].rstrip(',')
                a(f') {iname} (')
                a(f'    .WCLK(clk), .RCLK(clk),')
                a(f'    .WADDR(wr_addr[9:0]),')
                a(f'    .RADDR(rd_addr[9:0]),')
                a(f'    .WDATA(wr_data[{data_bit}]),')
                a(f'    .RDATA({iname}_rdata),')
                a(f'    .WE(2\'b11),')
                a(f'    .WCLKE(wr_range{r}),')
                a(f'    .WADDREN(1\'b1),')
                a(f'    .RADDREN(1\'b1),')
                a(f'    .RE(1\'b1), .RST(1\'b0)')
                a(f');')
                a('')

    # Assemble read-data buses
    a('// ── Assemble 32-bit read-data from individual bits ──────────────────────────')
    for r in range(2):
        bits_list = ', '.join(
            f'r{r}_l{lane}b{bit}_rdata'
            for lane in range(3, -1, -1)
            for bit in range(7, -1, -1)
        )
        a(f'wire [31:0] rdata_r{r} = {{{bits_list}}};')
    a('')

    # Registered read mux
    a('// ── Registered read mux ─────────────────────────────────────────────────────')
    a('always @(posedge clk) begin')
    a('    if (rd_range0)')
    a('        rd_data <= rdata_r0;')
    a('    else if (rd_range1)')
    a('        rd_data <= rdata_r1;')
    a('    else')
    a('        rd_data <= 32\'h0;  // addr >= 2048: return zero')
    a('end')
    a('')
    a('endmodule')

    return '\n'.join(lines) + '\n'


# ── church_ti60_f225.v patcher ───────────────────────────────────────────────

NEW_INSTANTIATION = '''\
  // cm_dmem_bram: direct EFX_RAM10 instantiation (bypasses $readmemb->VDB bug)
  wire [31:0] _dmem_rd_data;
  cm_dmem_bram u_dmem (
    .clk(clk),
    .wr_addr(dmem_wr__addr),
    .wr_data(dmem_wr__data),
    .wr_en(dmem_wr__en),
    .rd_addr(mem_addr[13:0]),
    .rd_data(_dmem_rd_data)
  );
  assign mem_rd_data = _dmem_rd_data;
'''


def patch_top(src: str, depth: int) -> str:
    """Patch church_ti60_f225.v to use cm_dmem_bram."""

    # ── Case 1: original form (initial begin) ────────────────────────────────
    dm = DECL_PAT.search(src)
    im = INIT_PAT.search(src)
    wp = WRITE_PAT.search(src)
    rp = READ_PAT.search(src)

    if dm and im and wp and rp:
        print('  Patching: original initial begin form')
        out = src
        out, n1 = DECL_PAT.subn('', out, count=1)
        out, n2 = INIT_PAT.subn('', out, count=1)
        out, n3 = WRITE_PAT.subn(NEW_INSTANTIATION, out, count=1)
        out, n4 = READ_PAT.subn('', out, count=1)
        if n1 == n2 == n3 == n4 == 1:
            return out
        raise RuntimeError(f'Substitution failed: n1={n1} n2={n2} n3={n3} n4={n4}')

    # ── Case 2: already patched with $readmemb form (3 separate regions) ───────
    decl_m  = DECL_BRAM_PAT.search(src)
    write_m = WRITE_BRAM_PAT.search(src)
    read_m  = READ_BRAM_PAT.search(src)

    if decl_m and write_m and read_m:
        print('  Patching: $readmemb byte-lane form (replacing 3 regions)')
        out = src
        out, n1 = DECL_BRAM_PAT.subn('', out, count=1)
        out, n2 = WRITE_BRAM_PAT.subn(NEW_INSTANTIATION, out, count=1)
        out, n3 = READ_BRAM_PAT.subn('', out, count=1)
        if n1 == n2 == n3 == 1:
            return out
        raise RuntimeError(f'BRAM substitution failed: n1={n1} n2={n2} n3={n3}')

    # ── Case 2b: partial $readmemb (decl present but write/read vary) ────────
    if decl_m:
        print('  Patching: $readmemb decl found; write/read blocks may vary.')
        print(f'    DECL={bool(decl_m)} WRITE={bool(write_m)} READ={bool(read_m)}')
        # Show surrounding context to help diagnose
        if not write_m:
            # Search for any dmem_b0 write to report what we see
            ctx = re.search(r'dmem_b0\[dmem_wr__addr\][^\n]*', src)
            print(f'    write context: {ctx.group(0)[:80] if ctx else "not found"}')
        if not read_m:
            ctx = re.search(r'dmem_b3\[mem_addr\][^\n]*', src)
            print(f'    read context:  {ctx.group(0)[:80] if ctx else "not found"}')
        raise RuntimeError(
            '$readmemb form detected but write or read pattern did not match.\n'
            'Paste the output of:\n'
            '  grep -n "dmem_b\\|mem_rd_data\\|_0_ <=\\|dmem_wr__en" '
            '~/church_project/SoC/church_ti60_f225.v | head -40\n'
            'so the patterns can be corrected.'
        )

    # ── Case 3: already patched with cm_dmem_bram ────────────────────────────
    if 'cm_dmem_bram' in src:
        print('  church_ti60_f225.v already uses cm_dmem_bram — no patch needed.')
        return src

    raise RuntimeError(
        'Could not identify dmem pattern in church_ti60_f225.v.\n'
        'Expected one of:\n'
        '  1. Original: "reg [31:0] dmem [N:0]" + "initial begin"\n'
        '  2. $readmemb: "reg [7:0] dmem_b0 [0:N]" + "$readmemb"\n'
        'Please check the file and update patch patterns if needed.'
    )


# ── Main ─────────────────────────────────────────────────────────────────────

def main():
    if len(sys.argv) > 1:
        soc_dir = os.path.abspath(sys.argv[1])
    else:
        # Default: parent of this script's directory
        script_dir = os.path.dirname(os.path.abspath(__file__))
        soc_dir = os.path.dirname(os.path.dirname(script_dir))

    vpath = os.path.join(soc_dir, 'church_ti60_f225.v')
    if not os.path.isfile(vpath):
        # Try one level up (script might be in hardware/soc_combined/)
        alt = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                           '../../..', 'church_ti60_f225.v')
        alt = os.path.normpath(alt)
        if os.path.isfile(alt):
            soc_dir = os.path.dirname(alt)
            vpath = alt
        else:
            print(f'ERROR: church_ti60_f225.v not found at {vpath}')
            print('  Pass the directory containing church_ti60_f225.v as argument.')
            sys.exit(1)

    print(f'SoC directory: {soc_dir}')
    print(f'Reading {vpath} ...')
    src = open(vpath).read()

    # ── Check for existing cm_dmem_bram patch ────────────────────────────────
    if 'cm_dmem_bram' in src:
        print('church_ti60_f225.v already uses cm_dmem_bram.')
        bak_path = vpath + '.pre_direct'
        if os.path.isfile(bak_path):
            print(f'  Reading original data from backup: {bak_path}')
            src = open(bak_path).read()
        else:
            print('Re-generating cm_dmem_bram.v only (no patch to .v needed).')
        # Try to extract depth from the module
        dm = re.search(r'reg \[31:0\] upper_mem \[2048:(\d+)\]', open(vpath).read())
        if dm:
            depth = int(dm.group(1)) + 1
        else:
            depth = 16384  # default
    else:
        # Determine depth from declaration
        dm = DECL_PAT.search(src)
        dm2 = re.search(r'reg \[7:0\] dmem_b0 \[0:(\d+)\]', src)
        if dm:
            depth = int(dm.group(1)) + 1
        elif dm2:
            depth = int(dm2.group(1)) + 1
        else:
            print('ERROR: cannot determine dmem depth.')
            sys.exit(1)

    print(f'  dmem depth: {depth} words ({depth * 4 // 1024} KB)')

    # ── Parse dmem initial values ─────────────────────────────────────────────
    im = INIT_PAT.search(src)
    dmem_data = {}
    if im:
        for line in im.group(1).splitlines():
            lm = re.match(r"    dmem\[(\d+)\] = 32'd(\d+);", line)
            if lm:
                dmem_data[int(lm.group(1))] = int(lm.group(2))
        nonzero = sum(1 for v in dmem_data.values() if v)
        print(f'  initial begin: {len(dmem_data)} entries, {nonzero} non-zero')
        max_addr = max(dmem_data.keys()) if dmem_data else 0
        print(f'  Highest non-zero address: {max_addr}')
    else:
        # Try to extract from .bin files if they exist
        b0 = os.path.join(soc_dir, 'cm_dmem_b0.bin')
        if os.path.isfile(b0):
            print(f'  No initial begin; reading from {b0} ...')
            lanes = [[], [], [], []]
            for n in range(4):
                fp = os.path.join(soc_dir, f'cm_dmem_b{n}.bin')
                with open(fp) as f:
                    for line in f:
                        v = int(line.strip(), 2) if line.strip() else 0
                        lanes[n].append(v)
            count = min(len(l) for l in lanes)
            for idx in range(count):
                w = (lanes[0][idx] |
                     (lanes[1][idx] << 8) |
                     (lanes[2][idx] << 16) |
                     (lanes[3][idx] << 24))
                if w:
                    dmem_data[idx] = w
            print(f'  Loaded {len(dmem_data)} non-zero entries from .bin files')
        else:
            print('WARNING: no initial begin and no .bin files found.')
            print('  cm_dmem_bram.v will have zero INIT (boot code missing).')

    # ── Validate: all data within addresses 0-2047 ────────────────────────────
    hi = max(dmem_data.keys()) if dmem_data else 0
    if hi >= 2048:
        print(f'WARNING: data at address {hi} >= 2048 — beyond EFX_RAM10 range.')
        print('  These will be in upper_mem (zero-initialized by synthesis).')
        print('  If this data is needed at boot, the range needs extending.')

    # ── Generate cm_dmem_bram.v ───────────────────────────────────────────────
    vmod_path = os.path.join(soc_dir, 'cm_dmem_bram.v')
    print(f'\nGenerating {vmod_path} ...')
    vmod_content = gen_cm_dmem_bram(dmem_data, depth)
    with open(vmod_path, 'w') as f:
        f.write(vmod_content)
    print(f'  Written: {vmod_path}  ({len(vmod_content.splitlines())} lines)')

    # ── Patch church_ti60_f225.v ─────────────────────────────────────────────
    if 'cm_dmem_bram' not in src:
        print(f'\nPatching {vpath} ...')
        bak = vpath + '.pre_direct'
        shutil.copy2(vpath, bak)
        print(f'  Backup: {bak}')
        try:
            patched = patch_top(src, depth)
        except RuntimeError as e:
            print(f'ERROR: {e}')
            sys.exit(1)
        with open(vpath, 'w') as f:
            f.write(patched)
        print(f'  Patched: {vpath}')
    else:
        print(f'\n{vpath} already uses cm_dmem_bram — not re-patched.')

    # cm_dmem_bram.v is listed as a separate source in church_soc_cm.xml —
    # do NOT also append it to church_ti60_f225.v or efx_map sees two
    # definitions of module cm_dmem_bram and aborts at Analysis.

    # ── Summary ───────────────────────────────────────────────────────────────
    print('\n' + '=' * 70)
    print('Done. Now run:')
    print()
    print('  cd ~/church_project/SoC')
    print('  nohup bash church-machine/hardware/soc_combined/run_efx_map.sh \\')
    print('       ~/church_project/SoC/church_soc_cm.xml > syn.log 2>&1 &')
    print('  tail -f syn.log    # ~30 min synthesis')
    print()
    print('  nohup bash church-machine/hardware/soc_combined/run_efx_pnr.sh \\')
    print('       ~/church_project/SoC/church_soc_cm.xml > pnr.log 2>&1 &')
    print('  tail -f pnr.log    # ~18 min P&R')
    print()
    print('  sudo openFPGALoader -b titanium_ti60_f225_jtag \\')
    print('       -f outflow/church_soc_cm.hex')
    print()
    print('Expected: CM boots without NULL_CAP fault, bridge sees LUMP_DONE: OK')
    print('=' * 70)


if __name__ == '__main__':
    main()
