import argparse
import os
import re
import signal
import subprocess


def run_with_timeout(cmd, timeout):
    try:
        proc = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            errors='replace',
            start_new_session=True
        )
        stdout, stderr = proc.communicate(timeout=timeout)
        return proc.returncode, stdout, stderr
    except subprocess.TimeoutExpired:
        os.killpg(proc.pid, signal.SIGTERM)
        try:
            stdout, stderr = proc.communicate(timeout=5)
        except subprocess.TimeoutExpired:
            os.killpg(proc.pid, signal.SIGKILL)
            stdout, stderr = proc.communicate()
        return -9, stdout, stderr
    
def get_io_widths(vpath, top_module):
    if top_module is not None:
        script = f"read_verilog -sv {vpath}; synth -top {top_module};"
    else:
        script = f"read_verilog -sv {vpath}; synth -auto-top;"
    cmd = ["yosys", "-p", script]
    #breakpoint()
    retcode, stdout, stderr = run_with_timeout(cmd, timeout=60)
    if retcode != 0:
        raise RuntimeError(f"Yosys failed with return code {retcode}:\nSTDOUT:\n{stdout}\nSTDERR:\n{stderr}")
    input_pattern = r"ABC RESULTS:\s+input signals:\s+(\d+)"
    output_pattern = r"ABC RESULTS:\s+output signals:\s+(\d+)"
    yosys_log = stdout.split("Re-integrating ABC results.")[-1]
    n_inputs = int(re.findall(input_pattern, yosys_log)[0])
    n_outputs = int(re.findall(output_pattern, yosys_log)[0])
    return n_inputs, n_outputs

def insert_io_notes(input_path, output_path, top_module):
    n_inputs, n_outputs = get_io_widths(input_path, top_module=top_module)
    note = f"// Ninputs {n_inputs}\n// Noutputs {n_outputs}\n\n"

    with open(input_path, "r") as f:
        verilog_code = f.read()
    
    modified_code = note + verilog_code

    with open(output_path, "w") as f:
        f.write(modified_code)

def main():
    parser = argparse.ArgumentParser(description="在verilog代码前插入IO位数注释.")
    parser.add_argument("input_path", help="Path to the input Verilog file.")
    parser.add_argument("-o", "--output_path", help="Path to the output Verilog file. Default using same path as input file.", default=None)
    parser.add_argument("-t", "--top_name", help="Top module name. Default using yosys auto-top command.", default=None)
    args = parser.parse_args()

    if args.output_path is None:
        args.output_path = args.input_path

    insert_io_notes(args.input_path, args.output_path, top_module=args.top_name)

if __name__ == "__main__":
    main()
