"""
Call tracing — capture function call trees with timing.

Usage:
    from sondeo import trace

    with trace() as t:
        my_function()

    t.print_tree()
"""

from __future__ import annotations

import sys
import threading
import time
from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass, field
from types import FrameType
from typing import Any

from sondeo.profiler import logging as logging_mod


@dataclass
class CallNode:
    """A node in the call tree."""

    name: str
    filename: str
    lineno: int
    elapsed_ms: float = 0.0
    children: list[CallNode] = field(default_factory=list)
    _start: float = field(default=0.0, repr=False)

    @property
    def total_children_ms(self) -> float:
        return sum(c.elapsed_ms for c in self.children)

    @property
    def self_ms(self) -> float:
        return self.elapsed_ms - self.total_children_ms

    def print_tree(self, indent: int = 0, min_ms: float = 0.1) -> None:
        """Print the call tree to stdout."""
        if self.elapsed_ms < min_ms:
            return
        prefix = "  " * indent
        loc = f"{self.filename}:{self.lineno}"
        print(f"{prefix}{self.name} ({self.elapsed_ms:.2f}ms) [{loc}]")
        for child in self.children:
            child.print_tree(indent + 1, min_ms)

    def to_dict(self) -> dict[str, Any]:
        """Convert to dict for JSON serialization."""
        return {
            "name": self.name,
            "file": self.filename,
            "line": self.lineno,
            "elapsed_ms": round(self.elapsed_ms, 3),
            "self_ms": round(self.self_ms, 3),
            "children": [c.to_dict() for c in self.children],
        }


@dataclass
class ThreadTrace:
    """Trace tree and metadata for one thread."""

    thread_id: int
    thread_name: str
    total_calls: int
    elapsed_ms: float
    children: list[CallNode] = field(default_factory=list)
    stalled_calls: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "thread_id": self.thread_id,
            "thread_name": self.thread_name,
            "total_calls": self.total_calls,
            "elapsed_ms": round(self.elapsed_ms, 3),
            "stalled_calls": list(self.stalled_calls),
            "children": [child.to_dict() for child in self.children],
        }


@dataclass
class TraceResult:
    """Result from a trace session."""

    root: CallNode = field(default_factory=lambda: CallNode("root", "", 0))
    total_calls: int = 0
    elapsed_ms: float = 0.0
    started_at: float | None = None
    ended_at: float | None = None
    correlation_id: str | None = None
    threads: list[ThreadTrace] = field(default_factory=list)
    contention_hints: list[str] = field(default_factory=list)

    def print_tree(self, min_ms: float = 0.1) -> None:
        """Print the full call tree."""
        for child in self.root.children:
            child.print_tree(min_ms=min_ms)

    def to_dict(self) -> dict[str, Any]:
        """Convert to dict for JSON serialization."""
        return {
            "total_calls": self.total_calls,
            "elapsed_ms": round(self.elapsed_ms, 3),
            "correlation_id": self.correlation_id,
            "contention_hints": list(self.contention_hints),
            "threads": [thread_trace.to_dict() for thread_trace in self.threads],
            "children": [c.to_dict() for c in self.root.children],
        }


@dataclass
class _ThreadState:
    thread_id: int
    thread_name: str
    root: CallNode
    stack: list[CallNode]
    total_calls: int = 0


_results: list[TraceResult] = []


def get_results() -> list[TraceResult]:
    """Get all collected trace results."""
    return list(_results)


def clear_results() -> None:
    """Clear all collected trace results."""
    _results.clear()


def _sum_elapsed(nodes: list[CallNode]) -> float:
    return sum(node.elapsed_ms for node in nodes)


def _build_contention_hints(
    *,
    states: list[_ThreadState],
    elapsed_ms: float,
    total_calls: int,
) -> list[str]:
    hints: list[str] = []

    if len(states) >= 2:
        per_thread_elapsed = [_sum_elapsed(state.root.children) for state in states]
        if per_thread_elapsed:
            max_elapsed = max(per_thread_elapsed)
            min_elapsed = min(per_thread_elapsed)
            if max_elapsed >= 40.0 and min_elapsed > 0.0 and (max_elapsed / min_elapsed) >= 4.0:
                hints.append(
                    "Call-time skew across threads is high. "
                    "Potential contention or imbalanced threaded workload."
                )

    stalled_threads = [state for state in states if len(state.stack) > 1]
    if stalled_threads:
        names = ", ".join(state.thread_name for state in stalled_threads)
        hints.append(
            "Some thread call stacks were still active when tracing stopped "
            f"({names}). Possible lock wait or blocked I/O."
        )

    if elapsed_ms >= 150.0 and total_calls <= max(1, len(states)):
        hints.append(
            "Trace duration was high with very low call throughput. "
            "Possible deadlock, lock contention, or blocking native call."
        )

    return hints


@contextmanager
def trace(
    *,
    filter_stdlib: bool = True,
    filter_prefix: str | None = None,
) -> Generator[TraceResult, None, None]:
    """Context manager that traces function calls within a block."""
    result = TraceResult()
    states: dict[int, _ThreadState] = {}
    state_lock = threading.Lock()

    def _thread_state() -> _ThreadState:
        thread = threading.current_thread()
        thread_id = thread.ident if thread.ident is not None else -1
        thread_name = thread.name
        state = states.get(thread_id)
        if state is not None:
            return state

        root = CallNode(name=f"thread:{thread_name}", filename="<thread>", lineno=0)
        state = _ThreadState(
            thread_id=thread_id,
            thread_name=thread_name,
            root=root,
            stack=[root],
            total_calls=0,
        )
        states[thread_id] = state
        return state

    def _trace(frame: FrameType, event: str, arg: Any) -> Any:
        filename = frame.f_code.co_filename

        if filter_stdlib and ("<" in filename or "lib/python" in filename.lower()):
            return _trace

        if filter_prefix and not filename.startswith(filter_prefix):
            return _trace

        with state_lock:
            state = _thread_state()

            if event == "call":
                node = CallNode(
                    name=frame.f_code.co_qualname
                    if hasattr(frame.f_code, "co_qualname")
                    else frame.f_code.co_name,
                    filename=filename,
                    lineno=frame.f_lineno,
                    _start=time.perf_counter(),
                )
                state.stack[-1].children.append(node)
                state.stack.append(node)
                state.total_calls += 1
                result.total_calls += 1
                return _trace

            if event == "return":
                if len(state.stack) > 1:
                    node = state.stack.pop()
                    node.elapsed_ms = (time.perf_counter() - node._start) * 1000.0

        return _trace

    parent_cid = logging_mod.current_correlation_id()
    with logging_mod.correlation_scope() as cid:
        result.correlation_id = cid
        result.started_at = time.time()
        start = time.perf_counter()

        logging_mod.emit_event(
            event_type="trace",
            level="debug",
            module=__name__,
            message="trace start",
            correlation_id=cid,
            data={
                "phase": "start",
                "parent_correlation_id": parent_cid,
            },
        )

        old_trace = sys.gettrace()
        old_thread_trace = threading.gettrace()
        sys.settrace(_trace)
        threading.settrace(_trace)
        try:
            yield result
        finally:
            sys.settrace(old_trace)
            threading.settrace(old_thread_trace)  # type: ignore[arg-type]
            result.elapsed_ms = (time.perf_counter() - start) * 1000.0
            result.ended_at = time.time()

            ordered_states = sorted(states.values(), key=lambda item: item.thread_name)
            result.threads = []
            result.root.children = []

            for state in ordered_states:
                elapsed_ms = _sum_elapsed(state.root.children)
                stalled = [node.name for node in state.stack[1:]]
                result.threads.append(
                    ThreadTrace(
                        thread_id=state.thread_id,
                        thread_name=state.thread_name,
                        total_calls=state.total_calls,
                        elapsed_ms=elapsed_ms,
                        children=list(state.root.children),
                        stalled_calls=stalled,
                    )
                )
                result.root.children.extend(state.root.children)

            result.contention_hints = _build_contention_hints(
                states=ordered_states,
                elapsed_ms=result.elapsed_ms,
                total_calls=result.total_calls,
            )
            _results.append(result)

            logging_mod.emit_event(
                event_type="trace",
                level="info",
                module=__name__,
                message="trace end",
                correlation_id=cid,
                data={
                    "phase": "end",
                    "total_calls": result.total_calls,
                    "elapsed_ms": round(result.elapsed_ms, 3),
                    "thread_count": len(result.threads),
                    "contention_hints": list(result.contention_hints),
                    "parent_correlation_id": parent_cid,
                },
            )
