"""Low-overhead, run-local performance instrumentation for simulations.

The profiling mode is intentionally read once when this module is imported.
``standard`` records low-frequency pipeline stages, while ``audit`` also wraps
hot RHS/property operations and computes exact-input reuse metrics.  With
profiling disabled, decorators return the original callable while classes are
being defined, so ordinary simulation calls do not pass through a wrapper.
"""

from __future__ import annotations

from collections.abc import Callable, Generator, Mapping
from contextlib import contextmanager
from contextvars import ContextVar, Token
from dataclasses import dataclass, field
from functools import wraps
import inspect
import os
import struct
from time import perf_counter_ns
from typing import Any, Literal, ParamSpec, TypeVar, cast


ProfileMode = Literal["off", "standard", "audit"]

_P = ParamSpec("_P")
_R = TypeVar("_R")
_MODE_RANK: Mapping[ProfileMode, int] = {"off": 0, "standard": 1, "audit": 2}


def _read_startup_mode() -> ProfileMode:
    raw_mode = os.getenv("SIMULATIONAPP_PROFILE", "off").strip().lower()
    aliases: dict[str, ProfileMode] = {
        "": "off",
        "0": "off",
        "false": "off",
        "no": "off",
        "off": "off",
        "1": "standard",
        "true": "standard",
        "yes": "standard",
        "on": "standard",
        "standard": "standard",
        "audit": "audit",
    }
    try:
        return aliases[raw_mode]
    except KeyError as exc:
        raise ValueError(
            "SIMULATIONAPP_PROFILE must be one of: off, standard, audit."
        ) from exc


PROFILE_MODE: ProfileMode = _read_startup_mode()


def _minimum_mode(value: str) -> ProfileMode:
    normalized = value.strip().lower()
    if normalized not in _MODE_RANK:
        raise ValueError("minimum_mode must be one of: off, standard, audit.")
    return cast(ProfileMode, normalized)


def _mode_enabled(minimum_mode: ProfileMode) -> bool:
    # ``off`` is an unconditional zero-wrapper mode, even if a caller passes
    # ``minimum_mode="off"`` by mistake.
    return (
        PROFILE_MODE != "off"
        and _MODE_RANK[PROFILE_MODE] >= _MODE_RANK[minimum_mode]
    )


@dataclass
class _TimingStats:
    calls: int = 0
    inclusive_ns: int = 0
    self_ns: int = 0
    max_ns: int = 0
    errors: int = 0

    def record(self, inclusive_ns: int, self_ns: int, error: bool) -> None:
        self.calls += 1
        self.inclusive_ns += inclusive_ns
        self.self_ns += self_ns
        self.max_ns = max(self.max_ns, inclusive_ns)
        if error:
            self.errors += 1

    def snapshot(self) -> dict[str, int]:
        return {
            "calls": self.calls,
            "inclusiveNs": self.inclusive_ns,
            "selfNs": self.self_ns,
            "maxNs": self.max_ns,
            "errors": self.errors,
        }


@dataclass
class _PropertyStats(_TimingStats):
    operation: str = ""
    layer: str = "semantic"
    medium: str = "unknown"
    exact_input_unique: int = 0
    exact_input_repeats: int = 0
    iteration_calls: int = 0
    iteration_total: int = 0
    iteration_max: int = 0
    iteration_converged: int = 0
    iteration_nonconverged: int = 0
    cache_lookups: int = 0
    cache_hits: int = 0
    cache_misses: int = 0

    def snapshot(self, *, audit: bool) -> dict[str, object]:
        result: dict[str, object] = super().snapshot()
        result.update(
            {
                "operation": self.operation,
                "layer": self.layer,
                "medium": self.medium,
                "cacheLookups": self.cache_lookups,
                "cacheHits": self.cache_hits,
                "cacheMisses": self.cache_misses,
            }
        )
        if audit:
            result.update(
                {
                    "exactInputUnique": self.exact_input_unique,
                    "exactInputRepeats": self.exact_input_repeats,
                    "iterationCalls": self.iteration_calls,
                    "iterationTotal": self.iteration_total,
                    "iterationMax": self.iteration_max,
                    "iterationConverged": self.iteration_converged,
                    "iterationNonconverged": self.iteration_nonconverged,
                }
            )
        return result


@dataclass
class _ActiveSpan:
    trace: PerformanceTrace
    name: str
    started_ns: int
    property_key: str | None = None
    property_operation: str | None = None
    property_outermost: bool = False
    child_ns: int = 0


@dataclass
class PerformanceTrace:
    """Mutable counters owned by exactly one :func:`profile_run` context."""

    mode: ProfileMode
    _phases: dict[str, _TimingStats] = field(default_factory=dict, repr=False)
    _properties: dict[str, _PropertyStats] = field(default_factory=dict, repr=False)
    _property_outermost_ns: int = field(default=0, repr=False)

    @property
    def enabled(self) -> bool:
        return self.mode != "off"

    def _property_stats(
        self,
        key: str,
        *,
        operation: str,
        layer: str,
        medium: str,
    ) -> _PropertyStats:
        stats = self._properties.get(key)
        if stats is None:
            stats = _PropertyStats(
                operation=operation,
                layer=layer,
                medium=medium,
            )
            self._properties[key] = stats
        return stats

    def _record_span(self, frame: _ActiveSpan, elapsed_ns: int, error: bool) -> None:
        self_ns = max(0, elapsed_ns - frame.child_ns)
        if frame.property_key is None:
            stats = self._phases.setdefault(frame.name, _TimingStats())
        else:
            layer, medium, operation = frame.property_key.split("|", 2)
            stats = self._property_stats(
                frame.property_key,
                operation=operation,
                layer=layer,
                medium=medium,
            )
            if frame.property_outermost:
                self._property_outermost_ns += elapsed_ns
        stats.record(elapsed_ns, self_ns, error)

    def _record_exact_input(
        self,
        key: str,
        *,
        operation: str,
        layer: str,
        medium: str,
        fingerprint: object,
    ) -> None:
        stats = self._property_stats(
            key,
            operation=operation,
            layer=layer,
            medium=medium,
        )
        shadow_key = (key, fingerprint)
        shadow = _PROPERTY_SHADOW.get()
        if shadow is None:
            # A trace normally installs its own set in ``profile_run``.  Keep
            # the ContextVar default immutable so no task can accidentally
            # share a process-global shadow set.
            shadow = set()
            _PROPERTY_SHADOW.set(shadow)
        if shadow_key in shadow:
            stats.exact_input_repeats += 1
        else:
            shadow.add(shadow_key)
            stats.exact_input_unique += 1

    def _record_iterations(
        self,
        key: str,
        *,
        operation: str,
        layer: str,
        medium: str,
        iterations: int,
        converged: bool,
    ) -> None:
        stats = self._property_stats(
            key,
            operation=operation,
            layer=layer,
            medium=medium,
        )
        iteration_count = max(0, int(iterations))
        stats.iteration_calls += 1
        stats.iteration_total += iteration_count
        stats.iteration_max = max(stats.iteration_max, iteration_count)
        if converged:
            stats.iteration_converged += 1
        else:
            stats.iteration_nonconverged += 1

    def _record_cache(
        self,
        key: str,
        *,
        operation: str,
        layer: str,
        medium: str,
        hits: int,
        misses: int,
    ) -> None:
        stats = self._property_stats(
            key,
            operation=operation,
            layer=layer,
            medium=medium,
        )
        hit_delta = max(0, hits)
        miss_delta = max(0, misses)
        stats.cache_hits += hit_delta
        stats.cache_misses += miss_delta
        stats.cache_lookups += hit_delta + miss_delta

    def snapshot(self) -> dict[str, object]:
        """Return a detached, JSON-serializable copy of all counters."""

        audit = self.mode == "audit"
        return {
            "mode": self.mode,
            "phases": {
                name: self._phases[name].snapshot()
                for name in sorted(self._phases)
            },
            "properties": {
                _public_property_key(key): self._properties[key].snapshot(audit=audit)
                for key in sorted(self._properties)
            },
            "propertyOutermostNs": self._property_outermost_ns,
        }


_CURRENT_TRACE: ContextVar[PerformanceTrace | None] = ContextVar(
    "simulation_performance_trace",
    default=None,
)
_ACTIVE_SPANS: ContextVar[tuple[_ActiveSpan, ...]] = ContextVar(
    "simulation_performance_spans",
    default=(),
)
_PROPERTY_SHADOW: ContextVar[set[object] | None] = ContextVar(
    "simulation_property_shadow",
    default=None,
)


def _public_property_key(key: str) -> str:
    layer, medium, operation = key.split("|", 2)
    return f"{layer}.{medium}.{operation}"


@contextmanager
def _tracked_span(
    trace: PerformanceTrace,
    name: str,
    *,
    property_key: str | None = None,
    property_operation: str | None = None,
    reset_property_shadow: bool = False,
) -> Generator[None, None, None]:
    stack = _ACTIVE_SPANS.get()
    property_outermost = property_key is not None and not any(
        item.property_key is not None for item in stack
    )
    frame = _ActiveSpan(
        trace=trace,
        name=name,
        started_ns=perf_counter_ns(),
        property_key=property_key,
        property_operation=property_operation,
        property_outermost=property_outermost,
    )
    stack_token = _ACTIVE_SPANS.set((*stack, frame))
    shadow_token: Token[set[object] | None] | None = None
    if reset_property_shadow and trace.mode == "audit":
        shadow_token = _PROPERTY_SHADOW.set(set())
    error = False
    try:
        yield
    except BaseException:
        error = True
        raise
    finally:
        elapsed_ns = max(0, perf_counter_ns() - frame.started_ns)
        _ACTIVE_SPANS.reset(stack_token)
        if shadow_token is not None:
            _PROPERTY_SHADOW.reset(shadow_token)
        if stack:
            stack[-1].child_ns += elapsed_ns
        trace._record_span(frame, elapsed_ns, error)


@contextmanager
def profile_run() -> Generator[PerformanceTrace, None, None]:
    """Create and bind an isolated trace for one simulation run.

    The yielded trace remains usable after the context exits, which lets the
    caller attach ``trace.snapshot()`` to a result without exposing live state.
    """

    trace = PerformanceTrace(mode=PROFILE_MODE)
    trace_token = _CURRENT_TRACE.set(trace)
    spans_token = _ACTIVE_SPANS.set(())
    shadow_token = _PROPERTY_SHADOW.set(set())
    try:
        if trace.enabled:
            with _tracked_span(trace, "simulation.total"):
                yield trace
        else:
            yield trace
    finally:
        _PROPERTY_SHADOW.reset(shadow_token)
        _ACTIVE_SPANS.reset(spans_token)
        _CURRENT_TRACE.reset(trace_token)


@contextmanager
def performance_span(
    name: str,
    minimum_mode: str = "standard",
    reset_property_shadow: bool = False,
) -> Generator[None, None, None]:
    """Time a block in the current run, or act as a no-op outside one."""

    minimum = _minimum_mode(minimum_mode)
    trace = _CURRENT_TRACE.get()
    if trace is None or not _mode_enabled(minimum):
        yield
        return
    with _tracked_span(
        trace,
        name,
        reset_property_shadow=reset_property_shadow,
    ):
        yield


def profile_phase(
    name: str,
    minimum_mode: str = "standard",
    reset_property_shadow: bool = False,
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
    """Decorate a simulation phase while preserving the off-mode callable."""

    minimum = _minimum_mode(minimum_mode)

    def decorate(function: Callable[_P, _R]) -> Callable[_P, _R]:
        if not _mode_enabled(minimum):
            return function

        if inspect.iscoroutinefunction(function):

            @wraps(function)
            async def async_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> Any:
                trace = _CURRENT_TRACE.get()
                if trace is None:
                    return await function(*args, **kwargs)
                with _tracked_span(
                    trace,
                    name,
                    reset_property_shadow=reset_property_shadow,
                ):
                    return await function(*args, **kwargs)

            return cast(Callable[_P, _R], async_wrapper)

        @wraps(function)
        def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
            trace = _CURRENT_TRACE.get()
            if trace is None:
                return function(*args, **kwargs)
            with _tracked_span(
                trace,
                name,
                reset_property_shadow=reset_property_shadow,
            ):
                return function(*args, **kwargs)

        return wrapper

    return decorate


def _medium_name(args: tuple[object, ...]) -> str:
    if not args:
        return "unknown"
    owner = args[0]
    configured_name = getattr(owner, "name", None)
    if isinstance(configured_name, str) and configured_name:
        return configured_name
    return type(owner).__name__


def _fingerprint(value: object) -> object:
    """Build a hashable, bit-exact token without retaining arbitrary objects."""

    if value is None or isinstance(value, (bool, int, str, bytes)):
        return (type(value).__name__, value)
    if isinstance(value, float):
        return ("float64", struct.pack("!d", value))
    if isinstance(value, tuple):
        return ("tuple", tuple(_fingerprint(item) for item in value))
    if isinstance(value, list):
        return ("list", tuple(_fingerprint(item) for item in value))
    if isinstance(value, Mapping):
        items = [(_fingerprint(key), _fingerprint(item)) for key, item in value.items()]
        items.sort(key=repr)
        return ("mapping", tuple(items))
    return (
        "object",
        type(value).__module__,
        type(value).__qualname__,
        id(value),
    )


def _input_fingerprint(
    signature: inspect.Signature | None,
    args: tuple[object, ...],
    kwargs: dict[str, object],
) -> object:
    if signature is not None:
        try:
            bound = signature.bind(*args, **kwargs)
            bound.apply_defaults()
            return tuple(
                (name, _fingerprint(value))
                for name, value in bound.arguments.items()
            )
        except TypeError:
            pass
    return (
        _fingerprint(args),
        tuple(sorted((name, _fingerprint(value)) for name, value in kwargs.items())),
    )


def _cache_counts(function: Callable[..., object]) -> tuple[int, int] | None:
    cache_info = getattr(function, "cache_info", None)
    if not callable(cache_info):
        return None
    try:
        info = cache_info()
        return int(info.hits), int(info.misses)
    except (AttributeError, TypeError, ValueError):
        return None


def _copy_cache_api(source: Callable[..., object], target: Callable[..., object]) -> None:
    for attribute in ("cache_clear", "cache_info", "cache_parameters"):
        value = getattr(source, attribute, None)
        if value is not None:
            setattr(target, attribute, value)


def profile_property(
    operation: str,
    layer: str = "semantic",
    minimum_mode: str = "audit",
    capture_inputs: bool = True,
    track_cache: bool = False,
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
    """Decorate one thermodynamic property operation."""

    minimum = _minimum_mode(minimum_mode)

    def decorate(function: Callable[_P, _R]) -> Callable[_P, _R]:
        if not _mode_enabled(minimum):
            return function
        try:
            signature: inspect.Signature | None = inspect.signature(function)
        except (TypeError, ValueError):
            signature = None

        def prepare(
            args: tuple[object, ...],
            kwargs: dict[str, object],
        ) -> tuple[PerformanceTrace | None, str, str, tuple[int, int] | None]:
            trace = _CURRENT_TRACE.get()
            medium = _medium_name(args)
            key = f"{layer}|{medium}|{operation}"
            if trace is not None and trace.mode == "audit" and capture_inputs:
                trace._record_exact_input(
                    key,
                    operation=operation,
                    layer=layer,
                    medium=medium,
                    fingerprint=_input_fingerprint(signature, args, kwargs),
                )
            before = _cache_counts(function) if trace is not None and track_cache else None
            return trace, medium, key, before

        def finish_cache(
            trace: PerformanceTrace | None,
            medium: str,
            key: str,
            before: tuple[int, int] | None,
        ) -> None:
            if trace is None or before is None:
                return
            after = _cache_counts(function)
            if after is None:
                return
            trace._record_cache(
                key,
                operation=operation,
                layer=layer,
                medium=medium,
                hits=after[0] - before[0],
                misses=after[1] - before[1],
            )

        if inspect.iscoroutinefunction(function):

            @wraps(function)
            async def async_wrapper(*args: _P.args, **kwargs: _P.kwargs) -> Any:
                object_args = cast(tuple[object, ...], args)
                object_kwargs = cast(dict[str, object], kwargs)
                trace, medium, key, before = prepare(object_args, object_kwargs)
                try:
                    if trace is None:
                        return await function(*args, **kwargs)
                    with _tracked_span(
                        trace,
                        f"property.{_public_property_key(key)}",
                        property_key=key,
                        property_operation=operation,
                    ):
                        return await function(*args, **kwargs)
                finally:
                    finish_cache(trace, medium, key, before)

            _copy_cache_api(function, async_wrapper)
            return cast(Callable[_P, _R], async_wrapper)

        @wraps(function)
        def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
            object_args = cast(tuple[object, ...], args)
            object_kwargs = cast(dict[str, object], kwargs)
            trace, medium, key, before = prepare(object_args, object_kwargs)
            try:
                if trace is None:
                    return function(*args, **kwargs)
                with _tracked_span(
                    trace,
                    f"property.{_public_property_key(key)}",
                    property_key=key,
                    property_operation=operation,
                ):
                    return function(*args, **kwargs)
            finally:
                finish_cache(trace, medium, key, before)

        _copy_cache_api(function, wrapper)
        return wrapper

    return decorate


def record_property_iterations(
    operation: str,
    iterations: int,
    converged: bool,
) -> None:
    """Record inverse-property solver iterations in audit mode."""

    trace = _CURRENT_TRACE.get()
    if trace is None or trace.mode != "audit":
        return
    for frame in reversed(_ACTIVE_SPANS.get()):
        if (
            frame.property_key is not None
            and frame.property_operation == operation
        ):
            layer, medium, _unused_operation = frame.property_key.split("|", 2)
            trace._record_iterations(
                frame.property_key,
                operation=operation,
                layer=layer,
                medium=medium,
                iterations=iterations,
                converged=converged,
            )
            return
    key = f"semantic|unknown|{operation}"
    trace._record_iterations(
        key,
        operation=operation,
        layer="semantic",
        medium="unknown",
        iterations=iterations,
        converged=converged,
    )


def record_property_cache(operation: str, *, hit: bool) -> None:
    """Record one run-local property-cache lookup in audit mode."""

    trace = _CURRENT_TRACE.get()
    if trace is None or trace.mode != "audit":
        return
    for frame in reversed(_ACTIVE_SPANS.get()):
        if (
            frame.property_key is not None
            and frame.property_operation == operation
        ):
            layer, medium, _unused_operation = frame.property_key.split("|", 2)
            trace._record_cache(
                frame.property_key,
                operation=operation,
                layer=layer,
                medium=medium,
                hits=1 if hit else 0,
                misses=0 if hit else 1,
            )
            return


__all__ = [
    "PROFILE_MODE",
    "PerformanceTrace",
    "performance_span",
    "profile_phase",
    "profile_property",
    "profile_run",
    "record_property_cache",
    "record_property_iterations",
]