CTStruct/ctstruct/tree.py

"""
ctstruct/tree.py
CTTree - 트리 전체를 관리하는 컨테이너
"""

from __future__ import annotations

from typing import Any, Callable, Iterator, Optional

from .node import CTNode
from .types import NodeType


class CTTree:
    """
    CTStruct 트리 컨테이너.

    루트 노드를 통해 전체 트리를 관리하며,
    파싱·렌더링·변환 파이프라인의 중심 객체.

    Examples:
        >>> tree = CTTree()
        >>> sec = tree.add(NodeType.SECTION)
        >>> sec.append(CTNode(NodeType.HEADING, "제목", {"level": 1}))
        >>> sec.append(CTNode(NodeType.PARAGRAPH, "본문 내용입니다."))
        >>> print(tree.pretty())
    """

    def __init__(self, title: str = "", metadata: Optional[dict[str, Any]] = None) -> None:
        self.root = CTNode(NodeType.ROOT)
        if title:
            self.root.set_attr("title", title)
        if metadata:
            meta_node = CTNode(NodeType.METADATA, attributes=metadata)
            self.root.append(meta_node)

    # ── 빌더 API ─────────────────────────────────────────────

    def add(self, node_type: NodeType,
            content: Optional[str] = None,
            **attrs: Any) -> CTNode:
        """루트에 직접 자식 노드를 추가하고 반환."""
        node = CTNode(node_type, content, dict(attrs))
        self.root.append(node)
        return node

    @staticmethod
    def make_heading(text: str, level: int = 1) -> CTNode:
        return CTNode(NodeType.HEADING, text, {"level": level})

    @staticmethod
    def make_paragraph(text: str) -> CTNode:
        return CTNode(NodeType.PARAGRAPH, text)

    @staticmethod
    def make_code_block(code: str, language: str = "") -> CTNode:
        return CTNode(NodeType.CODE_BLOCK, code, {"language": language})

    @staticmethod
    def make_list(items: list[str], ordered: bool = False) -> CTNode:
        from .types import ListStyle
        style = ListStyle.ORDERED if ordered else ListStyle.UNORDERED
        lst = CTNode(NodeType.LIST, attributes={"style": style.value})
        for item in items:
            lst.append(CTNode(NodeType.LIST_ITEM, item))
        return lst

    @staticmethod
    def make_table(headers: list[str], rows: list[list[str]]) -> CTNode:
        table = CTNode(NodeType.TABLE)
        header_row = CTNode(NodeType.TABLE_ROW, attributes={"is_header": True})
        for h in headers:
            header_row.append(CTNode(NodeType.TABLE_CELL, h, {"is_header": True}))
        table.append(header_row)
        for row in rows:
            tr = CTNode(NodeType.TABLE_ROW)
            for cell in row:
                tr.append(CTNode(NodeType.TABLE_CELL, cell))
            table.append(tr)
        return table

    # ── 탐색 ─────────────────────────────────────────────────

    def find(self, node_type: NodeType) -> list[CTNode]:
        """트리 전체에서 특정 타입 노드 검색."""
        return self.root.find(node_type)

    def find_by_id(self, node_id: str) -> Optional[CTNode]:
        return self.root.find_by_id(node_id)

    def walk(self, order: str = "pre") -> Iterator[CTNode]:
        return self.root.walk(order)

    def headings(self) -> list[CTNode]:
        return self.find(NodeType.HEADING)

    def toc(self) -> list[dict[str, Any]]:
        """
        문서 목차(TOC) 반환.
        Returns:
            [{"level": int, "text": str, "node_id": str}, ...]
        """
        return [
            {
                "level": h.attr("level", 1),
                "text": h.content or "",
                "node_id": h.node_id,
            }
            for h in self.headings()
        ]

    # ── 변환 ─────────────────────────────────────────────────

    def filter(self, predicate: Callable[[CTNode], bool]) -> "CTTree":
        """
        조건에 맞는 노드만 남긴 **새로운** 트리 반환.
        원본 트리는 불변.
        """
        import copy
        new_tree = copy.deepcopy(self)
        nodes_to_remove = [
            n for n in new_tree.walk()
            if not predicate(n) and not n.is_root
        ]
        for n in nodes_to_remove:
            if n.parent:
                n.parent.remove(n)
        return new_tree

    def map(self, transform: Callable[[CTNode], None]) -> "CTTree":
        """
        모든 노드에 transform 함수를 적용한 **새로운** 트리 반환.
        """
        import copy
        new_tree = copy.deepcopy(self)
        for node in new_tree.walk():
            transform(node)
        return new_tree

    def merge(self, other: "CTTree") -> "CTTree":
        """다른 트리의 루트 자식들을 현재 트리에 병합 (in-place)."""
        import copy
        for child in other.root.children:
            self.root.append(copy.deepcopy(child))
        return self

    # ── 직렬화 ───────────────────────────────────────────────

    def to_dict(self) -> dict[str, Any]:
        """트리 전체를 Python dict 로 직렬화."""
        def _serialize(node: CTNode) -> dict:
            return {
                "id": node.node_id,
                "type": node.node_type.name,
                "content": node.content,
                "attributes": node.attributes,
                "children": [_serialize(c) for c in node.children],
            }
        return _serialize(self.root)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "CTTree":
        """dict 에서 트리 복원."""
        tree = cls()

        def _deserialize(d: dict) -> CTNode:
            node = CTNode(
                node_type=NodeType[d["type"]],
                content=d.get("content"),
                attributes=d.get("attributes", {}),
                node_id=d.get("id", ""),
            )
            for child_data in d.get("children", []):
                node.append(_deserialize(child_data))
            return node

        tree.root = _deserialize(data)
        return tree

    # ── 표현 ─────────────────────────────────────────────────

    @property
    def node_count(self) -> int:
        return sum(1 for _ in self.walk())

    @property
    def depth(self) -> int:
        return max((n.depth for n in self.walk()), default=0)

    def pretty(self) -> str:
        return self.root.pretty()

    def __repr__(self) -> str:
        return f"CTTree(nodes={self.node_count}, depth={self.depth})"

    def __len__(self) -> int:
        return len(self.root.children)