"""
cortex/core/metrics.py — Metrics Collection (Issue #368).

Lightweight, pure-Python metrics collector for Cortex pipeline observability.
Tracks counters, histograms (with percentile support), and gauges — all keyed
by metric name and optional labels.

Design decisions:
    - Singleton per owner_id: MetricsCollector.get(owner_id) returns the same
      instance for the same owner, allowing global metric aggregation.
    - Thread-safe: uses asyncio.Lock for all mutations.
    - Zero-overhead when disabled: all public methods short-circuit immediately
      when enabled=False (no lock acquisition, no data structures touched).
    - No external dependencies: pure Python stdlib only.
    - Optional Prometheus text exposition format via export_prometheus().

Dependencies: None (leaf module).
Dependents: cortex.core.quality_gates, cortex.api.plugin.
"""
from __future__ import annotations

import asyncio
import math
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple


# ---------------------------------------------------------------------------
# Metric type definitions
# ---------------------------------------------------------------------------

@dataclass
class _Counter:
    """Monotonically increasing counter."""
    value: float = 0.0
    labels: Dict[str, float] = field(default_factory=dict)


@dataclass
class _Histogram:
    """Records individual observations for percentile calculation."""
    values: List[float] = field(default_factory=list)
    labels: Dict[str, List[float]] = field(default_factory=dict)


@dataclass
class _Gauge:
    """Point-in-time value (can go up or down)."""
    value: float = 0.0
    labels: Dict[str, float] = field(default_factory=dict)
    _is_set: bool = False  # True once set_gauge() has been called


# ---------------------------------------------------------------------------
# Default metric definitions
# ---------------------------------------------------------------------------

DEFAULT_COUNTERS = [
    "memories_stored",
    "memories_retrieved",
    "extractions_run",
    "reconciliations_run",
    "promotions_run",
    "boot_sequences",
    "re_retrievals",
    "sleep_cycles",
    "dream_cycles",
]

DEFAULT_HISTOGRAMS = [
    "retrieval_latency_ms",
    "extraction_latency_ms",
    "boot_latency_ms",
    "compilation_latency_ms",
]

DEFAULT_GAUGES = [
    "active_claims",
    "working_memory_size",
    "context_budget_usage_pct",
    "extraction_success_rate",
]


# ---------------------------------------------------------------------------
# MetricsCollector
# ---------------------------------------------------------------------------

class MetricsCollector:
    """
    Lightweight metrics collector for Cortex pipeline observability.

    Supports three metric types:
        - Counters: monotonically increasing values (increment only)
        - Histograms: individual observations with percentile calculation
        - Gauges: point-in-time values (can go up or down)

    All methods are async and thread-safe via asyncio.Lock.
    When ``enabled=False``, all operations are zero-cost no-ops.

    Usage::

        metrics = MetricsCollector.get("owner-1")
        await metrics.increment("memories_stored")
        await metrics.record("retrieval_latency_ms", 42.5)
        summary = await metrics.get_summary()
    """

    # Singleton registry: owner_id → MetricsCollector
    _instances: Dict[str, "MetricsCollector"] = {}

    @classmethod
    def get(cls, owner_id: str = "default", *, enabled: bool = True) -> "MetricsCollector":
        """Return the singleton MetricsCollector for *owner_id*.

        Creates a new instance on first call for a given owner_id.
        Subsequent calls return the same instance (singleton per owner).
        """
        if owner_id not in cls._instances:
            cls._instances[owner_id] = cls(owner_id=owner_id, enabled=enabled)
        return cls._instances[owner_id]

    @classmethod
    def reset_instances(cls) -> None:
        """Clear the singleton registry (for testing only)."""
        cls._instances.clear()

    def __init__(self, owner_id: str = "default", *, enabled: bool = True) -> None:
        self.owner_id = owner_id
        self.enabled = enabled
        self._lock = asyncio.Lock()
        self._counters: Dict[str, _Counter] = {}
        self._histograms: Dict[str, _Histogram] = {}
        self._gauges: Dict[str, _Gauge] = {}
        self._created_at = time.time()

        # Pre-register default metrics
        for name in DEFAULT_COUNTERS:
            self._counters[name] = _Counter()
        for name in DEFAULT_HISTOGRAMS:
            self._histograms[name] = _Histogram()
        for name in DEFAULT_GAUGES:
            self._gauges[name] = _Gauge()

    # ------------------------------------------------------------------ #
    # Public API
    # ------------------------------------------------------------------ #

    async def increment(self, counter_name: str, labels: Optional[Dict[str, str]] = None, amount: float = 1.0) -> None:
        """Increment a counter by *amount* (default 1).

        Creates the counter if it doesn't exist.
        """
        if not self.enabled:
            return
        label_key = _label_key(labels)
        async with self._lock:
            if counter_name not in self._counters:
                self._counters[counter_name] = _Counter()
            counter = self._counters[counter_name]
            if label_key:
                counter.labels[label_key] = counter.labels.get(label_key, 0.0) + amount
            else:
                counter.value += amount

    async def record(self, metric_name: str, value: float, labels: Optional[Dict[str, str]] = None) -> None:
        """Record an observation for a histogram, or set a gauge value.

        If *metric_name* is a registered gauge, sets its value.
        Otherwise, records as a histogram observation.
        Creates the metric if it doesn't exist.
        """
        if not self.enabled:
            return
        label_key = _label_key(labels)
        async with self._lock:
            # Check if it's a gauge first
            if metric_name in self._gauges:
                gauge = self._gauges[metric_name]
                gauge._is_set = True
                if label_key:
                    gauge.labels[label_key] = value
                else:
                    gauge.value = value
                return

            # Otherwise treat as histogram
            if metric_name not in self._histograms:
                self._histograms[metric_name] = _Histogram()
            hist = self._histograms[metric_name]
            if label_key:
                if label_key not in hist.labels:
                    hist.labels[label_key] = []
                hist.labels[label_key].append(value)
            else:
                hist.values.append(value)

    async def set_gauge(self, gauge_name: str, value: float, labels: Optional[Dict[str, str]] = None) -> None:
        """Explicitly set a gauge value.

        Creates the gauge if it doesn't exist.
        """
        if not self.enabled:
            return
        label_key = _label_key(labels)
        async with self._lock:
            if gauge_name not in self._gauges:
                self._gauges[gauge_name] = _Gauge()
            gauge = self._gauges[gauge_name]
            gauge._is_set = True
            if label_key:
                gauge.labels[label_key] = value
            else:
                gauge.value = value

    async def get_summary(self) -> Dict[str, Any]:
        """Return a snapshot of all metrics.

        Returns:
            Dict with keys 'counters', 'histograms', 'gauges', each containing
            the metric name → value/stats mapping.
        """
        async with self._lock:
            counters = {}
            for name, c in self._counters.items():
                if c.labels:
                    counters[name] = {"total": c.value, "by_label": dict(c.labels)}
                else:
                    counters[name] = c.value

            histograms = {}
            for name, h in self._histograms.items():
                stats = _histogram_stats(h.values) if h.values else {}
                if h.labels:
                    label_stats = {k: _histogram_stats(v) for k, v in h.labels.items() if v}
                    stats["by_label"] = label_stats
                histograms[name] = stats

            gauges = {}
            for name, g in self._gauges.items():
                if not g._is_set:
                    continue  # Skip gauges never explicitly set (cold-start grace)
                if g.labels:
                    gauges[name] = {"value": g.value, "by_label": dict(g.labels)}
                else:
                    gauges[name] = g.value

            return {
                "owner_id": self.owner_id,
                "counters": counters,
                "histograms": histograms,
                "gauges": gauges,
                "uptime_seconds": time.time() - self._created_at,
            }

    async def reset(self) -> None:
        """Reset all metrics to their initial state (for testing)."""
        async with self._lock:
            for c in self._counters.values():
                c.value = 0.0
                c.labels.clear()
            for h in self._histograms.values():
                h.values.clear()
                h.labels.clear()
            for g in self._gauges.values():
                g.value = 0.0
                g.labels.clear()
            self._created_at = time.time()

    def export_prometheus(self) -> str:
        """Export metrics in Prometheus text exposition format (synchronous).

        Returns a string suitable for serving at a /metrics endpoint.
        This is intentionally synchronous to avoid async overhead in the
        scrape path — it reads the current state without locking.
        """
        if not self.enabled:
            return ""

        lines: List[str] = []

        # Counters
        for name, c in self._counters.items():
            fqn = f"cortex_{name}"
            lines.append(f"# TYPE {fqn} counter")
            if c.value > 0 or not c.labels:
                lines.append(f"{fqn} {c.value}")
            for label_key, val in c.labels.items():
                lines.append(f'{fqn}{{{label_key}}} {val}')

        # Histograms (export as summary with quantiles)
        for name, h in self._histograms.items():
            fqn = f"cortex_{name}"
            lines.append(f"# TYPE {fqn} summary")
            if h.values:
                stats = _histogram_stats(h.values)
                lines.append(f'{fqn}{{quantile="0.5"}} {stats.get("p50", 0)}')
                lines.append(f'{fqn}{{quantile="0.95"}} {stats.get("p95", 0)}')
                lines.append(f'{fqn}{{quantile="0.99"}} {stats.get("p99", 0)}')
                lines.append(f"{fqn}_count {stats.get('count', 0)}")
                lines.append(f"{fqn}_sum {stats.get('sum', 0)}")

        # Gauges
        for name, g in self._gauges.items():
            fqn = f"cortex_{name}"
            lines.append(f"# TYPE {fqn} gauge")
            if g.value > 0 or not g.labels:
                lines.append(f"{fqn} {g.value}")
            for label_key, val in g.labels.items():
                lines.append(f'{fqn}{{{label_key}}} {val}')

        return "\n".join(lines) + "\n"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _label_key(labels: Optional[Dict[str, str]]) -> str:
    """Convert a labels dict to a stable string key for storage."""
    if not labels:
        return ""
    return ",".join(f'{k}="{v}"' for k, v in sorted(labels.items()))


def _histogram_stats(values: List[float]) -> Dict[str, Any]:
    """Calculate summary statistics for a list of observations."""
    if not values:
        return {}
    sorted_vals = sorted(values)
    n = len(sorted_vals)
    return {
        "count": n,
        "sum": sum(sorted_vals),
        "min": sorted_vals[0],
        "max": sorted_vals[-1],
        "mean": sum(sorted_vals) / n,
        "p50": _percentile(sorted_vals, 0.50),
        "p95": _percentile(sorted_vals, 0.95),
        "p99": _percentile(sorted_vals, 0.99),
    }


def _percentile(sorted_values: List[float], pct: float) -> float:
    """Calculate the p-th percentile from a pre-sorted list of values.

    Uses nearest-rank method for simplicity.
    """
    if not sorted_values:
        return 0.0
    n = len(sorted_values)
    idx = max(0, min(n - 1, int(math.ceil(pct * n)) - 1))
    return sorted_values[idx]
