#!/usr/bin/env python3

"""Check kernel source limits and report the built image's fixed memory."""

import pathlib

import struct

import sys

def source_sizes(root):

    """Count physical source lines and combine files that form one module."""

    totals = {"runtime": 0, "tests": 0, "tools": 0}

    modules = {}

    for path in sorted(root.rglob("*")):

        if path.suffix not in (".rad", ".ras"):

            continue

        relative = path.relative_to(root)

        category = "runtime"

        if "tests" in relative.parts or path.stem == "tests":

            category = "tests"

        elif relative.parts[0] == "tools":

            category = "tools"

        lines = len(path.read_bytes().splitlines())

        totals[category] += lines

        if category != "tests":

            module = str(relative.with_suffix(""))

            modules[module] = modules.get(module, 0) + lines

    if not totals["runtime"]:

        raise ValueError("kernel runtime sources are missing")

    for module, lines in sorted(modules.items()):

        print(f"module {module}: {lines} source lines")

        if lines >= 1000:

            raise ValueError(f"module {module} must have fewer than 1000 lines")

    for category, lines in totals.items():

        print(f"kernel {category}: {lines} source lines")

    if totals["runtime"] > 12000:

        raise ValueError("kernel runtime exceeds 12000 source lines")

def image_sizes(path):

    """Validate the native image extents and report each fixed memory segment."""

    with path.open("rb") as stream:

        header = stream.read(64)

    if len(header) != 64:

        raise ValueError("native image header is truncated")

    magic, version, entry = struct.unpack_from("<IIQ", header)

    if magic != 0x30444152 or version != 2:

        raise ValueError("expected a version 2 native image")

    segments = []

    payload = 64

    for index, name in enumerate(("code", "read-only", "writable")):

        address, initialized, memory = struct.unpack_from("<QII", header, 16 + index * 16)

        alignment = 4 if index == 0 else 8

        if initialized > memory or address + memory > 1 << 64:

            raise ValueError(f"invalid {name} segment extent")

        if memory and address % alignment:

            raise ValueError(f"misaligned {name} segment")

        segments.append((address, initialized, memory))

        payload += initialized

        print(f"image {name}: {memory} fixed bytes, {initialized} initialized bytes")

    code, initialized, _ = segments[0]

    if entry % 4 or initialized % 4 or not code <= entry < code + initialized:

        raise ValueError("entry is outside initialized instructions")

    for index, (start, _, size) in enumerate(segments):

        for other, _, length in segments[index + 1:]:

            if size and length and start < other + length and other < start + size:

                raise ValueError("native image segments overlap")

    if path.stat().st_size != payload:

        raise ValueError("native image payload size does not match its header")

    print(f"image total: {sum(segment[2] for segment in segments)} fixed bytes")

def main():

    """Check the worktree, or explicit source and image paths for a fixture."""

    if len(sys.argv) not in (1, 3):

        raise ValueError("usage: sizes [kernel-directory native-image]")

    root = pathlib.Path(sys.argv[1] if len(sys.argv) == 3 else "kernel")

    image = pathlib.Path(sys.argv[2] if len(sys.argv) == 3 else "bin/kernel.rv64")

    source_sizes(root)

    image_sizes(image)

if __name__ == "__main__":

    try:

        main()

    except (OSError, ValueError) as error:

        print(f"kernel size check: {error}", file=sys.stderr)

        sys.exit(1)