from __future__ import annotations

import asyncio
import inspect
import os
import sys
import threading
import time
from contextvars import ContextVar, Token
from dataclasses import dataclass, field
from types import FrameType


_ACTIVE_RECORDER: ContextVar["RequestProfiler | None"] = ContextVar(
    "waterfall_active_recorder",
    default=None,
)
_THREAD_PROFILE_STATE = threading.local()
_PACKAGE_ROOT = os.path.dirname(os.path.abspath(__file__))


@dataclass(slots=True)
class ProfileNode:
    name: str
    filename: str
    lineno: int
    started_ns: int
    ended_ns: int | None = None
    depth: int = 0
    children: list["ProfileNode"] = field(default_factory=list)

    @property
    def duration_ns(self) -> int:
        if self.ended_ns is None:
            return 0
        return max(0, self.ended_ns - self.started_ns)

    @property
    def self_time_ns(self) -> int:
        child_time = sum(child.duration_ns for child in self.children)
        return max(0, self.duration_ns - child_time)


@dataclass(slots=True)
class RequestTrace:
    method: str
    path: str
    status_code: int
    root: ProfileNode
    dropped_calls: int = 0

    @property
    def duration_ns(self) -> int:
        return self.root.duration_ns


def _current_task() -> asyncio.Task[object] | None:
    try:
        return asyncio.current_task()
    except RuntimeError:
        return None


def _install_dispatcher() -> None:
    depth = getattr(_THREAD_PROFILE_STATE, "depth", 0)
    if depth == 0:
        _THREAD_PROFILE_STATE.previous_profile = sys.getprofile()
        sys.setprofile(_dispatch_profile_event)
    _THREAD_PROFILE_STATE.depth = depth + 1


def _uninstall_dispatcher() -> None:
    depth = getattr(_THREAD_PROFILE_STATE, "depth", 0)
    if depth <= 1:
        sys.setprofile(
            getattr(_THREAD_PROFILE_STATE, "previous_profile", None)
        )
        _THREAD_PROFILE_STATE.depth = 0
        _THREAD_PROFILE_STATE.previous_profile = None
        return
    _THREAD_PROFILE_STATE.depth = depth - 1


def _dispatch_profile_event(frame: FrameType, event: str, arg: object) -> None:
    recorder = _ACTIVE_RECORDER.get()
    if recorder is not None:
        recorder.handle_event(frame, event, arg)

    previous_profile = getattr(_THREAD_PROFILE_STATE, "previous_profile", None)
    if previous_profile is not None:
        previous_profile(frame, event, arg)


def _format_frame_name(frame: FrameType) -> tuple[str, str, int]:
    code = frame.f_code
    filename = os.path.relpath(code.co_filename, os.getcwd())
    function_name = (
        code.co_qualname if hasattr(code, "co_qualname") else code.co_name
    )
    return function_name, filename, code.co_firstlineno


class RequestProfiler:
    def __init__(
        self,
        method: str,
        path: str,
        *,
        max_nodes: int = 5000,
    ) -> None:
        started_ns = time.perf_counter_ns()
        self.trace = RequestTrace(
            method=method,
            path=path,
            status_code=500,
            root=ProfileNode(
                name=f"{method} {path}",
                filename="<request>",
                lineno=0,
                started_ns=started_ns,
                depth=0,
            ),
        )
        self.max_nodes = max_nodes
        self.task: asyncio.Task[object] | None = None
        self._frame_nodes: dict[int, ProfileNode] = {}
        self._token: Token[RequestProfiler | None] | None = None

    def start(self) -> None:
        self.task = _current_task()
        self.trace.root.started_ns = time.perf_counter_ns()
        self._token = _ACTIVE_RECORDER.set(self)
        _install_dispatcher()

    def stop(self, *, status_code: int) -> RequestTrace:
        ended_ns = time.perf_counter_ns()
        self.trace.status_code = status_code

        try:
            if self._token is not None:
                _ACTIVE_RECORDER.reset(self._token)
        finally:
            _uninstall_dispatcher()

        for node in self._frame_nodes.values():
            if node.ended_ns is None:
                node.ended_ns = ended_ns
        self._frame_nodes.clear()
        self.trace.root.ended_ns = ended_ns
        return self.trace

    def handle_event(self, frame: FrameType, event: str, arg: object) -> None:
        if event not in {"call", "return"}:
            return

        if not self._is_current_request_task():
            return

        if event == "call":
            self._on_call(frame)
            return

        self._on_return(frame, arg)

    def _is_current_request_task(self) -> bool:
        if self.task is None:
            return True
        return _current_task() is self.task

    def _on_call(self, frame: FrameType) -> None:
        if not self._should_capture(frame):
            return

        if id(frame) in self._frame_nodes:
            return

        if len(self._frame_nodes) >= self.max_nodes:
            self.trace.dropped_calls += 1
            return

        parent = self._find_parent(frame.f_back)
        name, filename, lineno = _format_frame_name(frame)
        node = ProfileNode(
            name=name,
            filename=filename,
            lineno=lineno,
            started_ns=time.perf_counter_ns(),
            depth=parent.depth + 1,
        )
        parent.children.append(node)
        self._frame_nodes[id(frame)] = node

    def _on_return(self, frame: FrameType, arg: object) -> None:
        if self._is_coroutine_suspension(frame, arg):
            return

        node = self._frame_nodes.pop(id(frame), None)
        if node is not None and node.ended_ns is None:
            node.ended_ns = time.perf_counter_ns()

    def _find_parent(self, frame: FrameType | None) -> ProfileNode:
        while frame is not None:
            node = self._frame_nodes.get(id(frame))
            if node is not None:
                return node
            frame = frame.f_back
        return self.trace.root

    def _should_capture(self, frame: FrameType) -> bool:
        filename = os.path.abspath(frame.f_code.co_filename)
        return not filename.startswith(_PACKAGE_ROOT)

    def _is_coroutine_suspension(
        self,
        frame: FrameType,
        arg: object,
    ) -> bool:
        if not frame.f_code.co_flags & inspect.CO_COROUTINE:
            return False

        return asyncio.isfuture(arg) and not arg.done()
