#!/usr/bin/env python3
"""Parse Keil MDK-ARM .map file and generate structured RAM/Flash usage report.

Usage:
    python scripts/ram_report.py [MAP_FILE] [-o OUTPUT]

Output: Markdown tables to stdout or file, ready for inclusion in docs.
"""
import re
import sys
import os
import pathlib
import argparse
from datetime import datetime

MAP_DEFAULT = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
    "MDK-ARM", "STM32F407-Demo", "STM32F407-Demo.map",
)

# Rows in "Image component sizes" that re-aggregate objects (would double-count sums).
TOTALS_ROWS = {"Object Totals", "Library Totals", "User Totals"}

# ARM library aggregate rows (e.g. mc_w.l / mf_w.l / m_wm.l): the linker prints one
# aggregate line per library in addition to the individual member .o lines, so any
# "<lib>.l" row is a re-aggregation and must be dropped from per-file sums.
# (2026-10-11: sha256's 64-bit shifts pulled in m_wm.l — the hardcoded name list
#  missed it and the self-consistency check went +58; match the name pattern instead.)
_LIB_AGG_RE = re.compile(r'^[a-z]{1,2}_[a-z0-9]+\.l$')


def _is_aggregate(name):
    """Aggregate rows to drop from per-file sums: library-name summaries and Totals lines.

    '(incl. Padding)' / '(incl. Generated)' are linker metadata ledger rows; they ARE
    kept so that per-file sums reconcile exactly with Grand Totals (self-check Δ=0).
    """
    return name in TOTALS_ROWS or bool(_LIB_AGG_RE.match(name))


def parse_execution_region_header(line):
    """Parse: Execution Region RW_IRAM1 (Exec base: 0x20000000, ..., Size: 0x0001e8b0, Max: 0x00020000, ...)"""
    m = re.search(
        r'Execution Region (\w+) \((.*?)\)', line)
    if not m:
        return None
    name = m.group(1)
    body = m.group(2)
    # Split by comma first (handles multi-word keys like "Exec base")
    params = {}
    for part in body.split(","):
        part = part.strip()
        if ":" in part:
            k, v = part.split(":", 1)
            params[k.strip()] = v.strip()
    return {
        "name": name,
        "exec_base": int(params.get("Exec base", "0"), 16),
        "size": int(params.get("Size", "0"), 16),
        "max": int(params.get("Max", "0"), 16),
    }


def parse_map(path):
    text = pathlib.Path(path).read_text(encoding="utf-8", errors="replace")
    lines = text.splitlines()

    result = {
        "regions": [],
        "iram1_entries": [],
        "grand_totals": None,
        "totals": {},
        "components": [],
    }

    section = None  # track which section we're in
    region_name = None
    skip_header = 0

    i = 0
    while i < len(lines):
        ln = lines[i].rstrip()

        # Detect execution region headers
        if "Execution Region" in ln and "Exec base:" in ln:
            region = parse_execution_region_header(ln)
            if region:
                result["regions"].append(region)
                region_name = region["name"]
                section = "region_entries"
                skip_header = 2  # skip header line + blank
                i += 1
                continue

        # Detect "Image component sizes"
        if ln.strip() == "Image component sizes":
            section = "components"
            skip_header = 3  # skip header + blank + column header
            i += 1
            continue

        # Detect Grand Totals
        if "Grand Totals" in ln:
            parts = ln.split()
            # 219976 11898 11496 528 124576 5416149 Grand Totals
            nums = []
            for p in parts:
                try:
                    nums.append(int(p))
                except ValueError:
                    if "Grand" in p:
                        break
            if len(nums) >= 6:
                result["grand_totals"] = {
                    "code": nums[0], "inc_data": nums[1], "ro_data": nums[2],
                    "rw_data": nums[3], "zi_data": nums[4], "debug": nums[5],
                }
            section = None
            i += 1
            continue

        # Detect Total RO/RW/ROM lines
        m_tot = re.match(r'\s+Total (RO|RW|ROM)\s+Size.*?\(\s*(\d+)\s*\(\s*([\d.]+)kB\)', ln)
        if m_tot:
            result["totals"][m_tot.group(1)] = {
                "bytes": int(m_tot.group(2)),
                "kb": float(m_tot.group(3)),
            }
            i += 1
            continue

        # Parse region entries
        if section == "region_entries" and region_name == "RW_IRAM1":
            if skip_header > 0:
                skip_header -= 1
                i += 1
                continue
            if ln.strip() == "" or ln.strip().startswith("Exec Addr") or ln.strip().startswith("Load Addr"):
                i += 1
                continue
            if ln.strip().startswith("0x"):
                entry = _parse_region_entry(ln)
                if entry:
                    result["iram1_entries"].append(entry)
            elif not ln.strip().startswith("0x") and ln.strip():
                # Hit a new section or summary
                section = None

        # Parse component sizes
        if section == "components":
            if skip_header > 0:
                skip_header -= 1
                i += 1
                continue
            if ln.strip() == "" or "Object Name" in ln:
                i += 1
                continue
            comp = _parse_component_line(ln)
            if comp:
                result["components"].append(comp)

        i += 1

    return result


def _parse_region_entry(ln):
    """Parse a Memory Map entry line like:
    0x20000000   COMPRESSED   0x00000020   Data   RW          292    .data               freertos.o
    0x20000039   COMPRESSED   0x00000003   PAD
    """
    parts = ln.split()
    if len(parts) < 4:
        return None
    try:
        exec_addr = int(parts[0], 16)
    except ValueError:
        return None
    load_addr = parts[1]
    try:
        size = int(parts[2], 16)
    except ValueError:
        return None
    typ = parts[3]
    if typ == "PAD":
        return {"exec_addr": exec_addr, "load_addr": load_addr, "size": size,
                "type": "PAD", "attr": "", "idx": 0, "section": "", "object": ""}
    if len(parts) < 7:
        return None
    attr = parts[4] if parts[4] not in (".data", ".bss", ".text") else ""
    # Find idx (integer) and section (.something) and object (*.o)
    idx = 0
    section = ""
    obj_parts = []
    for j in range(4, len(parts)):
        p = parts[j]
        if p.isdigit():
            idx = int(p)
        elif p.startswith("."):
            section = p
        elif ".o" in p or ".l(" in p:
            obj_parts.append(p)
        elif section and not p.startswith("."):
            obj_parts.append(p)
    obj = " ".join(obj_parts) if obj_parts else ""
    return {"exec_addr": exec_addr, "load_addr": load_addr, "size": size,
            "type": typ, "attr": attr, "idx": idx, "section": section, "object": obj}


def _parse_component_line(ln):
    """Parse: 2738 536 262 1 4104 41449 ch395f.o"""
    parts = ln.split()
    if len(parts) < 7:
        return None
    try:
        code = int(parts[0])
        inc_data = int(parts[1])
        ro_data = int(parts[2])
        rw_data = int(parts[3])
        zi_data = int(parts[4])
        debug = int(parts[5])
    except ValueError:
        return None
    obj_name = " ".join(parts[6:])
    return {
        "code": code, "inc_data": inc_data, "ro_data": ro_data,
        "rw_data": rw_data, "zi_data": zi_data, "debug": debug,
        "object": obj_name,
    }


def fmt(n):
    """Format number with comma separator."""
    return f"{n:,}"


def hexfmt(n):
    return f"0x{n:08X}"


def generate_report(data, map_path):
    now = datetime.now().strftime("%Y-%m-%d %H:%M")
    lines = []
    w = lines.append

    w(f"# STM32F4-Base RAM/Flash Usage Report")
    w(f"")
    w(f"> Generated: {now}　Map: `{os.path.basename(map_path)}`")
    w("")

    # Grand totals
    gt = data["grand_totals"]
    if gt:
        total_ro = gt["code"] + gt["ro_data"]
        total_rw = gt["rw_data"] + gt["zi_data"]
        total_rom = total_ro + gt["rw_data"]
        w("## 1. Grand Totals (全局总量)")
        w("")
        w("| Metric | Value |")
        w("|---|---|")
        w(f"| Code (inc. data) | {fmt(gt['code'])} B ({gt['code']/1024:.1f} KB) |")
        w(f"| RO Data | {fmt(gt['ro_data'])} B ({gt['ro_data']/1024:.1f} KB) |")
        w(f"| RW Data (initialized) | {fmt(gt['rw_data'])} B ({gt['rw_data']/1024:.2f} KB) |")
        w(f"| ZI Data (zero-init) | {fmt(gt['zi_data'])} B ({gt['zi_data']/1024:.1f} KB) |")
        w(f"| **Total RO (Code + RO)** | **{fmt(total_ro)} B ({total_ro/1024:.1f} KB)** |")
        w(f"| **Total RW (RW + ZI) = SRAM used** | **{fmt(total_rw)} B ({total_rw/1024:.2f} KB)** |")
        w(f"| **Total ROM (Code + RO + RW init)** | **{fmt(total_rom)} B ({total_rom/1024:.1f} KB)** |")
        w("")

    # Execution regions
    w("## 2. 执行区水位 (Execution Regions)")
    w("")
    w("| Region | Base | Size | Max | Util% |")
    w("|---|---|---|---|---|")
    for r in data["regions"]:
        util = r["size"] / r["max"] * 100 if r["max"] else 0
        w(f"| {r['name']} | {hexfmt(r['exec_base'])} | {fmt(r['size'])} B ({r['size']/1024:.2f} KB) | {fmt(r['max'])} B ({r['max']/1024:.0f} KB) | {util:.1f}% |")
    w("")

    # IRAM1 entries (ZI + RW symbol detail)
    zii_entries = [e for e in data["iram1_entries"] if e["type"] in ("Zero", "Data") and e["size"] > 0]
    zii_entries.sort(key=lambda e: e["size"], reverse=True)
    w("## 3. RW_IRAM1 逐符号明细（ZI + RW，按大小降序）")
    w("")
    w("| # | Address | Size (B) | Size (KB) | Type | Section | Object |")
    w("|---|---|---|---|---|---|---|")
    for idx, e in enumerate(zii_entries, 1):
        kb = e["size"] / 1024
        w(f"| {idx} | {hexfmt(e['exec_addr'])} | {fmt(e['size'])} | {kb:.2f} | {e['type']} | {e['section']} | {e['object']} |")
    total_zii = sum(e["size"] for e in zii_entries)
    w(f"| | | **{fmt(total_zii)}** | **{total_zii/1024:.2f}** | | | **Total** |")
    w("")

    # Image component sizes (sorted by RW+ZI desc) — 仅保留 *.o 个体行，求和与 Grand Totals 对齐
    comps = [c for c in data["components"] if not _is_aggregate(c["object"])]
    comps.sort(key=lambda c: c["rw_data"] + c["zi_data"], reverse=True)
    top_n = [c for c in comps if c["rw_data"] + c["zi_data"] > 0]
    w("## 4. Image Component Sizes（RW + ZI > 0，按大小降序）")
    w("")
    w("| # | Object | Code | RO | RW | ZI | RW+ZI |")
    w("|---|---|---|---|---|---|---|")
    for idx, c in enumerate(top_n, 1):
        rw_zi = c["rw_data"] + c["zi_data"]
        w(f"| {idx} | {c['object']} | {fmt(c['code'])} | {fmt(c['ro_data'])} | {fmt(c['rw_data'])} | {fmt(c['zi_data'])} | {fmt(rw_zi)} |")
    total_comp_rw = sum(c["rw_data"] for c in comps)
    total_comp_zi = sum(c["zi_data"] for c in comps)
    total_comp_code = sum(c["code"] for c in comps)
    total_comp_ro = sum(c["ro_data"] for c in comps)
    w(f"| | **All .o files** | **{fmt(total_comp_code)}** | **{fmt(total_comp_ro)}** | **{fmt(total_comp_rw)}** | **{fmt(total_comp_zi)}** | **{fmt(total_comp_rw+total_comp_zi)}** |")
    w("")

    # Flash breakdown (sorted by Code+RO desc)
    flash_comps = sorted(comps, key=lambda c: c["code"] + c["ro_data"], reverse=True)
    flash_top = [c for c in flash_comps if c["code"] + c["ro_data"] > 100]
    w("## 5. Flash (Code + RO Data) Top-30（按大小降序）")
    w("")
    w("| # | Object | Code | RO Data | Code+RO |")
    w("|---|---|---|---|---|")
    for idx, c in enumerate(flash_top[:30], 1):
        c_ro = c["code"] + c["ro_data"]
        w(f"| {idx} | {c['object']} | {fmt(c['code'])} | {fmt(c['ro_data'])} | {fmt(c_ro)} |")
    total_code = sum(c["code"] for c in comps)
    total_ro_d = sum(c["ro_data"] for c in comps)
    w(f"| | **All .o files** | **{fmt(total_code)}** | **{fmt(total_ro_d)}** | **{fmt(total_code+total_ro_d)}** |")
    w("")

# Self-consistency check against Grand Totals
    # Per-file sums exclude "(incl. Padding/Generated)" linker metadata; their bytes are
    # counted in Grand Totals, so a small positive delta (~300 B here) is expected.
    if gt:
        tol = 8
        deltas = {
            "Code": total_comp_code - gt["code"],
            "RO": total_comp_ro - gt["ro_data"],
            "RW": total_comp_rw - gt["rw_data"],
            "ZI": total_comp_zi - gt["zi_data"],
        }
        ok = all(abs(d) <= tol for d in deltas.values())
        status = "PASS" if ok else "FAIL"
        if all(d == 0 for d in deltas.values()):
            note = "精确一致（含 '('incl. Generated)'-'(incl. Padding)' 台账行）"
        else:
            note = "差值=图像对齐填充（incl. Padding/Generated 行）"
        deltas_txt = " / ".join(f"{k}Δ={d:+d}" for k, d in deltas.items())
        w(f"> 自洽校验：per-file {fmt(total_comp_code)}/{fmt(total_comp_ro)}/{fmt(total_comp_rw)}/{fmt(total_comp_zi)}"
          f" vs Grand Totals {fmt(gt['code'])}/{fmt(gt['ro_data'])}/{fmt(gt['rw_data'])}/{fmt(gt['zi_data'])}"
          f"（{deltas_txt}；{note}）→ **{status}**")
        w("")

    return "\n".join(lines)


def main():
    ap = argparse.ArgumentParser(description="Parse Keil .map → Markdown RAM/Flash report")
    ap.add_argument("map_file", nargs="?", default=MAP_DEFAULT, help="Path to .map file")
    ap.add_argument("-o", "--output", help="Output file (default: stdout)")
    args = ap.parse_args()

    if not os.path.exists(args.map_file):
        print(f"ERROR: Map file not found: {args.map_file}", file=sys.stderr)
        sys.exit(1)

    data = parse_map(args.map_file)
    report = generate_report(data, args.map_file)

    if args.output:
        pathlib.Path(args.output).parent.mkdir(parents=True, exist_ok=True)
        # newline="\n"：Windows 下 write_text 默认把 \n 翻成 \r\n，会跟 .gitattributes
        # （* text=auto eol=lf）与 data/ 目录既有快照的行尾不一致（2026-10-22 实测）。
        pathlib.Path(args.output).write_text(report, encoding="utf-8", newline="\n")
        print(f"OK: Report written to {args.output}", file=sys.stderr)
    else:
        print(report)


if __name__ == "__main__":
    main()