from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable

from .profiler import ProfileNode, RequestTrace


@dataclass(slots=True)
class FlatNode:
    node: ProfileNode
    start_ms: float
    duration_ms: float
    percent_total: float
    depth: int = 0
    tree_prefix: str = ""


@dataclass(slots=True)
class AggregatedFlatNode:
    label: str
    tree_prefix: str
    depth: int
    start_ms: float
    duration_ms: float
    percent_total: float
    self_time_ms: float
    count: int


@dataclass(slots=True)
class DisplayNode:
    label: str
    tree_label: str
    start_ms: float
    duration_ms: float
    percent_total: float
    self_time_ms: float
    is_aggregate: bool = False
    count: int = 1


def _flatten(
    node: ProfileNode,
    root_started_ns: int,
    total_ns: int,
    ancestor_has_next: tuple[bool, ...] = (),
) -> list[FlatNode]:
    flat: list[FlatNode] = []
    duration_ms = node.duration_ns / 1_000_000
    start_ms = (node.started_ns - root_started_ns) / 1_000_000
    percent_total = (
        0.0 if total_ns == 0 else (node.duration_ns / total_ns) * 100
    )
    flat.append(
        FlatNode(
            node=node,
            start_ms=start_ms,
            duration_ms=duration_ms,
            percent_total=percent_total,
            depth=node.depth,
        )
    )
    child_count = len(node.children)
    for index, child in enumerate(node.children):
        is_last_child = index == child_count - 1
        child_tree_prefix = _tree_prefix(ancestor_has_next, is_last_child)
        child_nodes = _flatten(
            child,
            root_started_ns,
            total_ns,
            ancestor_has_next=ancestor_has_next + (not is_last_child,),
        )
        child_nodes[0].tree_prefix = child_tree_prefix
        flat.extend(child_nodes)
    return flat


def _render_bar(
    start_ms: float,
    duration_ms: float,
    total_ms: float,
    width: int,
) -> str:
    if width <= 0:
        return ""

    canvas = [" "] * width
    if total_ms <= 0:
        return "".join(canvas)

    start_col = min(width - 1, int((start_ms / total_ms) * width))
    span_col = max(1, int((duration_ms / total_ms) * width))
    end_col = min(width, start_col + span_col)
    for index in range(start_col, end_col):
        canvas[index] = "#"
    return "".join(canvas)


def _format_label(node: ProfileNode) -> str:
    if node.filename == "<request>":
        return node.name
    return f"{node.filename}:{node.lineno} {node.name}"


def _tree_prefix(
    ancestor_has_next: tuple[bool, ...],
    is_last_child: bool,
) -> str:
    parts: list[str] = []
    for has_next in ancestor_has_next:
        parts.append("|  " if has_next else "   ")
    parts.append("`- " if is_last_child else "|- ")
    return "".join(parts)


def _to_display_node(item: FlatNode) -> DisplayNode:
    label = _format_label(item.node)
    return DisplayNode(
        label=label,
        tree_label=f"{item.tree_prefix}{label}",
        start_ms=item.start_ms,
        duration_ms=item.duration_ms,
        percent_total=item.percent_total,
        self_time_ms=item.node.self_time_ns / 1_000_000,
    )


def _aggregate_tree_prefix(depth: int) -> str:
    if depth <= 0:
        return ""
    return "|  " * (depth - 1) + "+- "


def _aggregate_small_nodes(
    nodes: Iterable[FlatNode],
) -> list[AggregatedFlatNode]:
    grouped: dict[tuple[int, str], AggregatedFlatNode] = {}
    for item in nodes:
        label = _format_label(item.node)
        key = (item.depth, label)
        aggregated = grouped.get(key)
        if aggregated is None:
            grouped[key] = AggregatedFlatNode(
                label=label,
                tree_prefix=_aggregate_tree_prefix(item.depth),
                depth=item.depth,
                start_ms=item.start_ms,
                duration_ms=item.duration_ms,
                percent_total=item.percent_total,
                self_time_ms=item.node.self_time_ns / 1_000_000,
                count=1,
            )
            continue

        aggregated.start_ms = min(aggregated.start_ms, item.start_ms)
        aggregated.duration_ms += item.duration_ms
        aggregated.percent_total += item.percent_total
        aggregated.self_time_ms += item.node.self_time_ns / 1_000_000
        aggregated.count += 1

    return sorted(grouped.values(), key=lambda item: item.start_ms)


def _to_aggregated_display_node(item: AggregatedFlatNode) -> DisplayNode:
    label = f"{item.label} ({item.count} calls consolidated)"
    return DisplayNode(
        label=label,
        tree_label=f"{item.tree_prefix}{label}",
        start_ms=item.start_ms,
        duration_ms=item.duration_ms,
        percent_total=item.percent_total,
        self_time_ms=item.self_time_ms,
        is_aggregate=True,
        count=item.count,
    )


def render_trace(
    trace: RequestTrace,
    *,
    width: int = 60,
    hotspot_count: int = 5,
    collapse_below_percent: float = 0.1,
) -> str:
    total_ns = trace.duration_ns
    total_ms = total_ns / 1_000_000
    flat_nodes = _flatten(trace.root, trace.root.started_ns, total_ns)

    non_root_nodes = flat_nodes[1:]
    small_nodes = [
        item
        for item in non_root_nodes
        if item.percent_total < collapse_below_percent
    ]
    visible_nodes = [
        item
        for item in non_root_nodes
        if item.percent_total >= collapse_below_percent
    ]
    aggregated_small_nodes = _aggregate_small_nodes(small_nodes)

    display_nodes = [_to_display_node(item) for item in visible_nodes]
    display_nodes.extend(
        _to_aggregated_display_node(item) for item in aggregated_small_nodes
    )
    display_nodes.sort(key=lambda item: item.start_ms)

    hottest_node = max(
        display_nodes,
        key=lambda item: item.self_time_ms,
        default=None,
    )
    hottest_self = sorted(
        display_nodes,
        key=lambda item: item.self_time_ms,
        reverse=True,
    )[:hotspot_count]

    lines = [
        (
            f"Request {trace.status_code} {trace.method} "
            f"{trace.path} took {total_ms:.2f} ms"
        ),
        "Hotspots",
    ]

    if hottest_self:
        for item in hottest_self:
            lines.append(
                "  "
                f"{item.duration_ms:7.2f} ms total | "
                f"{item.self_time_ms:7.2f} ms self  "
                f"{item.tree_label}"
            )
    else:
        lines.append("  no Python frames were captured")

    lines.append("Waterfall")
    root_item = _to_display_node(flat_nodes[0])
    waterfall_nodes = [root_item, *display_nodes]
    for item in waterfall_nodes:
        marker = ""
        if hottest_node is item:
            marker = "  <<< hottest"

        bar = _render_bar(item.start_ms, item.duration_ms, total_ms, width)
        lines.append(
            f"{item.start_ms:8.2f} ms |{bar}|"
            f" {item.duration_ms:8.2f} ms"
            f" {item.percent_total:5.1f}% {item.tree_label}{marker}"
        )

    collapsed_count = sum(item.count for item in aggregated_small_nodes)
    if collapsed_count:
        lines.append(
            f"Collapsed low-impact calls: {collapsed_count} below "
            f"{collapse_below_percent:.1f}% each"
        )

    if trace.dropped_calls:
        lines.append(f"Dropped calls: {trace.dropped_calls}")

    return "\n".join(lines)
