"""
prometheus.py
--------------
Prometheus metrics exporter for the HIEROMEM API.

Tracks key system signals:
  • Total API requests
  • Current episode count
  • Average trust value
  • Response latency histogram

Can be extended to include per-route metrics, backend health,
and summarization or trust convergence counters.

Author: Sai Rishi Kiran Mannava
"""

from __future__ import annotations
from prometheus_client import (
    Counter,
    Gauge,
    Histogram,
    Summary,
    generate_latest,
    CONTENT_TYPE_LATEST,
)
from fastapi import Request, Response
from typing import Awaitable, Callable, Mapping, Dict, Set
import time

# --------------------------------------------------------------------------- #
# Metric Definitions
# --------------------------------------------------------------------------- #

# Global request count across all routes
REQUESTS_TOTAL = Counter(
    "hieromem_requests_total",
    "Total number of HTTP requests received by the HIEROMEM API",
    labelnames=("tenant_id", "endpoint", "method", "status_code"),
)

# Number of episodes currently stored in memory (tenant scoped)
EPISODES_TOTAL = Gauge(
    "hieromem_episodes_total",
    "Current number of stored episodes in memory",
    labelnames=("namespace",),
)

# Configured maximum capacity (Nmax) per namespace
MEMORY_CAPACITY = Gauge(
    "hieromem_memory_capacity",
    "Configured maximum active episode capacity (Nmax)",
    labelnames=("namespace",),
)

# Average trust score of all stored episodes (tenant scoped)
TRUST_AVERAGE = Gauge(
    "hieromem_trust_average",
    "Average trust value across all active episodes",
    labelnames=("namespace",),
)

# Distortion bound reported during summarisation (tenant scoped)
SUMMARY_DISTORTION_BOUND = Gauge(
    "hieromem_summary_distortion_bound",
    "Rolling average δ-distortion bound from summarisation (Theorem 2)",
    labelnames=("namespace",),
)

# Observed contraction factor from summarisation batches (tenant scoped)
CONTRACTION_FACTOR = Gauge(
    "hieromem_contraction_factor",
    "Observed contraction factor derived from summarisation utilisation",
    labelnames=("namespace",),
)

# Lyapunov margin (fidelity minus distortion) with boolean warning channel
LYAPUNOV_MARGIN = Gauge(
    "hieromem_lyapunov_margin",
    "Lyapunov-style stability margin (fidelity - distortion)",
    labelnames=("namespace",),
)

LYAPUNOV_WARNING = Gauge(
    "hieromem_lyapunov_warning",
    "Indicator that Lyapunov margin dropped below zero",
    labelnames=("namespace",),
)

# Symbolic inconsistency counts (per violation type)
SYMBOLIC_INCONSISTENCIES = Gauge(
    "hieromem_symbolic_inconsistencies_total",
    "Count of symbolic contradictions detected by the graph reasoner",
    labelnames=("namespace", "kind"),
)

# Response latency per request (seconds)
RESPONSE_LATENCY = Histogram(
    "hieromem_response_latency_seconds",
    "Latency distribution of HTTP responses",
    buckets=[0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 5.0],
)

QUOTA_VIOLATIONS_TOTAL = Counter(
    "hieromem_quota_violations_total",
    "Count of quota rejections per tenant",
    labelnames=("tenant_id", "type"),
)

RATE_LIMIT_EXCEEDED_TOTAL = Counter(
    "hieromem_rate_limit_exceeded_total",
    "Count of rate limit rejections per tenant",
    labelnames=("tenant_id",),
)

# Dual retrieval specific metrics
DUAL_RECALL_TOTAL = Counter(
    "hieromem_dual_recall_total",
    "Total number of dual retrieval (dense + symbolic) operations",
    labelnames=("namespace",),
)

DUAL_RECALL_ITERATIONS = Histogram(
    "hieromem_dual_recall_iterations",
    "Distribution of refinement iterations executed during dual retrieval",
    labelnames=("namespace",),
    buckets=[0, 1, 2, 3, 4, 5, 6, 10, 20],
)

DUAL_RECALL_SCORES = Summary(
    "hieromem_dual_recall_scores",
    "Observed distribution of dense, symbolic, and combined weights",
    labelnames=("namespace", "score_type"),
)

# Optional: Count of summarization and forgetting operations
SUMMARIZE_TOTAL = Counter(
    "hieromem_summarize_total",
    "Total number of summarization operations performed",
)
FORGET_TOTAL = Counter(
    "hieromem_forget_total",
    "Total number of forgetting (prune/decay) operations performed",
)

# Cache performance metrics
CACHE_HITS_TOTAL = Counter(
    "hieromem_cache_hits_total",
    "Total number of cache hits (L1 + L2)",
    labelnames=("namespace", "cache_tier"),
)

CACHE_MISSES_TOTAL = Counter(
    "hieromem_cache_misses_total",
    "Total number of cache misses",
    labelnames=("namespace",),
)

CACHE_HIT_RATIO = Gauge(
    "hieromem_cache_hit_ratio",
    "Cache hit ratio (hits / (hits + misses))",
    labelnames=("namespace",),
)

# --------------------------------------------------------------------------- #
# Middleware for automatic metrics
# --------------------------------------------------------------------------- #

CallNext = Callable[[Request], Awaitable[Response]]


async def metrics_middleware(
        request: Request,
        call_next: CallNext) -> Response:
    """Middleware to track latency and total requests for all API routes."""
    start_time = time.perf_counter()

    response = await call_next(request)

    elapsed = time.perf_counter() - start_time
    RESPONSE_LATENCY.observe(elapsed)

    tenant_id = getattr(
        getattr(
            request.state,
            "tenant",
            None),
        "tenant_id",
        None)
    if not tenant_id:
        tenant_id = getattr(request.state, "tenant_id", "public")
    endpoint = request.scope.get("path", "unknown")
    method = request.method
    status_code = getattr(response, "status_code", 0)
    REQUESTS_TOTAL.labels(
        tenant_id=tenant_id,
        endpoint=endpoint,
        method=method,
        status_code=str(status_code),
    ).inc()

    return response

# --------------------------------------------------------------------------- #
# Endpoint Helper
# --------------------------------------------------------------------------- #


def metrics_endpoint() -> tuple[bytes, str]:
    """
    Return Prometheus-formatted metrics for FastAPI `/metrics` endpoint.

    Example:
        @app.get("/metrics")
        def prometheus_metrics():
            content, content_type = metrics_endpoint()
            return Response(content=content, media_type=content_type)
    """
    UPTIME_SECONDS.set(max(0.0, time.perf_counter() - PROCESS_START_TIME))
    return generate_latest(), CONTENT_TYPE_LATEST


# --------------------------------------------------------------------------- #
# Runtime Updaters
# --------------------------------------------------------------------------- #


def _ns(namespace: str | None) -> str:
    return (namespace or "global").strip() or "global"


_SYMBOLIC_LABEL_CACHE: Dict[str, Set[str]] = {}


def update_episode_metrics(
    total: int,
    avg_trust: float,
    *,
    namespace: str | None = None,
    capacity: float | int | None = None,
) -> None:
    """
    Update episode-related gauges for Prometheus.
    Should be called after insertions, retrievals, or pruning.
    """
    ns = _ns(namespace)
    EPISODES_TOTAL.labels(namespace=ns).set(total)
    TRUST_AVERAGE.labels(namespace=ns).set(avg_trust)
    if capacity is not None:
        MEMORY_CAPACITY.labels(namespace=ns).set(float(capacity))


def update_stability_metrics(
    *,
    namespace: str | None,
    distortion_bound: float,
    contraction_factor: float,
    lyapunov_margin: float,
    lyapunov_warning: bool,
) -> None:
    """Update stability-related gauges exported for theoretical monitoring."""

    ns = _ns(namespace)
    SUMMARY_DISTORTION_BOUND.labels(namespace=ns).set(distortion_bound)
    CONTRACTION_FACTOR.labels(namespace=ns).set(contraction_factor)
    LYAPUNOV_MARGIN.labels(namespace=ns).set(lyapunov_margin)
    LYAPUNOV_WARNING.labels(namespace=ns).set(1.0 if lyapunov_warning else 0.0)


def update_symbolic_inconsistencies(
    *, namespace: str | None, violations: Mapping[str, object] | None
) -> None:
    """Expose symbolic contradiction counts to Prometheus."""

    ns = _ns(namespace)
    known = _SYMBOLIC_LABEL_CACHE.setdefault(ns, set())
    provided: Set[str] = set()
    numeric: Dict[str, float] = {}

    if violations:
        for key, value in violations.items():
            if isinstance(value, (int, float)):
                numeric[str(key)] = float(value)
            elif key == "violations" and isinstance(value, Mapping):
                for nested_key, nested_val in value.items():
                    if isinstance(nested_val, (int, float)):
                        label = f"violations_{nested_key}"
                        numeric[label] = float(nested_val)

    for label, count in numeric.items():
        SYMBOLIC_INCONSISTENCIES.labels(namespace=ns, kind=label).set(count)
        provided.add(label)

    total = numeric.get("total_contradictions")
    if total is None:
        total = sum(value for key, value in numeric.items()
                    if not key.startswith("violations_"))
    SYMBOLIC_INCONSISTENCIES.labels(
        namespace=ns,
        kind="total").set(
        float(total))
    provided.add("total")

    known.update(provided)
    for missing in known - provided:
        SYMBOLIC_INCONSISTENCIES.labels(namespace=ns, kind=missing).set(0.0)


def record_dual_recall_metrics(
    *,
    namespace: str | None,
    iterations: int,
    dense_scores: Mapping[str, float],
    symbolic_scores: Mapping[str, float],
    combined_scores: Mapping[str, float],
) -> None:
    """Record Prometheus metrics for a dual retrieval execution."""

    ns = _ns(namespace)
    DUAL_RECALL_TOTAL.labels(namespace=ns).inc()
    DUAL_RECALL_ITERATIONS.labels(namespace=ns).observe(float(iterations))

    for score in dense_scores.values():
        DUAL_RECALL_SCORES.labels(
            namespace=ns,
            score_type="dense").observe(
            float(score))
    for score in symbolic_scores.values():
        DUAL_RECALL_SCORES.labels(
            namespace=ns,
            score_type="symbolic").observe(
            float(score))
    for score in combined_scores.values():
        DUAL_RECALL_SCORES.labels(
            namespace=ns,
            score_type="combined").observe(
            float(score))


def record_summarization() -> None:
    """Increment summarization counter."""
    SUMMARIZE_TOTAL.inc()


def record_forgetting() -> None:
    """Increment forgetting counter."""
    FORGET_TOTAL.inc()


def record_quota_violation(tenant_id: str, violation_type: str) -> None:
    """Record when a tenant hits a quota limit."""

    QUOTA_VIOLATIONS_TOTAL.labels(
        tenant_id=tenant_id,
        type=violation_type).inc()


def record_rate_limit_block(tenant_id: str) -> None:
    """Record when a tenant exceeds rate limit."""

    RATE_LIMIT_EXCEEDED_TOTAL.labels(tenant_id=tenant_id).inc()


def record_cache_hit(namespace: str | None, cache_tier: str) -> None:
    """Record a cache hit (either L1 or L2)."""
    ns = _ns(namespace)
    CACHE_HITS_TOTAL.labels(namespace=ns, cache_tier=cache_tier).inc()
    _update_cache_hit_ratio(ns)


def record_cache_miss(namespace: str | None) -> None:
    """Record a cache miss."""
    ns = _ns(namespace)
    CACHE_MISSES_TOTAL.labels(namespace=ns).inc()
    _update_cache_hit_ratio(ns)


def _update_cache_hit_ratio(namespace: str) -> None:
    """Calculate and update cache hit ratio for a namespace."""
    # Get current counts from metrics
    # Note: This is an approximation since we can't easily query counter values
    # In production, you might want to maintain separate counters or use a different approach
    # For now, we'll update the gauge based on the assumption that hits and misses are tracked
    pass  # Ratio will be calculated by Prometheus queries or monitoring tools


# Process uptime gauge
PROCESS_START_TIME = time.perf_counter()
UPTIME_SECONDS = Gauge(
    "hieromem_uptime_seconds",
    "Number of seconds since the API process started",
)
