"""
Implements the depth-based memory profiling strategy described in the design doc.
This profiler uses pympler's Asizer at different depth levels to ensure
accurate memory measurement, especially for objects with cycles and shared refs.
"""

import time
from dataclasses import dataclass, field
from typing import Any, List, Optional

try:
    from pympler.asizeof import asized
except ImportError:
    asized = None


@dataclass
class DepthProfilerConfig:
    """Configuration for the depth-based profiler."""

    max_depth: int = 20
    plateau_threshold: int = 3
    min_size_for_display: int = 1024
    show_shared_nodes: bool = True


@dataclass
class MemoryNode:
    """Represents a node in the memory usage tree, generated by the profiler."""

    name: str
    type_name: str
    total_size: int
    exclusive_size: int = 0
    children: List["MemoryNode"] = field(default_factory=list)
    depth: int = 0
    is_shared: bool = False


class DepthBasedProfiler:
    """
    Implements a memory profiling strategy that is perfectly aligned with Pympler's
    internal traversal and sizing logic. It uses the `pympler.asizeof.asized`
    function to generate a complete, sized tree of object references, and then
    translates that into a `MemoryNode` tree for reporting. This approach
    guarantees correctness by delegating all traversal and sizing to Pympler.
    """

    def __init__(self, config: Optional[DepthProfilerConfig] = None):
        if asized is None:
            raise ImportError("pympler is required for DepthBasedProfiler.")
        self.config = config or DepthProfilerConfig()
        self.log_identifier = f"[{type(self).__name__}]"

    def profile(self, root_obj: Any, name: str = "root") -> MemoryNode:
        """
        Profiles an object using pympler.asizeof.asized to build a memory tree.
        """
        print(f"{self.log_identifier} Starting asized-based profile for '{name}'...")
        start_time = time.time()

        sized_tree = asized(
            root_obj,
            limit=self.config.max_depth,
            detail=self.config.max_depth,
        )

        memory_tree = self._translate_asized_tree(sized_tree, name, 0)

        end_time = time.time()
        print(
            f"{self.log_identifier} Profiling for '{name}' completed in {end_time - start_time:.2f}s. Total size: {memory_tree.total_size} bytes."
        )
        return memory_tree

    def _translate_asized_tree(
        self, sized_node: Any, node_name: str, depth: int
    ) -> MemoryNode:
        """
        Recursively translates a `pympler.asizeof.Asized` node and its referents
        into a `MemoryNode` tree.
        """
        type_name = sized_node.name

        is_shared = sized_node.size == 0 and sized_node.flat > 0

        memory_node = MemoryNode(
            name=node_name,
            type_name=type_name,
            total_size=sized_node.size,
            exclusive_size=sized_node.flat,
            depth=depth,
            is_shared=is_shared,
        )

        if depth >= self.config.max_depth:
            return memory_node

        for child_sized_node in sized_node.refs:
            child_name = child_sized_node.name

            child_node = self._translate_asized_tree(
                child_sized_node,
                child_name,
                depth + 1,
            )

            if child_node.total_size > 0 or (
                self.config.show_shared_nodes and child_node.is_shared
            ):
                memory_node.children.append(child_node)

        return memory_node
