#!/usr/bin/env python3
"""Apply small, deterministic Efinity synthesis adaptations to generated RTL."""

from __future__ import annotations

import argparse
import math
import re
from pathlib import Path


ROCKET_STORE_DATA = """    else begin
      automatic logic [3:0][63:0] _GEN_24 =
        {{ex_rs_1}, {{2{ex_rs_1[31:0]}}}, {{4{ex_rs_1[15:0]}}}, {{8{ex_rs_1[7:0]}}}};
      mem_reg_rs2 <= _GEN_24[ex_reg_mem_size];
    end"""

EFINITY_STORE_DATA = """    else begin
      // Efinity 2025.2 leaves the block-local packed automatic array above
      // undriven (VDB-1002).  Spell out the byte-lane replication instead.
      case (ex_reg_mem_size)
        2'h0: mem_reg_rs2 <= {8{ex_rs_1[7:0]}};
        2'h1: mem_reg_rs2 <= {4{ex_rs_1[15:0]}};
        2'h2: mem_reg_rs2 <= {2{ex_rs_1[31:0]}};
        default: mem_reg_rs2 <= ex_rs_1;
      endcase
    end"""

ATOMIC_LUT_DECLARATION = re.compile(
    r"^(?P<indent>\s*)automatic logic \[3:0\]\[3:0\] (?P<name>_GEN_[0-9]+) "
    r"= '\{4'hC, 4'h8, 4'hE, 4'h6\};$",
    re.MULTILINE,
)

NCO_ROM_START = '(*rom_style = "distributed" *) logic [PH_ACC_WL-1 : 0] freq_word_arr'
NCO_ROM_END = "logic [PH_ACC_WL-1 : 0] sin_phase, cos_phase;"


def parse_integer_parameter(text: str, name: str) -> int:
    match = re.search(rf"\bparameter\s+{re.escape(name)}\s*=\s*([0-9]+)\s*;", text)
    if match is None:
        raise RuntimeError(f"could not find integer parameter {name}")
    return int(match.group(1))


def parse_sv_integer(token: str) -> tuple[int, int | None]:
    token = token.strip().replace("_", "")
    match = re.fullmatch(r"(-)?([0-9]+)'([dDhHbB])([0-9a-fA-F]+)", token)
    if match is not None:
        sign, width_text, radix_text, digits = match.groups()
        radix = {"d": 10, "h": 16, "b": 2}[radix_text.lower()]
        value = int(digits, radix)
        if sign:
            value = -value
        return value, int(width_text)
    if re.fullmatch(r"-?[0-9]+", token):
        return int(token), None
    raise RuntimeError(f"unsupported SystemVerilog integer literal: {token!r}")


def parse_array_parameter(text: str, name: str) -> list[tuple[int, int | None]]:
    match = re.search(
        rf"\bparameter\s+\[[^\]]+\]\s+{re.escape(name)}\s*\[[^\]]+\]\s*=\s*\{{(.*?)\}}\s*;",
        text,
        re.DOTALL,
    )
    if match is None:
        raise RuntimeError(f"could not find array parameter {name}")
    return [parse_sv_integer(token) for token in match.group(1).split(",") if token.strip()]


def sv_case_function(
    name: str,
    input_width: str,
    output_width: str,
    values: list[int],
    literal_width: int,
) -> str:
    lines = [
        f"function automatic logic [{output_width}-1 : 0] {name};",
        f"    input logic [{input_width}-1 : 0] index;",
        "    begin",
        "        case (index)",
    ]
    for index, value in enumerate(values):
        literal = f"-{literal_width}'sd{-value}" if value < 0 else f"{literal_width}'d{value}"
        lines.append(f"            {index}: {name} = {literal};")
    lines.extend(
        [
            f"            default: {name} = '0;",
            "        endcase",
            "    end",
            "endfunction",
        ]
    )
    return "\n".join(lines)


def round_sv_real(value: float) -> int:
    """Match real-to-integral assignment rounding used by SystemVerilog."""
    return math.floor(value + 0.5) if value >= 0 else math.ceil(value - 0.5)


def build_nco_rom_functions(definitions: str) -> str:
    phase_width = parse_integer_parameter(definitions, "ACQ_NCO_PH_ACC_WL")
    sine_integer_bits = parse_integer_parameter(definitions, "ACQ_SIN_IL")
    sine_fraction_bits = parse_integer_parameter(definitions, "ACQ_SIN_FL")
    lut_address_width = parse_integer_parameter(definitions, "ACQ_SIN_LUT_DEPTH_WL")
    freq_values = [value for value, _ in parse_array_parameter(definitions, "FREQ_WORD_ARRAY")]
    pullin_values = [value for value, _ in parse_array_parameter(definitions, "PULLIN_FREQ_WORD_ARRAY")]
    sine_width = sine_integer_bits + sine_fraction_bits
    sine_values = [
        round_sv_real(
            math.sin(2.0 * math.pi * index / (1 << lut_address_width))
            * (1 << sine_fraction_bits)
        )
        for index in range(1 << lut_address_width)
    ]
    return "\n\n".join(
        [
            "// Efinity 2025.2 mis-elaborates the original initial/$sin ROMs.\n"
            "// These constant functions are generated from definitions.sv for Ti375 only.",
            sv_case_function(
                "efinity_freq_word", "FREQ_BIN_WL", "PH_ACC_WL", freq_values, phase_width
            ),
            sv_case_function(
                "efinity_pullin_word", "NUM_PULLIN_WL", "PH_ACC_WL", pullin_values, phase_width
            ),
            sv_case_function(
                "efinity_sine", "LUT_DEPTH_WL", "ACQ_SIN_IL + ACQ_SIN_FL", sine_values, sine_width
            ),
        ]
    )


def rewrite_acquisition_nco(source: Path, destination: Path, definitions: str) -> None:
    text = source.read_text(encoding="utf-8")
    start = text.find(NCO_ROM_START)
    end = text.find(NCO_ROM_END, start)
    if start < 0 or end < 0:
        raise RuntimeError(f"could not locate acquisition NCO ROM block in {source}")
    text = text[:start] + build_nco_rom_functions(definitions) + "\n\n" + text[end:]
    old_sum = (
        "freq_word_arr[freq_bin_idx_i] + "
        "pullin_freq_word_arr[pullin_freq_bin_idx_i]"
    )
    new_sum = (
        "efinity_freq_word(freq_bin_idx_i) + "
        "efinity_pullin_word(pullin_freq_bin_idx_i)"
    )
    if text.count(old_sum) != 2:
        raise RuntimeError(f"expected two NCO phase increments in {source}")
    text = text.replace(old_sum, new_sum)
    if text.count("sinwav_lut[truncated_sin_phase]") != 1:
        raise RuntimeError(f"expected one sine ROM read in {source}")
    if text.count("sinwav_lut[truncated_cos_phase]") != 1:
        raise RuntimeError(f"expected one cosine ROM read in {source}")
    text = text.replace("sinwav_lut[truncated_sin_phase]", "efinity_sine(truncated_sin_phase)")
    text = text.replace("sinwav_lut[truncated_cos_phase]", "efinity_sine(truncated_cos_phase)")
    destination.write_text(text, encoding="utf-8")


def rewrite_tracking_nco(source: Path, destination: Path, definitions: str) -> None:
    text = source.read_text(encoding="utf-8")
    declaration = (
        '(*rom_style = "distributed" *) logic [SIN_IL-1 : -SIN_FL] '
        "sinwav_lut [0 : 2**LUT_DEPTH_WL-1];"
    )
    values = [value for value, _ in parse_array_parameter(definitions, "TRCK_SINWAV_LUT")]
    expected_entries = parse_integer_parameter(definitions, "NUM_TRCK_SINWAV_LUT")
    if len(values) != expected_entries:
        raise RuntimeError(
            f"expected {expected_entries} tracking sine entries, got {len(values)}"
        )
    function = (
        "// Efinity 2025.2 leaves the original package-array initialized ROM undriven.\n"
        "// This Ti375-only constant function preserves the exact definitions.sv values.\n"
        + sv_case_function(
            "efinity_tracking_sine",
            "LUT_DEPTH_WL",
            "SIN_IL + SIN_FL",
            values,
            parse_integer_parameter(definitions, "TRCK_CARR_NCO_SIN_WL"),
        )
    )
    if text.count(declaration) != 1:
        raise RuntimeError(f"expected one tracking sine ROM declaration in {source}")
    text = text.replace(declaration, function)
    initializer = """initial begin
    sinwav_lut = TRCK_SINWAV_LUT;
end

"""
    if text.count(initializer) != 1:
        raise RuntimeError(f"expected one tracking sine ROM initializer in {source}")
    text = text.replace(initializer, "")
    if text.count("sinwav_lut[accum_truncated_sin]") != 1:
        raise RuntimeError(f"expected one tracking sine ROM read in {source}")
    if text.count("sinwav_lut[accum_truncated_cos]") != 1:
        raise RuntimeError(f"expected one tracking cosine ROM read in {source}")
    text = text.replace(
        "sinwav_lut[accum_truncated_sin]", "efinity_tracking_sine(accum_truncated_sin)"
    )
    text = text.replace(
        "sinwav_lut[accum_truncated_cos]", "efinity_tracking_sine(accum_truncated_cos)"
    )
    destination.write_text(text, encoding="utf-8")


def write_mem_file(path: Path, entries: list[tuple[int, int | None]], width: int) -> None:
    mask = (1 << width) - 1
    path.write_text(
        # Binary keeps each row exactly TWIDDLE_WIDTH bits.  A hexadecimal row
        # would round 15 bits up to four digits and Efinity warns about a
        # spurious 16-to-15-bit truncation even though the value is correct.
        "".join(f"{value & mask:0{width}b}\n" for value, _ in entries),
        encoding="ascii",
    )


def rewrite_twiddle_rom(source: Path, destination: Path) -> None:
    text = source.read_text(encoding="utf-8")
    old = """initial begin
    cos_rom = COS_ROM_ARRAY;
    sin_rom = SIN_ROM_ARRAY;
end"""
    new = """initial begin
    // Generated Ti375 files are resolved through Efinity's memory-init include path.
    $readmemb("ti375_twiddle_cos_15.mem", cos_rom);
    $readmemb("ti375_twiddle_sin_15.mem", sin_rom);
end"""
    if text.count(old) != 1:
        raise RuntimeError(f"expected one twiddle ROM initializer in {source}")
    destination.write_text(text.replace(old, new), encoding="utf-8")


def prepare_gps_rtl(source_dir: Path, output_dir: Path) -> None:
    definitions_path = source_dir / "common" / "definitions.sv"
    definitions = definitions_path.read_text(encoding="utf-8")
    output_dir.mkdir(parents=True, exist_ok=True)
    rewrite_acquisition_nco(
        source_dir / "acq_src" / "acq_nco_carr.sv",
        output_dir / "acq_nco_carr.sv",
        definitions,
    )
    rewrite_tracking_nco(
        source_dir / "trck_src" / "NCO_carr.sv",
        output_dir / "NCO_carr.sv",
        definitions,
    )
    rewrite_twiddle_rom(
        source_dir / "acq_src" / "twiddle_ram.sv",
        output_dir / "twiddle_ram.sv",
    )
    cos_entries = parse_array_parameter(definitions, "COS_ROM_ARRAY")
    sin_entries = parse_array_parameter(definitions, "SIN_ROM_ARRAY")
    expected_entries = parse_integer_parameter(definitions, "NUM_TWIDDLE_ROM_IDX")
    if len(cos_entries) != expected_entries or len(sin_entries) != expected_entries:
        raise RuntimeError(
            f"expected {expected_entries} twiddle entries, got "
            f"cos={len(cos_entries)} sin={len(sin_entries)}"
        )
    write_mem_file(output_dir / "ti375_twiddle_cos_15.mem", cos_entries, 15)
    write_mem_file(output_dir / "ti375_twiddle_sin_15.mem", sin_entries, 15)


def rewrite_rocket_store_data(path: Path) -> None:
    text = path.read_text(encoding="utf-8")
    if EFINITY_STORE_DATA in text:
        return
    if text.count(ROCKET_STORE_DATA) != 1:
        raise RuntimeError(f"expected one Rocket store-data array in {path}")
    path.write_text(text.replace(ROCKET_STORE_DATA, EFINITY_STORE_DATA), encoding="utf-8")


def rewrite_tilelink_atomic_lut(path: Path) -> None:
    text = path.read_text(encoding="utf-8")
    if "// Avoid Efinity VDB-1002 on the block-local packed LUT." in text:
        return
    match = ATOMIC_LUT_DECLARATION.search(text)
    if match is None:
        raise RuntimeError(f"expected one TileLink atomic packed LUT in {path}")
    if ATOMIC_LUT_DECLARATION.search(text, match.end()) is not None:
        raise RuntimeError(f"expected only one TileLink atomic packed LUT in {path}")
    indent = match.group("indent")
    temporary = match.group("name")
    old_read = f"{indent}cam_a_0_lut <= {temporary}[auto_in_a_bits_param[1:0]];"
    new_read = f"""{indent}// Avoid Efinity VDB-1002 on the block-local packed LUT.
{indent}case (auto_in_a_bits_param[1:0])
{indent}  2'h0: cam_a_0_lut <= 4'h6;
{indent}  2'h1: cam_a_0_lut <= 4'hE;
{indent}  2'h2: cam_a_0_lut <= 4'h8;
{indent}  default: cam_a_0_lut <= 4'hC;
{indent}endcase"""
    if text.count(old_read) != 1:
        raise RuntimeError(f"expected one TileLink atomic packed LUT read in {path}")
    text = text[: match.start()] + text[match.end() + 1 :]
    text = text.replace(old_read, new_read)
    path.write_text(text, encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--verilog-dir", type=Path, required=True)
    parser.add_argument("--gps-source-dir", type=Path, required=True)
    parser.add_argument("--gps-output-dir", type=Path, required=True)
    args = parser.parse_args()

    rewrite_rocket_store_data(args.verilog_dir / "Rocket.sv")
    rewrite_tilelink_atomic_lut(args.verilog_dir / "TLAtomicAutomata_pbus.sv")
    rewrite_tilelink_atomic_lut(args.verilog_dir / "TLAtomicAutomata_cbus.sv")
    prepare_gps_rtl(args.gps_source_dir, args.gps_output_dir)


if __name__ == "__main__":
    main()