# Source code for tabletree.src.lib.tabletree

    import json
    import os
    import traceback
    from collections import deque
    from typing import Any, Self, TypedDict, cast

    import msgpack

    from .row import RowBuilder, RowBuilderResource
    from .table import Table, TableBuilder

    class PartNames(TypedDict):
        partName: str
        partNumber: str
        DisplayPartName: str
        DisplayPartNumber: str

    class ModelInfo(TypedDict):
        part_names: PartNames

    class AppData(TypedDict):
        dsinspect: PartNames

    class TableTree(TypedDict):
        tables: dict[str, Table]
        formatVersion: int
        appdata: AppData

    class TableTreeBuilder:
        def __init__(self) -> None:
            """Create a tabletree."""
            self.root_builder: TableBuilder
            self.format_version: int = 3
            self.part_names: PartNames
            self._add_root_table_builder()

        def _add_root_table_builder(self):
            """Create a builder for the root table."""
            self.root_builder = TableBuilder().set_name("Top")
            self.root_builder.id = "root"

        def get_root_table_builder(self) -> TableBuilder:
            """Get the root table in tabletree."""
            return self.root_builder

        def set_model_info(self, model_info: ModelInfo) -> Self:
            """Set model info to tabletree.

            Args:
                model_info:
                    Model info of the geometry.
            """
            self.part_names = model_info["part_names"]
            return self

        def _build_all_tables(self) -> dict[str, Table]:
            """Build all the tables in tabletree."""
            tables: dict[str, Table] = {}
            builder_queue: deque[TableBuilder | RowBuilder] = deque(
                [self.root_builder]
            )
            while len(builder_queue) > 0:
                curr_builder = builder_queue.popleft()
                if isinstance(curr_builder, TableBuilder):
                    builder_id = curr_builder.id
                    tables[builder_id] = curr_builder.build()
                    for child in curr_builder.get_children():
                        builder_queue.append(child)

            return tables

        def build(self) -> TableTree:
            """Build the whole tabletree."""
            return TableTree(
                tables=self._build_all_tables(),
                formatVersion=self.format_version,
                appdata=AppData(dsinspect=self.part_names),
            )

    class TableTreeWriter:
        def __init__(self, tabletree: TableTree):
            """Create a writer for a tabletree.

            Args:
                tabletree: The tabletree to be written.
            """
            self.tabletree = tabletree

        def write(self, tabletree_outdir: str, blobs_outdir: str) -> bool:  # noqa: C901
            """Write the tabletree and blobs.

            Args:
                tabletree_outdir:
                    The output directory path for tabletree.
                blobs_outdir:
                    The output directory path for blobs.

            Returns:
                Whether files are successfully written or not.
            """

            def write_resource(res: RowBuilderResource, out_dir: str):
                out_path = os.path.join(out_dir, res["path"])
                out_parent_dir = os.path.realpath(
                    os.path.join(out_path, os.path.pardir)
                )
                os.makedirs(out_parent_dir, exist_ok=True)
                if res["format"] == "json":
                    with open(out_path, "w") as f:
                        json.dump(res["value"], f, allow_nan=False)
                elif res["format"] == "mp":
                    with open(out_path, "wb") as f:
                        f.write(msgpack.packb(res["value"], use_bin_type=True))

            def write_link_file(entry: str, path: str, out_dir: str):
                with (
                    open(os.path.join(out_dir, entry), "wb") as entry_buf,
                    open(path, "rb") as file_buf,
                ):
                    entry_buf.write(file_buf.read())

            try:
                os.makedirs(tabletree_outdir, exist_ok=True)
                os.makedirs(blobs_outdir, exist_ok=True)
                for table in self.tabletree["tables"].values():
                    if "_resources" in table:
                        for resource in (
                            cast(dict[str, Any], table).pop("_resources").values()
                        ):
                            write_resource(resource, blobs_outdir)

                    if "_link_file_paths" in table:
                        for entry, path in (
                            cast(dict[str, Any], table)
                            .pop("_link_file_paths")
                            .items()
                        ):
                            write_link_file(entry, path, tabletree_outdir)

                    for row in table["_rows"]:
                        if "_tables" in row:
                            continue

                        if "_resources" in row:
                            for resource in (
                                cast(dict[str, Any], row)
                                .pop("_resources")
                                .values()
                            ):
                                write_resource(resource, blobs_outdir)

                        if "_link_file_paths" in row:
                            for entry, path in (
                                cast(dict[str, Any], row)
                                .pop("_link_file_paths")
                                .items()
                            ):
                                write_link_file(entry, path, tabletree_outdir)

                with open(
                    os.path.join(tabletree_outdir, "tabletree.json"), "w"
                ) as f:
                    json.dump(self.tabletree, f, allow_nan=False)
                return True
            except Exception:
                traceback.print_exc()
                return False

        def serialize(self) -> str:
            """Serialize the whole tabletree into a string."""
            return json.dumps(self.tabletree, allow_nan=False)