"""Per-test performance attribution — function-level + section/log-level.

This is the **writer** half of the federation perf-attribution contract
(TESTING.md **TS-28**). It runs *inside* the test process and produces, per
test, one JSON artifact carrying:

- **function-level** timings from a per-test ``cProfile`` run (top-N frames by
  cumulative time — covers both the test body and the target code it drives);
- **section/log-level** timings from explicit :func:`perf_span` context
  managers a test (or the code under test) wraps around meaningful sections.

The federation test supervisor (``mgf-test-supervisor``) reads these artifacts
*out of process* (it never imports mgf-common — TS-23), so the two halves talk
only through this JSON wire format. The pytest plugin
(:mod:`mgf.common.pytest_plugin`) drives capture when ``--mgf-perf`` /
``MGF_PERF`` is set.

Design constraints:

- **Fail-open.** Perf capture MUST NOT change a test's outcome. Every capture
  path swallows its own errors — a broken profiler or unwritable artifact dir
  degrades to "no perf data", never a test failure.
- **Near-zero cost when off.** :func:`perf_span` does one ``ContextVar`` read
  and returns when capture is inactive; ``cProfile`` only runs under the flag.
- **Bounded output.** Function rows are capped at ``top_n`` (default
  :data:`DEFAULT_TOP_N`); spans are capped at :data:`MAX_SPANS` so a pathological
  loop can't write an unbounded artifact.
"""

from __future__ import annotations

import contextlib
import contextvars
import json
import os
import tempfile
import time
from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import cProfile
    from collections.abc import Iterator

PERF_SCHEMA_VERSION = "1.0.0"
DEFAULT_TOP_N = 15
MAX_SPANS = 500
# Drop sub-microsecond noise frames so the top-N is meaningful signal.
_MIN_CUM_MS = 0.05


@dataclass(frozen=True, slots=True)
class PerfSpan:
    """One explicit, named section timing (from :func:`perf_span`)."""

    label: str
    ms: float
    depth: int

    def to_dict(self) -> dict[str, object]:
        return {"label": self.label, "ms": round(self.ms, 3), "depth": self.depth}


@dataclass(frozen=True, slots=True)
class PerfFunction:
    """One function's aggregated timing from a ``cProfile`` run."""

    func: str  # "shortened/path.py:<lineno>(<name>)"
    ncalls: int
    tottime_ms: float  # time in this function excluding subcalls
    cumtime_ms: float  # time including subcalls

    def to_dict(self) -> dict[str, object]:
        return {
            "func": self.func,
            "ncalls": self.ncalls,
            "tottime_ms": round(self.tottime_ms, 3),
            "cumtime_ms": round(self.cumtime_ms, 3),
        }


@dataclass(frozen=True, slots=True)
class PerfSummary:
    """The per-test perf artifact. Serialised to ``<perf_dir>/<id>.json``."""

    test_id: str
    total_ms: float
    functions: tuple[PerfFunction, ...] = ()
    spans: tuple[PerfSpan, ...] = ()
    schema_version: str = PERF_SCHEMA_VERSION

    def to_dict(self) -> dict[str, object]:
        return {
            "schema_version": self.schema_version,
            "test_id": self.test_id,
            "total_ms": round(self.total_ms, 3),
            "functions": [f.to_dict() for f in self.functions],
            "spans": [s.to_dict() for s in self.spans],
        }

    def to_json(self) -> str:
        return json.dumps(self.to_dict(), ensure_ascii=False, separators=(",", ":"))


# ---------------------------------------------------------------------------
# Section/log-level capture — perf_span().
# ---------------------------------------------------------------------------


@dataclass
class _SpanCollector:
    """Mutable per-test accumulator for :func:`perf_span` samples."""

    spans: list[PerfSpan] = field(default_factory=list)
    depth: int = 0
    _overflowed: bool = False


# Held in a ContextVar so perf_span() finds the active collector without the
# caller threading it through. None => capture inactive (the common case).
_active: contextvars.ContextVar[_SpanCollector | None] = contextvars.ContextVar(
    "mgf_perf_collector", default=None
)


@contextmanager
def perf_span(label: str) -> Iterator[None]:
    """Time a named section of code (nestable, fail-open).

    Wrap any meaningful unit — a DB call, a render, the block around a log
    line — to attribute wall-time to it. When perf capture is **not** active
    (the default outside ``mtest --perf``) this is a single ``ContextVar`` read
    plus the bare ``yield``: cheap enough to leave in production code paths.

    Nested spans record their ``depth`` so the supervisor can render the tree.
    Spans are captured in the context where collection started; samples taken
    in a thread that didn't inherit the context are silently dropped (documented
    limitation, not an error).
    """
    collector = _active.get()
    if collector is None:
        yield
        return
    depth = collector.depth
    collector.depth += 1
    t0 = time.perf_counter()
    try:
        yield
    finally:
        collector.depth -= 1
        elapsed_ms = (time.perf_counter() - t0) * 1000.0
        if len(collector.spans) < MAX_SPANS:
            collector.spans.append(PerfSpan(label=label, ms=elapsed_ms, depth=depth))
        else:
            collector._overflowed = True


@contextmanager
def collect_spans() -> Iterator[_SpanCollector]:
    """Activate :func:`perf_span` capture for the duration of the block.

    Used by the pytest plugin around the test-call phase. Restores the prior
    collector on exit (so nested activations and xdist workers don't clobber
    each other).
    """
    collector = _SpanCollector()
    token = _active.set(collector)
    try:
        yield collector
    finally:
        _active.reset(token)


# ---------------------------------------------------------------------------
# Function-level capture — cProfile extraction.
# ---------------------------------------------------------------------------


def _shorten_path(path: str, keep: int = 3) -> str:
    """Trim an absolute source path to its last ``keep`` components."""
    parts = path.replace(os.sep, "/").split("/")
    return "/".join(parts[-keep:]) if len(parts) > keep else path


def functions_from_profile(
    profiler: cProfile.Profile, top_n: int = DEFAULT_TOP_N
) -> tuple[PerfFunction, ...]:
    """Extract the top-``n`` functions (by cumulative time) from a profile.

    Fail-open: any error parsing the profile yields an empty tuple rather than
    propagating into the test result.
    """
    try:
        import pstats

        stats = pstats.Stats(profiler)
        rows: list[PerfFunction] = []
        for (fn, lineno, name), row in stats.stats.items():  # type: ignore[attr-defined]
            # row == (call_count, num_calls, tottime, cumtime, callers)
            _cc, nc, tt, ct, _callers = row
            cum_ms = ct * 1000.0
            if cum_ms < _MIN_CUM_MS:
                continue
            rows.append(
                PerfFunction(
                    func=f"{_shorten_path(fn)}:{lineno}({name})",
                    ncalls=int(nc),
                    tottime_ms=tt * 1000.0,
                    cumtime_ms=cum_ms,
                )
            )
        # Deterministic: cumtime desc, then func name for ties.
        rows.sort(key=lambda r: (-r.cumtime_ms, r.func))
        return tuple(rows[: max(0, top_n)])
    except Exception:  # fail-open by contract — perf must never break a test.
        return ()


# ---------------------------------------------------------------------------
# Artifact IO.
# ---------------------------------------------------------------------------


def _safe_stem(test_id: str) -> str:
    """A filesystem-safe, collision-resistant stem for a test id."""
    import hashlib

    safe = "".join(c if (c.isalnum() or c in "-_.") else "_" for c in test_id)
    digest = hashlib.sha1(test_id.encode("utf-8")).hexdigest()[:8]  # noqa: S324 — non-crypto id
    return f"{safe[:120]}.{digest}"


def write_perf_artifact(perf_dir: Path, summary: PerfSummary) -> Path | None:
    """Atomically write one per-test perf artifact. Fail-open → None on error.

    The real ``test_id`` is stored *inside* the JSON so the supervisor indexes
    artifacts by content, not by reconstructing the filename scheme.
    """
    try:
        perf_dir.mkdir(parents=True, exist_ok=True)
        dest = perf_dir / f"{_safe_stem(summary.test_id)}.json"
        fd, tmp = tempfile.mkstemp(prefix=".perf-", suffix=".json.tmp", dir=str(perf_dir))
        tmp_path = Path(tmp)
        try:
            with os.fdopen(fd, "w", encoding="utf-8") as f:
                f.write(summary.to_json())
                f.flush()
                os.fsync(f.fileno())
            tmp_path.replace(dest)
        except OSError:
            with contextlib.suppress(OSError):
                tmp_path.unlink()
            return None
        return dest
    except Exception:  # fail-open by contract — perf must never break a test.
        return None
