from __future__ import annotations
from contextlib import contextmanager
import time
from typing import Any


def get_node_type(node: Any) -> str:
    """Returns the human-readable type or label of a node."""
    if hasattr(node, "bl_label") and node.bl_label:
        return str(node.bl_label)
    if hasattr(node, "bl_idname") and node.bl_idname:
        return str(node.bl_idname)
    return type(node).__name__


_current_profiler: NodeProfiler | None = None


def get_current_profiler() -> NodeProfiler | None:
    """Returns the currently active profiler instance, or None if inactive."""
    return _current_profiler


def set_current_profiler(profiler: NodeProfiler | None) -> None:
    """Sets the currently active profiler instance."""
    global _current_profiler
    _current_profiler = profiler


class NodeProfiler:
    """Tracks per-node execution counts and self-times."""

    def __init__(self) -> None:
        # Maps node_type -> [execution_count, total_self_time_seconds]
        self.stats: dict[str, list[float]] = {}
        self._stack: list[dict[str, Any]] = []

    @contextmanager
    def profile_node(self, node: Any):
        node_type = get_node_type(node)
        now = time.perf_counter()

        # Pause parent node's self-timer
        if self._stack:
            parent = self._stack[-1]
            parent["self_time"] += (now - parent["start_time"])

        frame = {
            "node_type": node_type,
            "start_time": now,
            "self_time": 0.0,
        }
        self._stack.append(frame)
        try:
            yield
        finally:
            end_time = time.perf_counter()
            popped = self._stack.pop()
            total_node_self_time = popped["self_time"] + (end_time - popped["start_time"])

            if node_type not in self.stats:
                self.stats[node_type] = [0.0, 0.0]
            self.stats[node_type][0] += 1.0
            self.stats[node_type][1] += total_node_self_time

            # Resume parent node's timer
            if self._stack:
                self._stack[-1]["start_time"] = time.perf_counter()

    def format_table(self) -> str:
        headers = ["node_type", "execution_count", "avg_execution_time", "tot_execution_time"]

        if not self.stats:
            col_widths = [20, 15, 18, 18]
            header_line = " | ".join(headers[i].ljust(col_widths[i]) for i in range(4))
            separator_line = "-+-".join("-" * col_widths[i] for i in range(4))
            return f"{header_line}\n{separator_line}\n(No nodes executed)"

        sorted_stats = sorted(self.stats.items(), key=lambda item: item[1][1], reverse=True)

        rows: list[list[str]] = []
        tot_count = 0
        tot_time_all = 0.0

        for node_type, (count, tot_time) in sorted_stats:
            icount = int(count)
            tot_count += icount
            tot_time_all += tot_time
            avg_time = tot_time / icount if icount > 0 else 0.0
            rows.append([
                node_type,
                str(icount),
                f"{avg_time:.6f}s",
                f"{tot_time:.6f}s",
            ])

        col_widths = [
            max(len(headers[i]), max(len(row[i]) for row in rows))
            for i in range(4)
        ]
        col_widths[0] = max(col_widths[0], len("Total"))
        col_widths[1] = max(col_widths[1], len(str(tot_count)))
        col_widths[3] = max(col_widths[3], len(f"{tot_time_all:.6f}s"))

        lines: list[str] = []
        header_line = " | ".join(headers[i].ljust(col_widths[i]) for i in range(4))
        separator_line = "-+-".join("-" * col_widths[i] for i in range(4))

        lines.append(header_line)
        lines.append(separator_line)
        for row in rows:
            lines.append(" | ".join([
                row[0].ljust(col_widths[0]),
                row[1].rjust(col_widths[1]),
                row[2].rjust(col_widths[2]),
                row[3].rjust(col_widths[3]),
            ]))
        lines.append(separator_line)
        total_row = [
            "Total".ljust(col_widths[0]),
            str(tot_count).rjust(col_widths[1]),
            "-".rjust(col_widths[2]),
            f"{tot_time_all:.6f}s".rjust(col_widths[3]),
        ]
        lines.append(" | ".join(total_row))

        return "\n".join(lines)