#!/usr/bin/env python3
"""Generate memory layout SVG diagrams for STM32F4-Base.

Usage:
    python tools/arch_diagram/gen_memory_svg.py [--type ram|flash|taskstack|all]

Output (default --all):
    docs/images/内存布局_RAM.svg
    docs/images/内存布局_FLASH.svg
    docs/images/任务栈预算.svg
"""
import argparse
import os
import sys

# ---------------------------------------------------------------------------
# Diagram data
# ---------------------------------------------------------------------------

# ── RAM breakdown (128KB = 131,072 B) ────────────────────────────────────────
SRAM_TOTAL = 131072
SRAM_USED  = 125104
SRAM_FREE  = SRAM_TOTAL - SRAM_USED

RAM_CATEGORIES = [
    ("FreeRTOS 堆\n(48KB heap_4)", 49152,  "#2196F3"),
    ("LVGL 内存池\n(40KB)",         40960,  "#4CAF50"),
    ("UI 帧缓冲\n(~13KB)",         13158,  "#FF9800"),
    ("FTL / 存储\n(~11KB)",        12416,  "#9C27B0"),
    ("网络 / CH395F\n(~7KB)",       7110,   "#F44336"),
    ("内核 / CMSIS\n(~3KB)",        3024,   "#00BCD4"),
    ("HAL 驱动\n(~1.4KB)",         1400,   "#607D8B"),
    ("其他静态\n(~0.5KB)",          498,    "#795548"),
    ("SRAM 空闲",                   SRAM_FREE, "#E0E0E0"),
]

# ── Flash breakdown (1MB = 1,048,576 B) ─────────────────────────────────────
FLASH_TOTAL = 1048576
FLASH_CODE  = 219976
FLASH_RO    = 11496
FLASH_RW    = 528
FLASH_USED  = FLASH_CODE + FLASH_RO + FLASH_RW
FLASH_FREE  = FLASH_TOTAL - FLASH_USED

FLASH_CATEGORIES = [
    ("代码段 (Code)",    FLASH_CODE,  "#2196F3"),
    ("只读数据 (RO)",    FLASH_RO,    "#4CAF50"),
    ("初始化数据 (RW)",  FLASH_RW,    "#FF9800"),
    ("Flash 空闲",       FLASH_FREE,  "#E0E0E0"),
]

# Flash code sub-breakdown (top categories from map analysis)
FLASH_CODE_SUB = [
    ("LVGL 库",        69364,  "#4CAF50"),
    ("应用代码/UI",    42540,  "#2196F3"),
    ("C 库 (microlib)", 4230,  "#795548"),
    ("FatFS + dhara",  11540,  "#9C27B0"),
    ("CH395F 网络",     8776,  "#F44336"),
    ("HAL 驱动",       9400,  "#607D8B"),
    ("FreeRTOS",       4240,  "#00BCD4"),
    ("其他",          69886,  "#E0E0E0"),  # remainder
]

# ── Task stack budgets ─────────────────────────────────────────────────────
TASK_STACKS = [
    ("defaultTask",    2048,  "Normal",    "FatFS 挂载、LED"),
    ("netTask",        4096,  "Normal",    "CH395F 轮询、消息处理"),
    ("ftpTask",        4096,  "Normal",    "lftpd FTP 服务"),
    ("testTask",       4096,  "BelowNormal", "测试套件"),
    ("uiTask",        16384,  "BelowNormal", "LVGL UI 主循环"),
    ("timeSyncTask",   2048,  "Normal",    "TCP 时间同步"),
]

HEAP_FREE = 49152 - (2048+4096+4096+4096+16384+2048+512+1024+80*8+448+24)
# = 49152 - 35304 = 13848 (approximate; stacks+TCBs+queue+eventgroup)

# ---------------------------------------------------------------------------
# SVG helpers
# ---------------------------------------------------------------------------

def _text_width(text, char_w=7.2):
    """Rough monospace text width estimate (px)."""
    return len(text) * char_w


def _wrap(text, max_chars=12):
    """Wrap label to multiple <tspan> lines."""
    words = text.split("\n") if "\n" in text else text.split()
    lines = []
    current = ""
    for w in words:
        test = (current + " " + w).strip() if current else w
        if len(test) > max_chars and current:
            lines.append(current)
            current = w
        else:
            current = test
    if current:
        lines.append(current)
    return lines


def svg_header(w, h):
    return (
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {w} {h}" '
        f'width="{w}" height="{h}" style="font-family:Consolas,\'Courier New\',monospace; '
        f'font-size:12px; background:#fff">\n'
        '<defs>\n'
        '<style>\n'
        '.title { font-size:14px; font-weight:bold; }\n'
        '.label { font-size:11px; }\n'
        '.legend-text { font-size:10px; }\n'
        '.bar-seg { stroke:#fff; stroke-width:0.5; }\n'
        '</style>\n'
        '</defs>\n'
    )


def svg_footer():
    return '</svg>\n'


def _draw_bar_segment(x, y, w, h, fill, label_inside, label_outside="",
                       label_font=10, min_inside_w=40):
    """Return SVG fragment for a bar segment with label."""
    parts = []
    parts.append(f'  <rect class="bar-seg" x="{x}" y="{y}" width="{w:.1f}" height="{h}" '
                 f'fill="{fill}" rx="2"/>\n')
    if w > min_inside_w and label_inside:
        # Center text inside bar
        lines = _wrap(label_inside, max_chars=int(w / 7))
        ty = y + h / 2 - (len(lines) - 1) * 6
        for line in lines:
            lx = x + w / 2
            parts.append(f'  <text x="{lx:.1f}" y="{ty:.1f}" text-anchor="middle" '
                         f'dominant-baseline="central" class="label" fill="#000">{line}</text>\n')
            ty += 12
    elif label_outside:
        lx = x + w + 4
        ty = y + h / 2
        parts.append(f'  <text x="{lx:.1f}" y="{ty:.1f}" text-anchor="start" '
                     f'dominant-baseline="central" class="legend-text" fill="#333">{label_outside}</text>\n')
    return "".join(parts)


# ---------------------------------------------------------------------------
# Diagram generators
# ---------------------------------------------------------------------------

def gen_ram_svg():
    """内存布局_RAM.svg — horizontal stacked bar with legend."""
    W, H = 820, 320
    BAR_X, BAR_Y, BAR_H = 60, 50, 50
    SCALE = (W - 80) / SRAM_TOTAL

    svg = svg_header(W, H)
    svg += f'<text x="{W/2}" y="24" text-anchor="middle" class="title">SRAM (128KB) 使用布局</text>\n'

    x = BAR_X
    for label, size, color in RAM_CATEGORIES:
        w = size * SCALE
        if w < 1:
            continue
        inside = label.replace("\n", " ") if w > 30 else ""
        outside = "" if w > 30 else label.replace("\n", " ")
        svg += _draw_bar_segment(x, BAR_Y, w, BAR_H, color, inside, outside)
        x += w

    # Legend
    lx, ly = 60, 130
    for label, size, color in RAM_CATEGORIES:
        pct = size / SRAM_TOTAL * 100
        kb = size / 1024
        display = label.replace("\n", " ")
        svg += f'  <rect x="{lx}" y="{ly}" width="12" height="12" fill="{color}" rx="2"/>\n'
        svg += f'  <text x="{lx+16}" y="{ly+10}" class="legend-text" fill="#333">'
        svg += f'{display}: {kb:.1f}KB ({pct:.1f}%)</text>\n'
        ly += 16
        if ly > H - 20:
            lx += 260
            ly = 130

    # Summary text
    svg += f'<text x="{BAR_X}" y="{BAR_Y+BAR_H+20}" class="label" fill="#333">'
    svg += f'已用: {SRAM_USED/1024:.1f}KB ({SRAM_USED/SRAM_TOTAL*100:.1f}%)'
    svg += f'　空闲: {SRAM_FREE/1024:.1f}KB ({SRAM_FREE/SRAM_TOTAL*100:.1f}%)'
    svg += f'　总量: 128KB</text>\n'

    svg += svg_footer()
    return svg


def gen_flash_svg():
    """内存布局_FLASH.svg — bar + code sub-breakdown."""
    W, H = 820, 240
    BAR_X, BAR_Y, BAR_H = 60, 50, 40
    SCALE = (W - 80) / FLASH_TOTAL

    svg = svg_header(W, H)
    svg += f'<text x="{W/2}" y="24" text-anchor="middle" class="title">Flash (1MB) 使用布局</text>\n'

    x = BAR_X
    for label, size, color in FLASH_CATEGORIES:
        w = size * SCALE
        if w < 1:
            continue
        inside = f"{label} {size/1024:.1f}KB" if w > 60 else ""
        outside = "" if w > 60 else f"{label} {size/1024:.1f}KB"
        svg += _draw_bar_segment(x, BAR_Y, w, BAR_H, color, inside, outside)
        x += w

    # Summary below bar
    svg += f'<text x="{BAR_X}" y="{BAR_Y+BAR_H+18}" class="label" fill="#333">'
    svg += f'已用: {FLASH_USED/1024:.1f}KB ({FLASH_USED/FLASH_TOTAL*100:.1f}%)'
    svg += f'　空闲: {FLASH_FREE/1024:.1f}KB ({FLASH_FREE/FLASH_TOTAL*100:.1f}%)'
    svg += f'　总量: 1024KB</text>\n'

    # Code sub-breakdown mini bar
    Y2 = 130
    svg += f'<text x="{BAR_X}" y="{Y2-6}" class="label" fill="#555">代码段明细 ({FLASH_CODE/1024:.1f}KB)</text>\n'
    SCALE2 = (W - 80) / FLASH_CODE
    x2 = BAR_X
    for label, size, color in FLASH_CODE_SUB:
        w = size * SCALE2
        if w < 1:
            continue
        inside = f"{label} {size/1024:.1f}KB" if w > 50 else ""
        outside = "" if w > 50 else f"{label} {size/1024:.1f}KB"
        svg += _draw_bar_segment(x2, Y2, w, 24, color, inside, outside)
        x2 += w

    svg += svg_footer()
    return svg


def gen_taskstack_svg():
    """任务栈预算.svg — horizontal bars for each task's stack allocation."""
    W, H = 740, 340
    BAR_X = 180
    MAX_STACK = 16384
    SCALE = (W - BAR_X - 40) / MAX_STACK
    ROW_H = 30
    Y_START = 60

    svg = svg_header(W, H)
    svg += f'<text x="{W/2}" y="24" text-anchor="middle" class="title">FreeRTOS 任务栈分配 (48KB 堆池)</text>\n'

    # Header
    svg += f'<text x="{BAR_X}" y="{Y_START-6}" class="legend-text" fill="#555">分配栈大小 (Bytes)</text>\n'

    y = Y_START
    for name, stack, prio, desc in TASK_STACKS:
        w = stack * SCALE
        color = "#2196F3" if "Normal" in prio else "#FF9800"

        # Task name label
        svg += f'  <text x="{BAR_X-4}" y="{y+ROW_H/2}" text-anchor="end" '
        svg += f'dominant-baseline="central" class="label" fill="#333">{name}</text>\n'

        # Bar
        svg += f'  <rect class="bar-seg" x="{BAR_X}" y="{y}" width="{w:.1f}" height="{ROW_H}" '
        svg += f'fill="{color}" rx="3"/>\n'

        # Size inside bar
        if w > 60:
            svg += f'  <text x="{BAR_X+w/2}" y="{y+ROW_H/2}" text-anchor="middle" '
            svg += f'dominant-baseline="central" class="label" fill="#fff">{stack}B ({stack/1024:.1f}KB)</text>\n'
        else:
            svg += f'  <text x="{BAR_X+w+4}" y="{y+ROW_H/2}" text-anchor="start" '
            svg += f'dominant-baseline="central" class="legend-text" fill="#333">{stack}B</text>\n'

        y += ROW_H + 4

    # FreeRTOS heap summary
    total_stacks = sum(s for _, s, _, _ in TASK_STACKS) + 512 + 1024  # app + idle + timer
    y += 4
    svg += f'<text x="{BAR_X}" y="{y}" class="label" fill="#555">'
    svg += f'堆总量: 49,152B (48KB)　'
    svg += f'任务栈合计: ~{total_stacks}B ({total_stacks/1024:.1f}KB)　'
    svg += f'空闲: ~{HEAP_FREE}B ({HEAP_FREE/1024:.1f}KB)</text>\n'

    # Color legend
    y += 18
    svg += f'  <rect x="{BAR_X}" y="{y}" width="10" height="10" fill="#2196F3" rx="2"/>\n'
    svg += f'  <text x="{BAR_X+14}" y="{y+8}" class="legend-text" fill="#333">Normal 优先级</text>\n'
    svg += f'  <rect x="{BAR_X+130}" y="{y}" width="10" height="10" fill="#FF9800" rx="2"/>\n'
    svg += f'  <text x="{BAR_X+144}" y="{y+8}" class="legend-text" fill="#333">BelowNormal 优先级</text>\n'

    svg += svg_footer()
    return svg


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="Generate memory layout SVG diagrams")
    parser.add_argument("--type", choices=["ram", "flash", "taskstack", "all"], default="all")
    args = parser.parse_args()

    out_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
                           "docs", "images")
    os.makedirs(out_dir, exist_ok=True)

    generators = {
        "ram": ("内存布局_RAM.svg", gen_ram_svg),
        "flash": ("内存布局_FLASH.svg", gen_flash_svg),
        "taskstack": ("任务栈预算.svg", gen_taskstack_svg),
    }

    if args.type == "all":
        targets = generators.values()
    else:
        targets = [generators[args.type]]

    for filename, gen_func in targets:
        path = os.path.join(out_dir, filename)
        svg = gen_func()
        with open(path, "w", encoding="utf-8") as f:
            f.write(svg)
        print(f"OK: {filename} ({len(svg)} bytes)")

    # XML self-check
    for filename, _ in targets:
        path = os.path.join(out_dir, filename)
        with open(path, "r", encoding="utf-8") as f:
            content = f.read()
        if not content.startswith("<svg "):
            print(f"FAIL: {filename} does not start with <svg", file=sys.stderr)
            sys.exit(1)
        if "</svg>" not in content:
            print(f"FAIL: {filename} missing </svg>", file=sys.stderr)
            sys.exit(1)
        if 'xmlns="http://www.w3.org/2000/svg"' not in content:
            print(f"FAIL: {filename} missing xmlns", file=sys.stderr)
            sys.exit(1)
    print("All SVG self-checks passed")


if __name__ == "__main__":
    main()