#!/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)