from __future__ import annotations

import math
import pickle

import lz4.frame as gzip
import numpy as np
import numpy.typing as npt

from py3dtiles.tilers.point.node.node import Node
from py3dtiles.typing import ExtraFieldsDescription
from py3dtiles.utils import split_aabb


class NodeCatalog:
    """NodeCatalog is a store of Node objects.

    Using a NodeCatalog allows to only store a children names
    in nodes, instead of storing a full recursive structure.
    """

    def __init__(
        self,
        nodes: bytes,
        name: bytes,
        root_aabb: npt.NDArray[np.float64],
        root_spacing: float,
        include_rgb: bool,
        extra_fields: list[ExtraFieldsDescription],
    ) -> None:
        self.nodes: dict[bytes, Node] = {}
        self.root_aabb = root_aabb
        self.root_spacing = root_spacing
        self.node_bytes: dict[bytes, bytes] = {}
        self.include_rgb: bool = include_rgb
        self.extra_fields: list[ExtraFieldsDescription] = extra_fields
        self._load_from_store(name, nodes)

    def get_node(self, name: bytes) -> Node:
        """Returns the node mathing the given name"""
        if name not in self.nodes:
            spacing = self.root_spacing / math.pow(2, len(name))
            aabb = self.root_aabb
            for i in name:
                aabb = split_aabb(aabb, int(i))
            node = Node(name, aabb, spacing, self.include_rgb, self.extra_fields)
            self.nodes[name] = node
        else:
            node = self.nodes[name]
        return node

    def dump(self, name: bytes, max_depth: int) -> bytes:
        """Serialize the stored nodes to a bytes list"""
        node = self.nodes[name]
        if node.dirty:
            self.node_bytes[name] = node.save_to_bytes()

        if node.children is not None and max_depth > 0:
            for n in node.children:
                self.dump(n, max_depth - 1)

        return pickle.dumps(self.node_bytes)

    def _load_from_store(self, name: bytes, data: bytes) -> Node:
        if len(data) > 0:
            out = pickle.loads(gzip.decompress(data))
            for n in out:
                spacing = self.root_spacing / math.pow(2, len(n))
                aabb = self.root_aabb
                for i in n:
                    aabb = split_aabb(aabb, int(i))
                node = Node(n, aabb, spacing, self.include_rgb, self.extra_fields)
                node.load_from_bytes(out[n])
                self.node_bytes[n] = out[n]
                self.nodes[n] = node
        else:
            spacing = self.root_spacing / math.pow(2, len(name))
            aabb = self.root_aabb
            for i in name:
                aabb = split_aabb(aabb, int(i))
            node = Node(name, aabb, spacing, self.include_rgb, self.extra_fields)
            self.nodes[name] = node

        return self.nodes[name]
