#!/usr/bin/env python3
import subprocess
import sys
import os
import struct

def run_command(cmd_list, error_msg):
    try:
        subprocess.run(cmd_list, check=True)
    except subprocess.CalledProcessError as e:
        print(f"{error_msg}: {e}", file=sys.stderr)
        sys.exit(1)

def compile_c_source(source_file, elf_file):
    # Compile the source file with no standard library.
    cmd = [
        "riscv64-unknown-elf-gcc",
        "-nostdlib",
        "-static",
        "-O2",
        "-o", elf_file,
        source_file
    ]
    print("Compiling source file...")
    run_command(cmd, "Compilation failed")

def convert_elf_to_bin(elf_file, bin_file):
    # Convert the compiled ELF to a raw binary.
    cmd = [
        "riscv64-unknown-elf-objcopy",
        "-O", "binary",
        elf_file,
        bin_file
    ]
    print("Converting ELF to raw binary...")
    run_command(cmd, "Objcopy conversion failed")

def generate_cpp_vector(bin_file, cpp_file, vector_name="program_data"):

    # Read the binary file.
    print("Reading raw binary...")
    with open(bin_file, "rb") as bf:
        data = bf.read()
    
    # Group the binary into 32-bit words (assuming little-endian)
    words = []
    for i in range(0, len(data), 4):
        chunk = data[i:i+4]
        # If the last chunk is not 4 bytes, pad with zeros.
        if len(chunk) < 4:
            chunk = chunk.ljust(4, b'\x00')
        # Unpack little-endian unsigned int.
        word = struct.unpack("<I", chunk)[0]
        words.append(word)

    # Create the C++ file with the vector.
    print("Writing C++ vector file...")
    with open(cpp_file, "w") as cf:
        cf.write("// Generated C++ vector containing program data.\n")
        cf.write("#include <vector>\n")
        cf.write("#include <cstdint>\n\n")
        cf.write(f"std::vector<uint32_t> {vector_name} = {{\n")

        # Format: 8 words per line.
        for i, word in enumerate(words):
            # Print in hex with 0x... formatting.
            cf.write(f"  0x{word:08X}, ")
            if (i + 1) % 8 == 0:
                cf.write("\n")
        cf.write("\n};\n")
    print(f"C++ file '{cpp_file}' generated successfully.")

def main():
    if len(sys.argv) < 2:
        print("Usage: ./build_and_generate.py <source.c> [<output.cpp>]", file=sys.stderr)
        sys.exit(1)
    
    source_file = sys.argv[1]
    # Default names for intermediate and output files.
    elf_file = "program.elf"
    bin_file = "program.bin"
    cpp_file = sys.argv[2] if len(sys.argv) >= 3 else "program_data.cpp"

    # Check that the source file exists.
    if not os.path.exists(source_file):
        print(f"Source file '{source_file}' does not exist.", file=sys.stderr)
        sys.exit(1)

    compile_c_source(source_file, elf_file)
    convert_elf_to_bin(elf_file, bin_file)
    generate_cpp_vector(bin_file, cpp_file)

    # Optionally, clean up the intermediate files.
    os.remove(elf_file)
    os.remove(bin_file)

if __name__ == "__main__":
    main()