#!/usr/bin/env python3
"""Show flash, static RAM and ELF-section use for embedded firmware.

Usage:
    scripts/analyze-fw.py <firmware.elf> <memory.x>
"""

import re
from pathlib import Path
import subprocess
import sys


REGION = re.compile(
    r"^\s*([A-Za-z_][A-Za-z0-9_]*)\s*:\s*"
    r"ORIGIN\s*=\s*([^,]+),\s*LENGTH\s*=\s*([^\s/}]+)",
    re.MULTILINE,
)


def number(text):
    """Parse linker numbers such as 0x2001_0000, 176K, and 1M."""
    text = text.strip().replace("_", "")
    scale = {"K": 1024, "M": 1024 * 1024}.get(text[-1:].upper(), 1)
    return int(text[:-1] if scale != 1 else text, 0) * scale


def run(command):
    """Run a tool and return stdout."""
    return subprocess.run(command, check=True, text=True, capture_output=True).stdout


def llvm_size():
    """Find the llvm-size installed by the active Rust toolchain."""
    sysroot = run(["rustc", "--print", "sysroot"]).strip()
    host = next(
        line.split(": ", 1)[1]
        for line in run(["rustc", "-vV"]).splitlines()
        if line.startswith("host: ")
    )
    tool = Path(sysroot) / "lib" / "rustlib" / host / "bin" / "llvm-size"
    if not tool.is_file():
        raise RuntimeError(
            "llvm-size is unavailable; install it with: rustup component add llvm-tools-preview"
        )
    return tool


def regions(memory_x):
    """Return every named linker region as (name, origin, length)."""
    return [
        (name, number(origin), number(length))
        for name, origin, length in REGION.findall(memory_x.read_text())
    ]


def sections(output):
    """Return llvm-size sections as (name, bytes, address)."""
    result = []
    for line in output.splitlines():
        fields = line.split()
        if len(fields) == 3 and fields[0].startswith("."):
            result.append((fields[0], int(fields[1], 0), int(fields[2], 0)))
    return result


def region_for(address, layout):
    """Return the linker region containing an address, if any."""
    for name, origin, length in layout:
        if origin <= address < origin + length:
            return name
    return "outside"


def main():
    if len(sys.argv) != 3:
        raise ValueError(f"usage: {sys.argv[0]} <firmware.elf> <memory.x>")

    elf = Path(sys.argv[1])
    memory_x = Path(sys.argv[2])
    if not elf.is_file():
        raise ValueError(f"ELF does not exist: {elf}")
    if not memory_x.is_file():
        raise ValueError(f"linker layout does not exist: {memory_x}")

    layout = regions(memory_x)
    if not layout:
        raise ValueError(f"no MEMORY regions found in: {memory_x}")
    section_list = sections(run([str(llvm_size()), "-A", str(elf)]))
    if not section_list:
        raise ValueError(f"no sections found in: {elf}")

    print(f"\nFirmware: {elf}")
    print(f"ELF file size: {elf.stat().st_size:,} bytes")
    print("\nLinker regions:")
    print("  Region       Origin        Used / Capacity              Remaining")
    for name, origin, length in layout:
        used = sum(
            size
            for _, size, address in section_list
            if size and origin <= address < origin + length
        )
        remaining = length - used
        print(
            f"  {name:<12} 0x{origin:08X} {used / 1024:>7.2f} KiB /"
            f" {length / 1024:>7.2f} KiB  {remaining / 1024:>7.2f} KiB"
            f" ({100 * used / length:5.2f}%)"
        )

    print("\nELF sections:")
    print("  Region       Address       Size  Section")
    for name, size, address in sorted(section_list, key=lambda section: section[2]):
        if size:
            print(f"  {region_for(address, layout):<12} 0x{address:08X} {size:>10,}  {name}")


if __name__ == "__main__":
    try:
        main()
    except (IndexError, OSError, RuntimeError, ValueError, subprocess.CalledProcessError) as error:
        print(f"error: {error}", file=sys.stderr)
        raise SystemExit(1) from error