"""
Time profiling — decorators and context managers for measuring execution time.

Usage:
    from sondeo import profile, timer

    @profile
    def my_function():
        ...

    with timer("block name"):
        ...

    async with timer("async block"):
        ...
"""

from __future__ import annotations

import asyncio
import cProfile
import functools
import inspect
import pstats
import threading
import time
from collections.abc import Callable
from dataclasses import dataclass, field
from io import StringIO
from types import TracebackType
from typing import Any, TypeVar, overload

from sondeo.profiler import logging as logging_mod

F = TypeVar("F", bound=Callable[..., Any])


@dataclass
class TimingResult:
    """Result from a timing measurement."""

    name: str
    elapsed_ms: float
    calls: int = 1
    started_at: float | None = None
    ended_at: float | None = None
    correlation_id: str | None = None
    children: list[TimingResult] = field(default_factory=list)
    thread_id: int | None = None
    thread_name: str | None = None
    task_id: int | None = None
    task_name: str | None = None
    is_async: bool = False
    contention_hint: str | None = None

    @property
    def elapsed_s(self) -> float:
        return self.elapsed_ms / 1000.0

    def __repr__(self) -> str:
        return f"TimingResult({self.name!r}, {self.elapsed_ms:.2f}ms)"


@dataclass
class AsyncTaskStat:
    """Aggregate timing stats for a single asyncio task."""

    task_id: int
    task_name: str
    total_ms: float
    calls: int
    slowest_ms: float
    functions: list[str]


# Global results store
_results: list[TimingResult] = []


def _current_thread_info() -> tuple[int, str]:
    thread = threading.current_thread()
    ident = thread.ident if thread.ident is not None else -1
    return ident, thread.name


def _current_task_info() -> tuple[int | None, str | None]:
    try:
        task = asyncio.current_task()
    except RuntimeError:
        return None, None

    if task is None:
        return None, None
    return id(task), task.get_name()


def _derive_contention_hint(result: TimingResult) -> str | None:
    if result.is_async and result.elapsed_ms >= 100.0:
        return "Slow coroutine detected (>100ms). Check for blocking I/O or CPU work in async code."

    if not result.is_async and result.thread_name not in (None, "MainThread"):
        if result.elapsed_ms >= 150.0:
            return (
                "Long-running worker-thread span detected. "
                "If threads are waiting on locks, inspect contention hotspots."
            )

    return None


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


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


def get_async_task_stats(*, min_ms: float = 0.0, top_n: int | None = None) -> list[AsyncTaskStat]:
    """Aggregate async timings by task to highlight slow coroutines."""
    aggregate: dict[int, AsyncTaskStat] = {}

    for result in _results:
        if not result.is_async:
            continue
        if result.task_id is None or result.task_name is None:
            continue
        if result.elapsed_ms < min_ms:
            continue

        stat = aggregate.get(result.task_id)
        if stat is None:
            stat = AsyncTaskStat(
                task_id=result.task_id,
                task_name=result.task_name,
                total_ms=0.0,
                calls=0,
                slowest_ms=0.0,
                functions=[],
            )
            aggregate[result.task_id] = stat

        stat.total_ms += result.elapsed_ms
        stat.calls += 1
        stat.slowest_ms = max(stat.slowest_ms, result.elapsed_ms)
        if result.name not in stat.functions:
            stat.functions.append(result.name)

    stats = sorted(
        aggregate.values(),
        key=lambda item: (item.total_ms, item.slowest_ms),
        reverse=True,
    )
    if top_n is not None:
        return stats[:top_n]
    return stats


def get_concurrency_hints(*, slow_task_ms: float = 50.0) -> list[str]:
    """Return heuristic hints about thread contention and slow coroutine tasks."""
    hints: list[str] = []

    # Timing skew across threads can indicate contention or workload imbalance.
    per_thread_total: dict[int, float] = {}
    for result in _results:
        if result.thread_id is None:
            continue
        per_thread_total[result.thread_id] = (
            per_thread_total.get(result.thread_id, 0.0) + result.elapsed_ms
        )

    if len(per_thread_total) >= 2:
        totals = list(per_thread_total.values())
        max_total = max(totals)
        min_total = min(totals)
        if max_total >= 40.0 and min_total > 0.0 and (max_total / min_total) >= 4.0:
            hints.append(
                "Large runtime skew across threads detected. "
                "Potential lock contention or uneven work distribution."
            )

    async_stats = get_async_task_stats(min_ms=slow_task_ms, top_n=1)
    if async_stats:
        slowest = async_stats[0]
        hints.append(
            f"Slow coroutine task '{slowest.task_name}' observed (max {slowest.slowest_ms:.2f}ms)."
        )

    for result in _results:
        if result.contention_hint and result.contention_hint not in hints:
            hints.append(result.contention_hint)

    return hints


class _TimerContext:
    """Sync + async compatible timer context manager."""

    def __init__(self, name: str) -> None:
        self._name = name
        self._result: TimingResult | None = None
        self._scope: Any | None = None
        self._start_perf: float | None = None
        self._parent_cid: str | None = None
        self._started = False

    def _begin(self, *, is_async: bool) -> TimingResult:
        if self._started:
            raise RuntimeError("timer context already started")

        self._parent_cid = logging_mod.current_correlation_id()
        self._scope = logging_mod.correlation_scope()
        cid = self._scope.__enter__()
        thread_id, thread_name = _current_thread_info()
        task_id, task_name = _current_task_info()

        self._result = TimingResult(
            name=self._name,
            elapsed_ms=0.0,
            started_at=time.time(),
            correlation_id=cid,
            thread_id=thread_id,
            thread_name=thread_name,
            task_id=task_id,
            task_name=task_name,
            is_async=is_async,
        )
        self._start_perf = time.perf_counter()
        self._started = True

        logging_mod.emit_event(
            event_type="timer",
            level="debug",
            module=__name__,
            message=f"{self._name} start",
            correlation_id=cid,
            data={
                "phase": "start",
                "name": self._name,
                "mode": "async" if is_async else "sync",
                "thread_id": thread_id,
                "thread_name": thread_name,
                "task_id": task_id,
                "task_name": task_name,
                "parent_correlation_id": self._parent_cid,
            },
        )
        return self._result

    def _finish(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        if (
            not self._started
            or self._result is None
            or self._scope is None
            or self._start_perf is None
        ):
            return

        self._result.elapsed_ms = (time.perf_counter() - self._start_perf) * 1000.0
        self._result.ended_at = time.time()
        self._result.contention_hint = _derive_contention_hint(self._result)
        _results.append(self._result)

        logging_mod.emit_event(
            event_type="timer",
            level="info",
            module=__name__,
            message=f"{self._name} end",
            correlation_id=self._result.correlation_id,
            data={
                "phase": "end",
                "name": self._name,
                "elapsed_ms": round(self._result.elapsed_ms, 3),
                "mode": "async" if self._result.is_async else "sync",
                "thread_id": self._result.thread_id,
                "thread_name": self._result.thread_name,
                "task_id": self._result.task_id,
                "task_name": self._result.task_name,
                "contention_hint": self._result.contention_hint,
                "parent_correlation_id": self._parent_cid,
            },
        )

        self._scope.__exit__(exc_type, exc, tb)
        self._started = False

    def __enter__(self) -> TimingResult:
        return self._begin(is_async=False)

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        self._finish(exc_type, exc, tb)

    async def __aenter__(self) -> TimingResult:
        return self._begin(is_async=True)

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> bool:
        self._finish(exc_type, exc, tb)
        return False


def timer(name: str = "block") -> _TimerContext:
    """Context manager that measures execution time of sync or async blocks."""
    return _TimerContext(name)


def _finalize_profile(
    *,
    prof: cProfile.Profile | None,
    func: F,
    sort_by: str,
    top_n: int,
    print_stats: bool,
    started_at: float,
    start_perf: float,
    cid: str,
    parent_cid: str | None,
    is_async: bool,
) -> None:
    elapsed = (time.perf_counter() - start_perf) * 1000.0

    stats: pstats.Stats | None = None
    total_calls = 1
    stream = StringIO()
    if prof is not None:
        stats = pstats.Stats(prof, stream=stream)
        stats.sort_stats(sort_by)
        total_calls = int(getattr(stats, "total_calls", 0))

    thread_id, thread_name = _current_thread_info()
    task_id, task_name = _current_task_info()

    timing = TimingResult(
        name=func.__qualname__,
        elapsed_ms=elapsed,
        calls=total_calls,
        started_at=started_at,
        ended_at=time.time(),
        correlation_id=cid,
        thread_id=thread_id,
        thread_name=thread_name,
        task_id=task_id,
        task_name=task_name,
        is_async=is_async,
    )
    timing.contention_hint = _derive_contention_hint(timing)
    if stats is not None:
        timing._stats = stats  # type: ignore[attr-defined]
    _results.append(timing)

    logging_mod.emit_event(
        event_type="profile",
        level="info",
        module=func.__module__,
        message=f"{func.__qualname__} end",
        correlation_id=cid,
        data={
            "phase": "end",
            "function": func.__qualname__,
            "elapsed_ms": round(elapsed, 3),
            "calls": total_calls,
            "profiler": "cprofile" if stats is not None else "wallclock",
            "mode": "async" if is_async else "sync",
            "thread_id": thread_id,
            "thread_name": thread_name,
            "task_id": task_id,
            "task_name": task_name,
            "contention_hint": timing.contention_hint,
            "parent_correlation_id": parent_cid,
        },
    )

    if print_stats and stats is not None:
        stats.print_stats(top_n)
        print(stream.getvalue())


@overload
def profile(fn: F) -> F: ...


@overload
def profile(
    *,
    sort_by: str = "cumulative",
    top_n: int = 20,
    print_stats: bool = False,
) -> Callable[[F], F]: ...


def profile(
    fn: F | None = None,
    *,
    sort_by: str = "cumulative",
    top_n: int = 20,
    print_stats: bool = False,
) -> F | Callable[[F], F]:
    """Decorator that profiles sync and async functions with cProfile."""

    def decorator(func: F) -> F:
        if inspect.iscoroutinefunction(func):

            @functools.wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                parent_cid = logging_mod.current_correlation_id()
                with logging_mod.correlation_scope() as cid:
                    prof: cProfile.Profile | None = cProfile.Profile()
                    started_at = time.time()
                    start = time.perf_counter()
                    thread_id, thread_name = _current_thread_info()
                    task_id, task_name = _current_task_info()

                    logging_mod.emit_event(
                        event_type="profile",
                        level="debug",
                        module=func.__module__,
                        message=f"{func.__qualname__} start",
                        correlation_id=cid,
                        data={
                            "phase": "start",
                            "function": func.__qualname__,
                            "mode": "async",
                            "thread_id": thread_id,
                            "thread_name": thread_name,
                            "task_id": task_id,
                            "task_name": task_name,
                            "parent_correlation_id": parent_cid,
                        },
                    )

                    using_cprofile = True
                    try:
                        if prof is not None:
                            prof.enable()
                    except ValueError:
                        # cProfile cannot run concurrently in the same thread.
                        using_cprofile = False
                        prof = None

                    try:
                        result = await func(*args, **kwargs)
                    finally:
                        if using_cprofile and prof is not None:
                            prof.disable()
                        _finalize_profile(
                            prof=prof,
                            func=func,
                            sort_by=sort_by,
                            top_n=top_n,
                            print_stats=print_stats,
                            started_at=started_at,
                            start_perf=start,
                            cid=cid,
                            parent_cid=parent_cid,
                            is_async=True,
                        )
                return result

            return async_wrapper  # type: ignore[return-value]

        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            parent_cid = logging_mod.current_correlation_id()
            with logging_mod.correlation_scope() as cid:
                prof: cProfile.Profile | None = cProfile.Profile()
                started_at = time.time()
                start = time.perf_counter()
                thread_id, thread_name = _current_thread_info()
                task_id, task_name = _current_task_info()

                logging_mod.emit_event(
                    event_type="profile",
                    level="debug",
                    module=func.__module__,
                    message=f"{func.__qualname__} start",
                    correlation_id=cid,
                    data={
                        "phase": "start",
                        "function": func.__qualname__,
                        "mode": "sync",
                        "thread_id": thread_id,
                        "thread_name": thread_name,
                        "task_id": task_id,
                        "task_name": task_name,
                        "parent_correlation_id": parent_cid,
                    },
                )

                try:
                    using_cprofile = True
                    try:
                        if prof is not None:
                            prof.enable()
                    except ValueError:
                        using_cprofile = False
                        prof = None

                    result = func(*args, **kwargs)
                finally:
                    if using_cprofile and prof is not None:
                        prof.disable()
                    _finalize_profile(
                        prof=prof,
                        func=func,
                        sort_by=sort_by,
                        top_n=top_n,
                        print_stats=print_stats,
                        started_at=started_at,
                        start_perf=start,
                        cid=cid,
                        parent_cid=parent_cid,
                        is_async=False,
                    )
            return result

        return wrapper  # type: ignore[return-value]

    if fn is not None:
        return decorator(fn)
    return decorator
