#!/usr/bin/env python3
import argparse
import re
from pathlib import Path


def parse_ports(verilog_path):
    text = Path(verilog_path).read_text()
    module_name = Path(verilog_path).stem
    start = text.index(f"module {module_name}(")
    end = text.index(");", start)
    header = re.sub(r"//.*", "", text[start:end])
    pieces = re.split(r"\b(input|output)\b", header)[1:]
    ports = []
    for direction, body in zip(pieces[0::2], pieces[1::2]):
        body = body.replace("\n", " ")
        body = re.split(r"\b(input|output)\b", body)[0]
        default_width = ""
        width_match = re.match(r"\s*(\[[^\]]+\])\s*(.*)", body)
        if width_match:
            default_width = width_match.group(1)
            rest = width_match.group(2)
        else:
            rest = body
        for item in rest.split(","):
            item = item.strip()
            if not item:
                continue
            item_match = re.match(r"(\[[^\]]+\])\s*(\w+)$", item)
            if item_match:
                width = item_match.group(1)
                name = item_match.group(2)
            else:
                width = default_width
                name = item.split()[-1]
            if re.fullmatch(r"\w+", name):
                ports.append((direction, width, name))
    return ports


def decl(width, name):
    return f"  wire {width + ' ' if width else ''}rocket_{name};"


def const_for(width):
    if not width:
        return "1'b0"
    match = re.fullmatch(r"\[(\d+):(\d+)\]", width.strip())
    bits = abs(int(match.group(1)) - int(match.group(2))) + 1
    return f"{bits}'d0"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--rocket-top", required=True)
    parser.add_argument("--out", required=True)
    args = parser.parse_args()

    ports = parse_ports(args.rocket_top)
    port_by_name = {name: (direction, width) for direction, width, name in ports}
    assigns = []

    for direction, width, name in ports:
        if direction != "input":
            continue
        if name == "debug_systemjtag_reset":
            value = "~jtag_trst_n | core_reset"
        elif name.endswith("_clock") or name == "debug_clock":
            value = "core_clk"
        elif name.endswith("_reset") or name == "debug_reset":
            value = "core_reset"
        elif name == "gps_acq_clk":
            value = "gps_acq_clk"
        elif name == "gps_trck_clk":
            value = "gps_trck_clk"
        elif name == "gps_rf_clk":
            value = "gps_rf_clk"
        elif name == "gps_afe_clk":
            value = "gps_rf_clk"
        elif name == "gps_rf_data":
            value = "gps_rf_data"
        elif name == "gps_sample_align":
            value = "1'b0"
        elif name == "gps_sample_counter":
            value = "32'd0"
        elif name == "pmod_irq":
            value = "pmod_irq"
        elif name == "resetctrl_hartIsInReset_0":
            value = "core_reset"
        elif name == "debug_systemjtag_jtag_TCK":
            value = "jtag_tck"
        elif name == "debug_systemjtag_jtag_TMS":
            value = "jtag_tms"
        elif name == "debug_systemjtag_jtag_TDI":
            value = "jtag_tdi"
        elif name == "debug_systemjtag_mfr_id":
            value = "11'h489"
        elif name == "debug_systemjtag_part_number":
            value = "16'h0001"
        elif name == "debug_systemjtag_version":
            value = "4'h1"
        elif name == "debug_dmactiveAck":
            value = "rocket_debug_dmactive"
        elif name == "uart_rxd":
            value = "uart_rxd"
        elif name == "spi_dq_i":
            value = "spi_dq_i"
        else:
            value = const_for(width)
        assigns.append(f"  assign rocket_{name} = {value};")

    port_map = ",\n".join(f"    .{name}(rocket_{name})" for _, _, name in ports)
    top_name = Path(args.rocket_top).stem
    content = f"""`timescale 1ns/1ps

module asic_rocket_top (
  input  wire       board_clock,
  input  wire       board_reset,
  input  wire       gps_rf_clk,
  input  wire [3:0] gps_rf_data,
  input  wire       pmod_irq,
  input  wire       uart_rxd,
  output wire       uart_txd,
  output wire       spi_sck,
  output wire       spi_csn,
  input  wire [3:0] spi_dq_i,
  output wire [3:0] spi_dq_o,
  output wire [3:0] spi_dq_oe,
  input  wire       jtag_tck,
  input  wire       jtag_tms,
  input  wire       jtag_tdi,
  input  wire       jtag_trst_n,
  output wire       jtag_tdo_data,
  output wire       jtag_tdo_driven,
  output wire       debug_ndreset,
  output wire       debug_dmactive
);

  wire core_clk;
  wire gps_acq_clk;
  wire gps_trck_clk;
  wire pll_locked;
  wire core_reset;

  asic_pll_model u_asic_pll (
    .board_clock(board_clock),
    .board_reset(board_reset),
    .core_clk(core_clk),
    .gps_acq_clk(gps_acq_clk),
    .gps_trck_clk(gps_trck_clk),
    .locked(pll_locked)
  );

  assign core_reset = board_reset | ~pll_locked;

{chr(10).join(decl(width, name) for _, width, name in ports)}

{chr(10).join(assigns)}

  assign uart_txd = {"rocket_uart_txd" if "uart_txd" in port_by_name else "1'b1"};
  assign spi_sck = {"rocket_spi_sck" if "spi_sck" in port_by_name else "1'b0"};
  assign spi_csn = {"rocket_spi_csn" if "spi_csn" in port_by_name else "1'b1"};
  assign spi_dq_o = {"rocket_spi_dq_o" if "spi_dq_o" in port_by_name else "4'h0"};
  assign spi_dq_oe = {"rocket_spi_dq_oe" if "spi_dq_oe" in port_by_name else "4'h0"};
  assign jtag_tdo_data = {"rocket_debug_systemjtag_jtag_TDO_data" if "debug_systemjtag_jtag_TDO_data" in port_by_name else "1'b0"};
  assign jtag_tdo_driven = {"rocket_debug_systemjtag_jtag_TDO_driven" if "debug_systemjtag_jtag_TDO_driven" in port_by_name else "1'b0"};
  assign debug_ndreset = {"rocket_debug_ndreset" if "debug_ndreset" in port_by_name else "1'b0"};
  assign debug_dmactive = {"rocket_debug_dmactive" if "debug_dmactive" in port_by_name else "1'b0"};

  {top_name} u_rocket (
{port_map}
  );
endmodule
"""
    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    Path(args.out).write_text(content)


if __name__ == "__main__":
    main()