#!/usr/bin/env python3

#
# Generates a Verilog include file on stdout.
#
# Usage:  
#
#     curv-svh-from-plusargs [-o output_file.svh] +TB_ARG1="arg1" +TB_ARG2=hello +SYN_ARG1=72 +WILL_BE_IGNORED=123
#
# This will put all args (except WILL_BE_IGNORED) into both the sim and synth blocks, but the ones that are only specified
# for TB_ will be blank in synthesis, and vice-versa.
#
# Here's a more complex use.
#
#     curv-svh-from-plusargs [-o output_file.svh] curv-svh-from-plusargs \
#       +TB_COMMON_ARG="testbench_version_of_val" +SYN_COMMON_ARG=synthesis_version_of_val \
#       +TB_SHOULD_BE_BLANK_IN_SYNTHESIS="testbench_only_val" \
#       +SYN_SHOULD_BE_BLANK_IN_SIMULATION="synthesis_only_val" \
#       +TBONLY_SHOULD_BE_OMITTED_IN_SYNTHESIS="testbench_only_val2__omitted_in_synthesis" \
#       +SYNONLY_SHOULD_BE_OMITTED_IN_SIMULATION="synthesis_only_val2__omitted_in_simulation" \
#       +WILL_BE_IGNORED=123
#       +BOTH_XXX="both_val"   # will appear in both simulation and synthesis as "BOTH_XXX=both_val"
#
# This does the same, but with the twist that TBONLY_ args are omitted entirely from the syntehsis block, and
# vice-versa for SYNONLY_ args.
#

import sys
import re
import argparse
import os
import unittest

REQUIRED_PREFIX_FOR_SIMULATION = "TB_"
REQUIRED_PREFIX_FOR_SIMULATION_ONLY = "TBONLY_"
REQUIRED_PREFIX_FOR_SYNTHESIS  = "SYN_"
REQUIRED_PREFIX_FOR_SYNTHESIS_ONLY = "SYNONLY_"
REQUIRED_PREFIX_FOR_BOTH = "BOTH_"

# used for include guards if you write to stdout instead of using -o
DEFAULT_OUTPUT_FILE_NAME = f"generated-by-{os.path.basename(__file__).replace('.py', '.vh')}"

def get_file_parts(output_file_name: str) -> tuple[str, str, str]:
    include_guard = "__" + os.path.basename(output_file_name).replace(".", "_").replace("-", "_").upper() + "__"
    # Display the installed CLI name in the generated file header for clarity
    script_name = "curv-svh-from-plusargs"
    file_preamble = f"""`ifndef {include_guard}
`define {include_guard}

//
// Automatically generated by '{script_name}' during make
//

`ifdef SIMULATION

"""

    file_middle = """
`else // SYNTHESIS

"""

    file_postamble = f"""
`endif

`endif // {include_guard}
"""

    return file_preamble, file_middle, file_postamble

def normalize_plusargs_tokens(tokens: list[str]) -> list[str]:
    normalized = []
    i = 0
    n = len(tokens)
    while i < n:
        t0 = tokens[i]
        # Pass through non-plusargs tokens unchanged
        if not t0.startswith("+"):
            normalized.append(t0)
            i += 1
            continue

        key_part = t0
        value_part = None

        if "=" in key_part:
            # Cases: "+KEY=value" or "+KEY=" followed by value
            if key_part.endswith("="):
                if i + 1 < n:
                    value_part = tokens[i + 1]
                    # Drop the trailing '=' from key_part
                    key_part = key_part[:-1]
                    i += 2
                else:
                    value_part = ""
                    key_part = key_part[:-1]
                    i += 1
            else:
                k, v = key_part.split("=", 1)
                key_part = k
                value_part = v
                i += 1
        else:
            # No '=' in current token: look ahead to assemble
            if i + 1 < n:
                t1 = tokens[i + 1]
                if t1 == "=":
                    if i + 2 < n:
                        value_part = tokens[i + 2]
                        i += 3
                    else:
                        value_part = ""
                        i += 2
                elif t1.startswith("="):
                    value_part = t1[1:]
                    i += 2
                else:
                    value_part = t1
                    i += 2
            else:
                value_part = ""
                i += 1

        ktrim = key_part.strip()
        vtrim = value_part.strip() if value_part is not None else ""
        # Normalize to no whitespace around '=' as requested
        normalized.append(f"{ktrim}={vtrim}")

    return normalized

def parse_args(args):
    argparser = argparse.ArgumentParser()
    argparser.add_argument("-o", "--output", type=str, required=False, default=None, help="output file (if blank, stdout is used)")
    argparser.add_argument("-t", "--test", action="store_true", help="run unit tests and exit")
    argparser.add_argument("plusargs", nargs="*", help="plusargs to parse; TB_ prefix for simulation, SYN_ prefix for synthesis, BOTH_ prefix for both")
    args = argparser.parse_args(args)
    # Reassemble split plusargs like ['+KEY', '=', 'VALUE'] → ['+KEY = VALUE']
    args.plusargs = normalize_plusargs_tokens(args.plusargs or [])

    return args

class TestParseArgs(unittest.TestCase):
    def test_old_style_no_whitespace(self):
        ns = parse_args(["+TB_KEY=value"])
        self.assertEqual(ns.plusargs, ["+TB_KEY=value"]) 

    def test_key_then_equals_value(self):
        ns = parse_args(["+TB_KEY", "=value"]) 
        self.assertEqual(ns.plusargs, ["+TB_KEY=value"]) 

    def test_key_equals_then_value(self):
        ns = parse_args(["+TB_KEY=", "value"]) 
        self.assertEqual(ns.plusargs, ["+TB_KEY=value"]) 

    def test_key_space_eq_space_value(self):
        ns = parse_args(["+TB_KEY", "=", "value"]) 
        self.assertEqual(ns.plusargs, ["+TB_KEY=value"]) 

    def test_preserve_quotes(self):
        ns = parse_args(["+TB_KEY", "=", '"some path with spaces"'])
        self.assertEqual(ns.plusargs, ['+TB_KEY="some path with spaces"'])

def run_tests() -> int:
    suite = unittest.TestLoader().loadTestsFromTestCase(TestParseArgs)
    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(suite)
    return 0 if result.wasSuccessful() else 1

def get_plusargs_dict(args: argparse.Namespace, plusargs_simulation: dict[str, str], plusargs_synthesis: dict[str, str], plusargs_simulation_only: dict[str, str], plusargs_synthesis_only: dict[str, str]):
    # Process each argument
    for i, arg in enumerate(args.plusargs, 1):
        # Match key="value" or key=value pattern
        match_simulation = re.match(fr'\+{REQUIRED_PREFIX_FOR_SIMULATION}(.+?)=(?:"(.+?)"|(.+))', arg)
        match_synthesis  = re.match(fr'\+{REQUIRED_PREFIX_FOR_SYNTHESIS}(.+?)=(?:"(.+?)"|(.+))', arg)
        match_simulation_only = re.match(fr'\+{REQUIRED_PREFIX_FOR_SIMULATION_ONLY}(.+?)=(?:"(.+?)"|(.+))', arg)
        match_synthesis_only = re.match(fr'\+{REQUIRED_PREFIX_FOR_SYNTHESIS_ONLY}(.+?)=(?:"(.+?)"|(.+))', arg)
        match_both = re.match(fr'\+{REQUIRED_PREFIX_FOR_BOTH}(.+?)=(?:"(.+?)"|(.+))', arg)
        if match_simulation:
            key = match_simulation.group(1)
            # Group 2 is the quoted value, group 3 is the unquoted value
            value = match_simulation.group(2) if match_simulation.group(2) is not None else match_simulation.group(3)
            # strip leading + if present
            if key.startswith("+"):
                key = key[1:]
            # save to dict
            plusargs_simulation[key] = value
        elif match_synthesis:
            key = match_synthesis.group(1)
            # Group 2 is the quoted value, group 3 is the unquoted value
            value = match_synthesis.group(2) if match_synthesis.group(2) is not None else match_synthesis.group(3)
            # strip leading + if present
            if key.startswith("+"):
                key = key[1:]
            # save to dict
            plusargs_synthesis[key] = value
        elif match_simulation_only:
            key = match_simulation_only.group(1)
            # Group 2 is the quoted value, group 3 is the unquoted value
            value = match_simulation_only.group(2) if match_simulation_only.group(2) is not None else match_simulation_only.group(3)
            # save to dict
            plusargs_simulation_only[key] = value
        elif match_synthesis_only:
            key = match_synthesis_only.group(1)
            # Group 2 is the quoted value, group 3 is the unquoted value
            value = match_synthesis_only.group(2) if match_synthesis_only.group(2) is not None else match_synthesis_only.group(3)
            # save to dict
            plusargs_synthesis_only[key] = value
        elif match_both:
            key = match_both.group(1)
            # Group 2 is the quoted value, group 3 is the unquoted value
            value = match_both.group(2) if match_both.group(2) is not None else match_both.group(3)
            # strip leading + if present
            if key.startswith("+"):
                key = key[1:]
            # save to dict
            plusargs_synthesis[key] = value
            plusargs_simulation[key] = value


def main():
    # Quick check for -t before heavy processing
    if any(a in ("-t", "--test") for a in sys.argv[1:]):
        exit_code = run_tests()
        sys.exit(exit_code)

    args = parse_args(sys.argv[1:])

    plusargs_simulation = {}
    plusargs_synthesis  = {}
    plusargs_simulation_only = {}
    plusargs_synthesis_only = {}
    get_plusargs_dict(args, plusargs_simulation, plusargs_synthesis, plusargs_simulation_only, plusargs_synthesis_only)

    # calculate the max key length for padding
    if len(plusargs_simulation) > 0 or len(plusargs_synthesis) > 0 or len(plusargs_simulation_only) > 0 or len(plusargs_synthesis_only) > 0:
        # Find longest key to determine padding
        MIN_SPACING = 10
        combined_keys = list(plusargs_simulation.keys()) + list(plusargs_synthesis.keys()) + list(plusargs_simulation_only.keys()) + list(plusargs_synthesis_only.keys())
        max_key_len = max(len(key) for key in combined_keys) + MIN_SPACING
    else:
        print("no plusargs!")
        return 1
    
    # find the keys that are in both dictionaries, in simulation only, or synthesis only
    common_keys = set(plusargs_simulation.keys()) & set(plusargs_synthesis.keys())
    simulation_only_blank_in_synthesis_keys = set(plusargs_simulation.keys()) - set(plusargs_synthesis.keys())
    synthesis_only_blank_in_simulation_keys = set(plusargs_synthesis.keys()) - set(plusargs_simulation.keys())

    # if the argument appears only in simulation,  make a blank version for synthesis
    # if the argument appears only in synthesis,  make a blank version for simulation
    # if the argument appears in both, make a version for each with the correct respective value
    final_kvs = {}
    for key in common_keys:
        tb_value = plusargs_simulation[key]
        syn_value = plusargs_synthesis[key]
        final_kvs[key] = {'tb': tb_value, 'syn': syn_value}
    for key in simulation_only_blank_in_synthesis_keys:
        tb_value = plusargs_simulation[key]
        final_kvs[key] = {'tb': tb_value, 'syn': ""}
    for key in synthesis_only_blank_in_simulation_keys:
        syn_value = plusargs_synthesis[key]
        final_kvs[key] = {'tb': "", 'syn': syn_value}
    for key in plusargs_simulation_only.keys():
        tb_value = plusargs_simulation_only[key]
        final_kvs[key] = {'tb': tb_value, 'syn': None}
    for key in plusargs_synthesis_only.keys():
        syn_value = plusargs_synthesis_only[key]
        final_kvs[key] = {'tb': None, 'syn': syn_value}

    simulation_str = ""
    synthesis_str = ""
    # output sub strs for the simulation and synthesis blocks
    indent = " " * 4
    if len(final_kvs) > 0:
        for key, value in sorted(final_kvs.items()):
            padding = " " * (max_key_len - len(key))
            tb_value = value['tb']
            syn_value = value['syn']
            if tb_value is not None:
                simulation_str += f'{indent}`define {key}{padding} `"{tb_value}`"\n'
            if syn_value is not None:
                synthesis_str += f'{indent}`define {key}{padding} `"{syn_value}`"\n'

    #
    # output file contents
    #
    file_preamble, file_middle, file_postamble = get_file_parts(output_file_name=args.output if args.output else DEFAULT_OUTPUT_FILE_NAME)
    output_str = file_preamble
    output_str += simulation_str
    output_str += file_middle
    output_str += synthesis_str
    output_str += file_postamble

    if args.output:
        with open(args.output, "w") as f:
            f.write(output_str)
    else:
        print(output_str)
    return 0

if __name__ == "__main__":
    sys.exit(main())
 
