"""
Profiling utilities.
"""

import contextlib
import time
import types
import typing as T
from typing import override

import pandas as pd
import torch
from tensordict import TensorDict

from unipercept.types import Device

__all__ = [
    "LatencyAccumulator",
    "LatencyHandle",
    "LatencyProfile",
    "LatencyRecorder",
    "ProfileAccumulator",
    "latency_iter",
    "latency_profile",
    "profile",
    "profile_iter",
]

type Nanos = int
type LatencyMemory = T.MutableMapping[str, Nanos]

LATENCY_KEY_SEPARATOR: T.Final[str] = "."
# Backwards-compatible aliases for external callers
TimingsMemory = LatencyMemory
PROFILE_KEY_SEPARATOR = LATENCY_KEY_SEPARATOR


class LatencyProfile:
    """
    Context manager that profiles a block of code and stores the elapsed time.

    If the `dest` is `None` and the strict flag is set to `True`, an error is raised.
    Otherwise, the the null context manager is returned.
    """

    __slots__ = ("_start", "_time", "key", "memory")

    memory: LatencyMemory | None
    key: str
    _time: Nanos
    _start: Nanos | None

    @T.overload
    def __new__(
        cls,
        dest: None,
        key: str,
        /,
        *,
        strict: T.Literal[False],
    ) -> contextlib.AbstractContextManager[None]: ...

    @T.overload
    def __new__(
        cls,
        dest: None,
        key: str,
        /,
        *,
        strict: T.Literal[True],
    ) -> T.NoReturn: ...

    @T.overload
    def __new__(
        cls,
        dest: LatencyMemory,
        key: str,
        /,
        *,
        strict: bool = False,
    ) -> T.Self: ...

    def __new__(
        cls,
        dest: T.MutableMapping[str, Nanos] | None,
        key: str,
        /,
        *,
        strict: bool = False,
    ) -> contextlib.AbstractContextManager[None] | T.Self:
        """Return a profiling context that records nanosecond latency if enabled."""
        if dest is None:
            if strict:
                msg = f"No memory was passed and strict mode is enabled ({strict=})."
                raise RuntimeError(msg)
            return contextlib.nullcontext()

        # Construct the class
        self = super().__new__(cls)
        self.memory = dest
        self.key = key
        self._start = None

        return self

    @property
    def _running(self) -> bool:
        return self._start is not None

    @property
    def _delta(self) -> Nanos:
        if self._start is None:
            return 0
        return time.perf_counter_ns() - self._start

    def __enter__(self) -> "LatencyHandle":
        """Start timing and return a handle for nested measurements."""
        if self.memory is not None:
            self._start = time.perf_counter_ns()
        return LatencyHandle(self)

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: types.TracebackType | None,
    ) -> None:
        """Stop timing and persist the measured latency."""
        if self.memory is None:
            return
        if not self._running:
            return
        self.memory[self.key] = self._delta
        self._start = None


class LatencyHandle(T.MutableMapping[str, Nanos]):
    """Redirect all operations to the parent mapping, adding a prefixed key."""

    __slots__ = ("_ref",)

    def __init__(self, ref: LatencyProfile) -> None:
        self._ref = ref

    @property
    def _prefix(self) -> str:
        return self._ref.key + LATENCY_KEY_SEPARATOR

    @property
    def memory(self) -> LatencyMemory:
        """Return the backing latency memory, raising if unavailable."""
        mem = self._ref.memory
        if mem is None:
            msg = "Leaf is not attached to a profile context with memory."
            raise ValueError(msg)
        return mem

    def cancel(self) -> None:
        """Cancel the ongoing profiling context."""
        self._ref._start = None

    @override
    def __setitem__(self, key: str, value: Nanos) -> None:
        self.memory[self._prefix + key] = value

    @override
    def __getitem__(self, key: str) -> Nanos:
        return self.memory[self._prefix + key]

    @override
    def __delitem__(self, key: str) -> None:
        del self.memory[self._prefix + key]

    @override
    def __iter__(self) -> T.Iterator[str]:
        return iter(self.memory)

    @override
    def __len__(self) -> int:
        return len(self.memory)

    @override
    def keys(self) -> T.KeysView[str]:
        return self.memory.keys()

    @override
    def values(self) -> T.ValuesView[int]:
        return self.memory.values()

    @override
    def items(self) -> T.ItemsView[str, int]:
        return self.memory.items()


@T.overload
def latency_iter[R: T.Any](
    dest: LatencyMemory | None,
    key: str,
    src: T.Iterator[R] | T.Iterable[R],
    *,
    warmup: int,
    strict: T.Literal[False],
) -> T.Iterable[R]: ...


@T.overload
def latency_iter[R: T.Any](
    dest: None,
    key: str,
    src: T.Iterator[R] | T.Iterable[R],
    *,
    warmup: int,
    strict: T.Literal[True],
) -> T.NoReturn: ...


@T.overload
def latency_iter[R: T.Any](
    dest: LatencyMemory,
    key: str,
    src: T.Iterator[R] | T.Iterable[R],
    *,
    warmup: int,
    strict: bool = False,
) -> T.Iterable[R]: ...


def latency_iter[R: T.Any](
    dest: LatencyMemory | None,
    key: str,
    src: T.Iterator[R] | T.Iterable[R],
    *,
    warmup: int = 0,
    strict: bool = False,
) -> T.Iterable[R]:
    r"""Profile item retrieval latency for an iterable.

    Records nanosecond latency per item and yields the items unchanged.

    Parameters
    ----------
    timings : TimingsMemory | None
        The memory to store the timings in.
    source : Iterator[_R] | Iterble[_R]
        The source to iterate over. If it is not already an iterator, we turn it into
        one using `iter(source)`.
    key : str
        The key to store the timings under in the target memory.
    warmup : int
        The number of iterations to skip before recording the timings.
    """
    assert warmup >= 0, f"Warmup ({warmup=}) must be a non-negative integer."
    if dest is None:
        if strict:
            msg = f"No memory was passed and strict mode is enabled ({strict=})."
            raise RuntimeError(msg)
        yield from src
        return

    iterator = src if isinstance(src, T.Iterator) else iter(src)
    remaining_warmup = warmup

    while True:
        with LatencyProfile(dest, key, strict=True) as ph:
            try:
                item = next(iterator)
            except StopIteration:
                ph.cancel()
                break

            if remaining_warmup > 0:
                ph.cancel()
                remaining_warmup -= 1
        yield item


class ProfileRecord(T.NamedTuple):
    r"""A single long-format record used to store profiling data."""

    key: str
    time: int | float


class LatencyRecorder(T.MutableMapping[str, int]):
    """Mapping that records latencies and offers helper entrypoints."""

    __slots__ = ("records",)

    _KEY_TIME: T.Final[str] = "seconds"
    _KEY_NAME: T.Final[str] = "name"

    records: list[ProfileRecord]

    def __init__(self) -> None:
        self.records = []

    def profile(self, key: str, *, strict: bool = False) -> LatencyProfile:
        """Return a latency profiling context bound to this recorder."""
        return LatencyProfile(self, key, strict=strict)

    def iter(
        self,
        key: str,
        src: T.Iterator[T.Any] | T.Iterable[T.Any],
        *,
        warmup: int = 0,
        strict: bool = False,
    ) -> T.Iterable[T.Any]:
        """Profile item retrieval latency for an iterable, bound to this recorder."""
        return latency_iter(self, key, src, warmup=warmup, strict=strict)

    def run(
        self, warmup: int = 0, incomplete_ok: bool = True, enabled: bool = True
    ) -> T.Iterator[dict[str, int] | None]:
        """
        Return an iterator of dicts to which profiling data can be written.

        The first populated dict defines the expected keys for subsequent iterations.
        """
        if not enabled:
            while True:
                yield None
        skip = warmup
        while True:
            target = {}
            yield target
            if len(target.keys()) == 0:
                continue

            new_keys = frozenset(target.keys())
            cur_keys = frozenset(self.keys())
            if len(self.records) > 0 and new_keys != cur_keys and not incomplete_ok:
                missing = new_keys - cur_keys
                msg = (
                    f"Received {missing} keys that were not present in the previous "
                    "iteration."
                )
                raise ValueError(msg)
            if skip > 0:
                skip -= 1
            else:
                for k, v in target.items():
                    self[k] = v
            del target

    @override
    def __getitem__(self, key: str) -> float:
        return self.to_means()[key]

    @override
    def __setitem__(self, key: str, value: int) -> None:
        assert isinstance(value, int), "Value must be time in nanoseconds"
        rec = ProfileRecord(key, value)
        self.records.append(rec)

    @override
    def __delitem__(self, key: str) -> None:
        idxs_to_remove = [i for i, r in enumerate(self.records) if r.key == key]
        for i in reversed(idxs_to_remove):
            del self.records[i]

    @override
    def __iter__(self) -> T.Iterator[str]:
        return iter(self.to_means())

    @override
    def __len__(self) -> int:
        return len(self.records)

    def reset(self) -> None:
        """Clear all recorded latencies."""
        self.records.clear()

    def to_records(self, *, unit: bool = True) -> T.Iterable[ProfileRecord]:
        """Return records in either seconds (default) or raw nanoseconds."""
        return (
            ProfileRecord(r.key, r.time * 1e-9 if unit else r.time)
            for r in self.records
        )

    def to_dataframe(self, *, unit: bool = True) -> pd.DataFrame:
        """Return the profiling data as a pandas DataFrame."""
        return pd.DataFrame.from_records(
            self.to_records(unit=unit),
            columns=[self._KEY_NAME, self._KEY_TIME],
        )

    def to_tensordict(
        self, *, device: Device | None = None, unit: bool = True
    ) -> TensorDict:
        """Return the profiling data grouped in a TensorDict."""
        return TensorDict(
            {
                key: torch.tensor(
                    times,
                    device=device,
                    dtype=torch.double if unit else torch.long,
                    requires_grad=False,
                )
                for key, times in self.to_groups(unit=unit).items()
            },
            device=device,
            batch_size=[],
        )

    def to_summary(self, *, unit: bool = True) -> pd.DataFrame:
        """Return summarizing statistics for every key in the accumulator."""
        df = self.to_dataframe(unit=unit)

        return df.groupby(self._KEY_NAME)[self._KEY_TIME].agg(
            ["count", "sum", "mean", "std", "min", "max", "median"]
        )

    def to_means(self) -> dict[str, float]:
        """Return mean values per recorded key."""
        return {k: sum(v) / len(v) for k, v in self.to_groups().items()}

    def to_groups(self, *, unit: bool = True) -> dict[str, list[int | float]]:
        """Group the records by their key."""
        groups = {}
        for r in self.records:
            t = r.time * 1e-9 if unit else r.time
            groups.setdefault(r.key, []).append(t)
        return groups

    @override
    def items(self) -> T.ItemsView[str, float]:
        return self.to_means().items()

    @override
    def keys(self) -> T.AbstractSet[str]:
        return frozenset(r.key for r in self.records)

    @override
    def values(self) -> T.ValuesView[float]:
        return self.to_means().values()

    @override
    def __repr__(self) -> str:
        return f"{self.__class__.__name__}({', '.join(self.keys())})"
