visual_export.py

"""Export a small FolderVisualizer scan result as a standalone SVG tree."""

from __future__ import annotations

import argparse
import json
import os
import re
import unicodedata
import xml.etree.ElementTree as ET
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from launcher import (
    JSON_SIZE_THRESHOLD,
    MAX_CHILDREN_THRESHOLD,
    NODE_COUNT_THRESHOLD,
)


PROJECT_DIR = Path(__file__).resolve().parent
DEFAULT_INPUT_PATH = PROJECT_DIR / "outputs" / "scan-result.json"
DEFAULT_OUTPUT_DIR = PROJECT_DIR / "outputs" / "exports"

NODE_WIDTH = 240
NODE_HEIGHT = 40
COLUMN_GAP = 96
DEPTH_STEP = NODE_WIDTH + COLUMN_GAP
ROW_GAP = 52
CANVAS_PADDING = 40
MIN_CANVAS_HEIGHT = 520
TEXT_MAX_UNITS = 26
FILE_STEM_SUFFIX_LENGTH = 5

SVG_NAMESPACE = "http://www.w3.org/2000/svg"
FONT_FAMILY = "Segoe UI, Microsoft YaHei, Noto Sans CJK SC, sans-serif"
INVALID_FILENAME_CHARACTERS = re.compile(r'[<>:"/\\|?*\x00-\x1f]')

ET.register_namespace("", SVG_NAMESPACE)


class ExportError(RuntimeError):
    """A user-facing SVG export error."""


@dataclass
class NodePosition:
    node: dict[str, Any]
    depth: int
    parent_index: int | None
    x: float
    y: float = 0
    child_indices: list[int] = field(default_factory=list)


@dataclass(frozen=True)
class TreeStatistics:
    file_count: int
    directory_count: int
    node_count: int
    max_children: int
    max_depth: int
    json_size: int


@dataclass(frozen=True)
class ExportResult:
    output_path: Path
    width: int
    height: int
    file_count: int
    directory_count: int
    node_count: int
    max_depth: int
    svg_size: int


def _svg_tag(name: str) -> str:
    return f"{{{SVG_NAMESPACE}}}{name}"


def _required_nonnegative_int(data: dict[str, Any], key: str) -> int:
    value = data.get(key)
    if not isinstance(value, int) or isinstance(value, bool) or value < 0:
        raise ExportError(f'JSON 字段 "{key}" 必须是非负整数')
    return value


def load_scan_result(
    path: Path = DEFAULT_INPUT_PATH,
) -> tuple[dict[str, Any], TreeStatistics]:
    """Load and fully validate the tree before SVG generation."""

    path = path.resolve()
    try:
        json_size = path.stat().st_size
    except FileNotFoundError as error:
        raise ExportError(f"找不到扫描结果：{path}") from error
    except OSError as error:
        raise ExportError(f"无法读取扫描结果信息：{error}") from error

    try:
        with path.open("r", encoding="utf-8") as input_file:
            data = json.load(input_file)
    except json.JSONDecodeError as error:
        raise ExportError(f"scan-result.json 格式损坏：{error}") from error
    except OSError as error:
        raise ExportError(f"无法读取 scan-result.json：{error}") from error

    if not isinstance(data, dict):
        raise ExportError("scan-result.json 顶层必须是 JSON 对象")
    root = data.get("root")
    if not isinstance(root, dict):
        raise ExportError("scan-result.json 缺少有效的 root 节点")
    if root.get("type") != "directory":
        raise ExportError("扫描结果的根节点必须是目录")

    file_count = _required_nonnegative_int(data, "file_count")
    directory_count = _required_nonnegative_int(data, "directory_count")
    total_size = _required_nonnegative_int(data, "total_size")
    if _required_nonnegative_int(root, "size") != total_size:
        raise ExportError("根节点大小与 total_size 不一致")

    actual_files = 0
    actual_directories = 0
    max_children = 0
    max_depth = 0
    stack: list[tuple[dict[str, Any], int]] = [(root, 0)]
    while stack:
        node, depth = stack.pop()
        max_depth = max(max_depth, depth)
        name = node.get("name")
        path_value = node.get("path")
        node_type = node.get("type")
        children = node.get("children", [])

        if not isinstance(name, str):
            raise ExportError(f"深度 {depth} 的节点缺少有效名称")
        if not isinstance(path_value, str):
            raise ExportError(f'节点 "{name}" 缺少有效路径')
        if not isinstance(children, list) or not all(
            isinstance(child, dict) for child in children
        ):
            raise ExportError(f'节点 "{path_value}" 的 children 必须是节点列表')

        if node_type == "file":
            actual_files += 1
            if children:
                raise ExportError(f'文件节点 "{path_value}" 不能包含子节点')
        elif node_type == "directory":
            actual_directories += 1
            max_children = max(max_children, len(children))
        else:
            raise ExportError(
                f'节点 "{path_value}" 的 type 必须是 file 或 directory'
            )

        for child in reversed(children):
            stack.append((child, depth + 1))

    if actual_files != file_count:
        raise ExportError(
            f"文件数量不一致：JSON 为 {file_count}，树中实际为 {actual_files}"
        )
    if actual_directories != directory_count:
        raise ExportError(
            "目录数量不一致："
            f"JSON 为 {directory_count}，树中实际为 {actual_directories}"
        )

    statistics = TreeStatistics(
        file_count=file_count,
        directory_count=directory_count,
        node_count=file_count + directory_count,
        max_children=max_children,
        max_depth=max_depth,
        json_size=json_size,
    )
    _require_small_tree(statistics)
    return data, statistics


def _require_small_tree(statistics: TreeStatistics) -> None:
    reasons: list[str] = []
    if statistics.node_count > NODE_COUNT_THRESHOLD:
        reasons.append(
            f"节点数量 {statistics.node_count:,} > {NODE_COUNT_THRESHOLD:,}"
        )
    if statistics.max_children > MAX_CHILDREN_THRESHOLD:
        reasons.append(
            "最大直接子节点数量 "
            f"{statistics.max_children:,} > {MAX_CHILDREN_THRESHOLD:,}"
        )
    if statistics.json_size > JSON_SIZE_THRESHOLD:
        reasons.append(
            f"JSON 大小 {statistics.json_size / 1024**2:.2f} MiB > "
            f"{JSON_SIZE_THRESHOLD / 1024**2:.2f} MiB"
        )
    if reasons:
        details = "；".join(reasons)
        raise ExportError(
            "当前扫描结果属于大目录模式，阶段6.4.1不生成完整SVG："
            f"{details}"
        )


def layout_tree(root: dict[str, Any]) -> tuple[list[NodePosition], int, int]:
    """Lay out the tree iteratively using the current tidy-tree rules."""

    positions: list[NodePosition] = []
    next_leaf_y = 0
    max_depth = 0
    max_y = 0.0
    tasks: list[tuple[str, dict[str, Any] | int, int | None, int]] = [
        ("enter", root, None, 0)
    ]

    while tasks:
        action, value, parent_index, depth = tasks.pop()
        if action == "enter":
            node = value
            if not isinstance(node, dict):
                raise ExportError("布局阶段遇到无效节点")
            position_index = len(positions)
            position = NodePosition(
                node=node,
                depth=depth,
                parent_index=parent_index,
                x=CANVAS_PADDING + depth * DEPTH_STEP,
            )
            positions.append(position)
            if parent_index is not None:
                positions[parent_index].child_indices.append(position_index)
            max_depth = max(max_depth, depth)

            tasks.append(("exit", position_index, parent_index, depth))
            children = node.get("children", [])
            for child in reversed(children):
                tasks.append(("enter", child, position_index, depth + 1))
        else:
            position_index = value
            if not isinstance(position_index, int):
                raise ExportError("布局阶段遇到无效位置")
            position = positions[position_index]
            if not position.child_indices:
                position.y = next_leaf_y
                next_leaf_y += ROW_GAP
            else:
                first_child = positions[position.child_indices[0]]
                last_child = positions[position.child_indices[-1]]
                position.y = (first_child.y + last_child.y) / 2
            max_y = max(max_y, position.y)

    y_offset = CANVAS_PADDING + NODE_HEIGHT / 2
    for position in positions:
        position.y += y_offset

    width = (
        CANVAS_PADDING * 2
        + (max_depth + 1) * NODE_WIDTH
        + max_depth * COLUMN_GAP
    )
    height = max(
        MIN_CANVAS_HEIGHT,
        int(max_y + NODE_HEIGHT + CANVAS_PADDING * 2),
    )
    return positions, width, height


def _text_units(value: str) -> int:
    units = 0
    for character in value:
        width = unicodedata.east_asian_width(character)
        units += 2 if width in ("F", "W", "A") else 1
    return units


def _take_prefix(value: str, maximum_units: int) -> str:
    if maximum_units <= 0:
        return ""
    result: list[str] = []
    used = 0
    for character in value:
        units = _text_units(character)
        if used + units > maximum_units:
            break
        result.append(character)
        used += units
    return "".join(result)


def _take_suffix_characters(value: str, count: int) -> str:
    return value[-count:] if count > 0 else ""


def _truncate_tail(value: str, maximum_units: int = TEXT_MAX_UNITS) -> str:
    if _text_units(value) <= maximum_units:
        return value
    return _take_prefix(value, maximum_units - 1) + "…"


def display_name(
    name: str, node_type: str, maximum_units: int = TEXT_MAX_UNITS
) -> str:
    """Truncate folders at the tail and files in the middle."""

    if _text_units(name) <= maximum_units:
        return name
    if node_type == "directory":
        return _truncate_tail(name, maximum_units)

    extension_start = name.rfind(".")
    if extension_start <= 0 or extension_start == len(name) - 1:
        return _truncate_tail(name, maximum_units)

    stem = name[:extension_start]
    extension = name[extension_start:]
    stem_suffix = _take_suffix_characters(stem, FILE_STEM_SUFFIX_LENGTH)
    required_suffix = stem_suffix + extension
    available_prefix_units = (
        maximum_units - 1 - _text_units(required_suffix)
    )
    prefix = _take_prefix(stem, max(0, available_prefix_units))

    if prefix:
        return prefix + "…" + required_suffix
    if _text_units(required_suffix) <= maximum_units - 1:
        return "…" + required_suffix

    # An unusually long extension is kept intact as required, even if it
    # exceeds the normal text budget.
    return "…" + extension


def _add_title(group: ET.Element, name: str, path: str) -> None:
    title = ET.SubElement(group, _svg_tag("title"))
    title.text = f"{name}\n{path}"


def _add_folder_icon(group: ET.Element) -> None:
    icon = ET.SubElement(
        group,
        _svg_tag("g"),
        {"transform": "translate(11 10)", "aria-hidden": "true"},
    )
    ET.SubElement(
        icon,
        _svg_tag("path"),
        {
            "d": "M 1 4 H 8 L 10.5 7 H 21 V 19 H 1 Z",
            "fill": "#5f91c9",
            "stroke": "#376fae",
            "stroke-width": "1.2",
            "stroke-linejoin": "round",
        },
    )
    ET.SubElement(
        icon,
        _svg_tag("line"),
        {
            "x1": "2.5",
            "y1": "9",
            "x2": "19.5",
            "y2": "9",
            "stroke": "#d8e9fb",
            "stroke-width": "1",
        },
    )


def _add_file_icon(group: ET.Element) -> None:
    icon = ET.SubElement(
        group,
        _svg_tag("g"),
        {"transform": "translate(13 9)", "aria-hidden": "true"},
    )
    ET.SubElement(
        icon,
        _svg_tag("path"),
        {
            "d": "M 1 1 H 11 L 17 7 V 22 H 1 Z",
            "fill": "#ffffff",
            "stroke": "#7b8798",
            "stroke-width": "1.3",
            "stroke-linejoin": "round",
        },
    )
    ET.SubElement(
        icon,
        _svg_tag("path"),
        {
            "d": "M 11 1 V 7 H 17",
            "fill": "#edf1f6",
            "stroke": "#7b8798",
            "stroke-width": "1.3",
            "stroke-linejoin": "round",
        },
    )
    for y_position in (11, 15, 19):
        ET.SubElement(
            icon,
            _svg_tag("line"),
            {
                "x1": "4",
                "y1": str(y_position),
                "x2": "13",
                "y2": str(y_position),
                "stroke": "#b2bcc9",
                "stroke-width": "1",
            },
        )


def _connector_path(
    parent: NodePosition,
    children: list[NodePosition],
) -> str:
    start_x = parent.x + NODE_WIDTH
    branch_x = start_x + COLUMN_GAP / 2
    if len(children) == 1:
        return f"M {start_x:g} {parent.y:g} H {children[0].x:g}"

    commands = [
        f"M {start_x:g} {parent.y:g} H {branch_x:g}",
        f"M {branch_x:g} {children[0].y:g} V {children[-1].y:g}",
    ]
    commands.extend(
        f"M {branch_x:g} {child.y:g} H {child.x:g}"
        for child in children
    )
    return " ".join(commands)


def build_svg(
    data: dict[str, Any], statistics: TreeStatistics
) -> tuple[ET.ElementTree, int, int]:
    """Create a script-free, resource-free SVG document."""

    root_node = data["root"]
    positions, width, height = layout_tree(root_node)
    svg = ET.Element(
        _svg_tag("svg"),
        {
            "version": "1.1",
            "viewBox": f"0 0 {width} {height}",
            "width": str(width),
            "height": str(height),
            "role": "img",
            "aria-label": (
                f"{root_node['name']} 文件树，"
                f"{statistics.directory_count} 个目录，"
                f"{statistics.file_count} 个文件"
            ),
        },
    )
    document_title = ET.SubElement(svg, _svg_tag("title"))
    document_title.text = f"FolderVisualizer：{root_node['name']}"
    ET.SubElement(
        svg,
        _svg_tag("rect"),
        {
            "x": "0",
            "y": "0",
            "width": str(width),
            "height": str(height),
            "fill": "#ffffff",
        },
    )

    connectors_group = ET.SubElement(
        svg,
        _svg_tag("g"),
        {"id": "connectors", "aria-hidden": "true"},
    )
    for position in positions:
        if not position.child_indices:
            continue
        child_positions = [positions[index] for index in position.child_indices]
        ET.SubElement(
            connectors_group,
            _svg_tag("path"),
            {
                "d": _connector_path(position, child_positions),
                "fill": "none",
                "stroke": "#c3ccd8",
                "stroke-width": "1",
                "stroke-linecap": "round",
                "stroke-linejoin": "round",
            },
        )

    nodes_group = ET.SubElement(svg, _svg_tag("g"), {"id": "nodes"})
    for index, position in enumerate(positions):
        node = position.node
        node_type = node["type"]
        is_root = index == 0
        group = ET.SubElement(
            nodes_group,
            _svg_tag("g"),
            {
                "transform": (
                    f"translate({position.x:g} "
                    f"{position.y - NODE_HEIGHT / 2:g})"
                ),
                "data-node-type": node_type,
            },
        )
        _add_title(group, node["name"], node["path"])

        if is_root:
            fill = "#e7f1ff"
            stroke = "#84aee0"
            text_color = "#164a84"
            font_weight = "600"
        elif node_type == "directory":
            fill = "#f3f8ff"
            stroke = "#cbdcf3"
            text_color = "#164a84"
            font_weight = "600"
        else:
            fill = "#ffffff"
            stroke = "#dce3ec"
            text_color = "#344054"
            font_weight = "400"

        ET.SubElement(
            group,
            _svg_tag("rect"),
            {
                "x": "0",
                "y": "0",
                "width": str(NODE_WIDTH),
                "height": str(NODE_HEIGHT),
                "rx": "8",
                "fill": fill,
                "stroke": stroke,
                "stroke-width": "1",
            },
        )
        if node_type == "directory":
            _add_folder_icon(group)
        else:
            _add_file_icon(group)

        text = ET.SubElement(
            group,
            _svg_tag("text"),
            {
                "x": "43",
                "y": "25",
                "fill": text_color,
                "font-family": FONT_FAMILY,
                "font-size": "14",
                "font-weight": font_weight,
            },
        )
        text.text = display_name(node["name"], node_type)

    return ET.ElementTree(svg), width, height


def _safe_filename(value: str) -> str:
    cleaned = INVALID_FILENAME_CHARACTERS.sub("_", value).strip(" .")
    return cleaned or "folder"


def default_output_path(
    data: dict[str, Any], output_dir: Path = DEFAULT_OUTPUT_DIR
) -> Path:
    root_name = str(data["root"].get("name", "folder"))
    return output_dir.resolve() / f"{_safe_filename(root_name)}-tree.svg"


def export_small_tree(
    input_path: Path = DEFAULT_INPUT_PATH,
    output_path: Path | None = None,
    output_dir: Path = DEFAULT_OUTPUT_DIR,
) -> ExportResult:
    """Load a small scan result and atomically publish its complete SVG tree."""

    data, statistics = load_scan_result(input_path)
    final_output_path = (
        output_path.resolve()
        if output_path is not None
        else default_output_path(data, output_dir)
    )
    final_output_path.parent.mkdir(parents=True, exist_ok=True)
    temporary_path = final_output_path.with_name(
        f"{final_output_path.name}.tmp"
    )
    temporary_path.unlink(missing_ok=True)

    try:
        tree, width, height = build_svg(data, statistics)
        ET.indent(tree, space="  ")
        tree.write(
            temporary_path,
            encoding="utf-8",
            xml_declaration=True,
            short_empty_elements=True,
        )
        os.replace(temporary_path, final_output_path)
    except (OSError, ValueError, TypeError) as error:
        temporary_path.unlink(missing_ok=True)
        raise ExportError(f"无法生成 SVG：{error}") from error

    return ExportResult(
        output_path=final_output_path,
        width=width,
        height=height,
        file_count=statistics.file_count,
        directory_count=statistics.directory_count,
        node_count=statistics.node_count,
        max_depth=statistics.max_depth,
        svg_size=final_output_path.stat().st_size,
    )


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="将小目录 scan-result.json 导出为完整横向 SVG 文件树"
    )
    parser.add_argument(
        "--input",
        type=Path,
        default=DEFAULT_INPUT_PATH,
        help=f"扫描结果 JSON（默认 {DEFAULT_INPUT_PATH}）",
    )
    output_group = parser.add_mutually_exclusive_group()
    output_group.add_argument(
        "--output",
        type=Path,
        help="完整输出文件路径；默认根据根目录名称自动生成",
    )
    output_group.add_argument(
        "--output-dir",
        type=Path,
        default=DEFAULT_OUTPUT_DIR,
        help=f"输出目录（默认 {DEFAULT_OUTPUT_DIR}）",
    )
    return parser.parse_args()


def main() -> int:
    args = _parse_args()
    try:
        result = export_small_tree(
            input_path=args.input,
            output_path=args.output,
            output_dir=args.output_dir,
        )
    except (ExportError, KeyboardInterrupt) as error:
        if isinstance(error, KeyboardInterrupt):
            print("\nSVG 导出已取消。")
            return 130
        print(f"错误: {error}")
        return 1

    print("小目录完整树 SVG 导出完成")
    print(f"输出文件: {result.output_path}")
    print(f"画布大小: {result.width:,} × {result.height:,}")
    print(f"节点总数: {result.node_count:,}")
    print(f"文件数量: {result.file_count:,}")
    print(f"目录数量: {result.directory_count:,}")
    print(f"最大深度: {result.max_depth}")
    print(f"SVG 大小: {result.svg_size:,} 字节")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())