#!/usr/bin/env python3
"""Validate and show flash, static RAM and ELF-section use.

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

import argparse
import json
import re
from collections import namedtuple
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,
)

Section = namedtuple("Section", "name size address allocated executable")
LoadSegment = namedtuple(
    "LoadSegment", "virtual_address physical_address file_size memory_size"
)


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):
    return subprocess.run(command, check=True, text=True, capture_output=True).stdout


def llvm_tool(name):
    """Find one LLVM binary 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" / name
    if not tool.is_file():
        raise RuntimeError(
            f"{name} 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 elf_layout(output):
    """Return allocated-section and load-segment metadata from llvm-readobj JSON."""
    documents = json.loads(output)
    if len(documents) != 1:
        raise ValueError("llvm-readobj returned no unique ELF document")
    document = documents[0]

    section_list = []
    for wrapped in document.get("Sections", []):
        section = wrapped["Section"]
        flags = {flag["Name"] for flag in section["Flags"]["Flags"]}
        section_list.append(
            Section(
                section["Name"]["Name"],
                section["Size"],
                section["Address"],
                "SHF_ALLOC" in flags,
                "SHF_EXECINSTR" in flags,
            )
        )

    load_segments = []
    for wrapped in document.get("ProgramHeaders", []):
        segment = wrapped["ProgramHeader"]
        if segment["Type"]["Name"] == "PT_LOAD":
            load_segments.append(
                LoadSegment(
                    segment["VirtualAddress"],
                    segment["PhysicalAddress"],
                    segment["FileSize"],
                    segment["MemSize"],
                )
            )
    return section_list, load_segments


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 region_for_span(address, size, layout):
    """Return the one linker region containing the complete nonempty span."""
    end = address + size
    for name, origin, length in layout:
        if origin <= address and end <= origin + length:
            return name
    return None


def merged_span_size(spans):
    """Return bytes covered by the union of half-open address spans."""
    merged = []
    for start, end in sorted(spans):
        if not merged or start > merged[-1][1]:
            merged.append([start, end])
        else:
            merged[-1][1] = max(merged[-1][1], end)
    return sum(end - start for start, end in merged)


def region_usage(name, layout, section_list, load_segments):
    """Return runtime or load-image bytes occupying one linker region.

    Section VMAs count runtime RAM. PT_LOAD physical spans count flash storage,
    including RAM initializers. Their union avoids counting flash sections twice.
    """
    spans = [
        (section.address, section.address + section.size)
        for section in section_list
        if section.allocated
        and section.size
        and region_for_span(section.address, section.size, layout) == name
    ]
    spans.extend(
        (
            segment.physical_address,
            segment.physical_address + segment.file_size,
        )
        for segment in load_segments
        if segment.file_size
        and region_for_span(segment.physical_address, segment.file_size, layout) == name
    )
    return merged_span_size(spans)


def validation_errors(layout, section_list, load_segments):
    """Return every placement error that makes a raw firmware image unsafe."""
    errors = []
    for section in section_list:
        if not section.allocated or section.size == 0:
            continue
        region = region_for_span(section.address, section.size, layout)
        if region is None:
            errors.append(
                f"allocated section {section.name} spans "
                f"0x{section.address:08X}..0x{section.address + section.size:08X} "
                "outside declared memory"
            )

    for index, segment in enumerate(load_segments):
        if (
            segment.memory_size
            and region_for_span(segment.virtual_address, segment.memory_size, layout)
            is None
        ):
            errors.append(
                f"PT_LOAD[{index}] virtual span "
                f"0x{segment.virtual_address:08X}.."
                f"0x{segment.virtual_address + segment.memory_size:08X} "
                "is outside declared memory"
            )
        if (
            segment.file_size
            and region_for_span(segment.physical_address, segment.file_size, layout)
            is None
        ):
            errors.append(
                f"PT_LOAD[{index}] physical span "
                f"0x{segment.physical_address:08X}.."
                f"0x{segment.physical_address + segment.file_size:08X} "
                "is outside declared memory"
            )
    return errors


def flash_operation_section_errors(layout, section_list):
    """Validate the output section that contains the F722 Flash operation."""
    matches = [section for section in section_list if section.name == ".data"]
    if len(matches) != 1:
        return ["firmware must contain exactly one .data section"]
    section = matches[0]
    if not section.allocated or not section.executable or section.size == 0:
        return [".data must contain executable SRAM Flash-operation code"]
    if region_for_span(section.address, section.size, layout) != "RAM":
        return ["executable .data must reside completely in RAM"]
    return []


INSTRUCTION = re.compile(r"^\s*([0-9A-Fa-f]+):\s+([A-Za-z][A-Za-z0-9.]*)\s*(.*?)\s*$")
DIRECT_BRANCH = re.compile(
    r"^(?:b(?:eq|ne|cs|hs|cc|lo|mi|pl|vs|vc|hi|ls|ge|lt|gt|le)?"
    r"(?:\.[nw])?|blx?|bx|cbz|cbnz)$"
)


def sram_control_flow_errors(section, disassembly):
    """Reject external control flow and PC-relative loads from SRAM code."""
    errors = []
    instructions = 0
    end = section.address + section.size
    for line in disassembly.splitlines():
        match = INSTRUCTION.match(line)
        if match is None:
            continue
        instructions += 1
        address = int(match.group(1), 16)
        mnemonic = match.group(2).lower()
        operands = match.group(3).split("@", 1)[0].strip().lower()
        if not section.address <= address < end:
            errors.append(f"instruction 0x{address:08X} is outside .data")
            continue

        if DIRECT_BRANCH.fullmatch(mnemonic):
            target = re.search(r"\b0x([0-9a-f]+)\b", operands)
            if target is None:
                if mnemonic == "bx" and operands == "lr":
                    continue
                errors.append(
                    f"0x{address:08X} uses indirect control transfer {mnemonic}"
                )
                continue
            destination = int(target.group(1), 16)
            if not section.address <= destination < end:
                errors.append(
                    f"0x{address:08X} branches outside SRAM .data to "
                    f"0x{destination:08X}"
                )
            continue

        if mnemonic in ("tbb", "tbh"):
            errors.append(f"0x{address:08X} uses indirect control transfer {mnemonic}")
        elif mnemonic == "pop" and re.search(r"\bpc\b", operands):
            continue
        elif re.match(r"^pc(?:\s|,|$)", operands):
            errors.append(f"0x{address:08X} writes the program counter indirectly")
        elif mnemonic.startswith("ldr") and re.search(r"\[\s*pc\b", operands):
            errors.append(f"0x{address:08X} performs a PC-relative load")

    if instructions == 0:
        errors.append("llvm-objdump found no instructions in executable .data")
    return errors


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("elf", type=Path)
    parser.add_argument("memory_x", type=Path)
    args = parser.parse_args()
    elf, memory_x = args.elf, args.memory_x
    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}")
    readobj = llvm_tool("llvm-readobj")
    section_list, load_segments = elf_layout(
        run(
            [
                str(readobj),
                "--sections",
                "--program-headers",
                "--elf-output-style=JSON",
                str(elf),
            ]
        )
    )
    if not section_list:
        raise ValueError(f"no sections found in: {elf}")
    if not load_segments:
        raise ValueError(f"no loadable segments 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 = region_usage(name, layout, section_list, load_segments)
        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 section in sorted(section_list, key=lambda item: item.address):
        if section.size:
            location = (
                region_for(section.address, layout)
                if section.allocated
                else "non-alloc"
            )
            print(
                f"  {location:<12} "
                f"0x{section.address:08X} {section.size:>10,}  {section.name}"
            )

    errors = validation_errors(layout, section_list, load_segments)
    if errors:
        raise ValueError("firmware memory validation failed:\n  " + "\n  ".join(errors))

    errors = flash_operation_section_errors(layout, section_list)
    if errors:
        raise ValueError(
            "Flash-operation placement validation failed:\n  " + "\n  ".join(errors)
        )
    data = next(section for section in section_list if section.name == ".data")
    objdump = llvm_tool("llvm-objdump")
    errors = sram_control_flow_errors(
        data,
        run(
            [
                str(objdump),
                "-d",
                "--no-show-raw-insn",
                "--section=.data",
                str(elf),
            ]
        ),
    )
    if errors:
        raise ValueError(
            "Flash-operation control-flow validation failed:\n  " + "\n  ".join(errors)
        )


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