#!/usr/bin/env python3
# Simple kernel symbol dumper
# Usage: ksyms.py --input <kernel binary> --output <output file>

import argparse
import subprocess
import sys

excluded_symbols = [
    "text_start_addr",
    "text_end_addr",
    "rodata_start_addr",
    "rodata_end_addr",
    "data_start_addr",
    "data_end_addr",
]


def parse_nm_line(line: str):
    parts = line.split()
    if len(parts) < 3:
        return None

    addr_str, type, name = parts[0], parts[1], parts[2]

    # Only get symbols in text
    if type.lower() != "t":
        return None

    if name in excluded_symbols:
        return None

    addr = int(addr_str, 16)
    return addr, name


def run_nm(binary: str):
    result = subprocess.run(["nm", "-n", binary], capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"nm failed: {result.stderr}")
    return result.stdout.splitlines()


parser = argparse.ArgumentParser(description="Generate kernel symbols from a binary")
parser.add_argument("--input", help="Input kernel binary")
parser.add_argument("--output", help="Output Assembly file (.s)")
parser.add_argument("--empty", action="store_true", help="Generate empty ksyms")
args = parser.parse_args()

with open(args.output, "w") as f:
    f.write("/* Auto-generated by build/ksyms.py - do not edit */\n")
    f.write(".section .rodata\n")

    if args.empty:
        f.write(".global ksyms_array\n")
        f.write(".global ksyms_count\n")
        f.write(".align 8\n")
        f.write("ksyms_array:\n")
        f.write("ksyms_count:\n")
        f.write("    .quad 0\n")
        sys.exit(0)

    lines = run_nm(args.input)
    symbols = []
    for line in lines:
        parsed = parse_nm_line(line)
        if parsed:
            symbols.append(parsed)

    # Sort by address
    symbols.sort(key=lambda x: x[0])

    for i, (addr, name) in enumerate(symbols):
        escaped = name.replace("\\", "\\\\").replace('"', '\\"')
        f.write(f'.Lstr{i}: .asciz "{escaped}"\n')

    f.write("\n.global ksyms_array\n")
    f.write(".global ksyms_count\n")
    f.write(".align 8\n")
    f.write("ksyms_array:\n")

    for i, (addr, name) in enumerate(symbols):
        name_len = len(name.encode("utf-8"))
        f.write(f"    .quad 0x{addr:016x}\n")
        f.write(f"    .quad .Lstr{i}\n")
        f.write(f"    .quad {name_len}\n")

    f.write("\nksyms_count:\n")
    f.write(f"    .quad {len(symbols)}\n")