"""Operational metric definitions for observability.

Defines the 14 operational metric keys tracked across the plan
lifecycle and a lightweight ``MetricEntry`` model for recording
individual metric observations.

``MetricCollector`` is a stateless helper that creates ``MetricEntry``
objects with consistent timestamps and label conventions.

``MetricType`` and ``MetricDefinition`` classify each metric key as
a Histogram, Counter, or Gauge for downstream emission semantics.

Based on ``docs/specification.md`` Observability section and
Forgejo issues #500 / #579.
"""

from __future__ import annotations

from datetime import datetime
from enum import StrEnum
from typing import Any

from pydantic import BaseModel, ConfigDict, Field


class OperationalMetricKey(StrEnum):
    """Enumeration of the 14 operational metric keys.

    Metrics are grouped by subsystem:

    **Plan-level**
        ``PLAN_DURATION_MS``, ``PLAN_TOTAL_COST_USD``,
        ``PLAN_DECISION_COUNT``, ``SUBPLAN_COUNT``

    **Actor-level**
        ``ACTOR_INVOCATION_COUNT``, ``ACTOR_LATENCY_MS``

    **Tool-level**
        ``TOOL_INVOCATION_COUNT``, ``TOOL_ERROR_RATE``

    **Context-level**
        ``CONTEXT_BUILD_TIME_MS``, ``CONTEXT_TOKEN_COUNT``

    **LLM-level**
        ``LLM_CALL_COUNT``, ``LLM_TOTAL_TOKENS``,
        ``LLM_TOTAL_COST_USD``, ``LLM_AVG_LATENCY_MS``
    """

    PLAN_DURATION_MS = "plan_duration_ms"
    PLAN_TOTAL_COST_USD = "plan_total_cost_usd"
    PLAN_DECISION_COUNT = "plan_decision_count"
    ACTOR_INVOCATION_COUNT = "actor_invocation_count"
    ACTOR_LATENCY_MS = "actor_latency_ms"
    TOOL_INVOCATION_COUNT = "tool_invocation_count"
    TOOL_ERROR_RATE = "tool_error_rate"
    CONTEXT_BUILD_TIME_MS = "context_build_time_ms"
    CONTEXT_TOKEN_COUNT = "context_token_count"
    LLM_CALL_COUNT = "llm_call_count"
    LLM_TOTAL_TOKENS = "llm_total_tokens"
    LLM_TOTAL_COST_USD = "llm_total_cost_usd"
    LLM_AVG_LATENCY_MS = "llm_avg_latency_ms"
    SUBPLAN_COUNT = "subplan_count"


class MetricType(StrEnum):
    """Semantic type of a metric for emission semantics.

    Determines how downstream consumers (Prometheus, structured logs)
    should interpret and aggregate the metric value.

    ``HISTOGRAM`` — distribution of observed values (e.g. latencies).
    ``COUNTER`` — monotonically increasing cumulative count.
    ``GAUGE`` — point-in-time value that can go up or down.
    """

    HISTOGRAM = "histogram"
    COUNTER = "counter"
    GAUGE = "gauge"


class MetricDefinition(BaseModel):
    """Links an ``OperationalMetricKey`` to its ``MetricType``.

    Attributes:
        key: The operational metric key.
        metric_type: Histogram, Counter, or Gauge.
        unit: Human-readable unit string (e.g. ``"ms"``, ``"usd"``).
        description: Short description of the metric.
    """

    key: OperationalMetricKey = Field(..., description="The operational metric key")
    metric_type: MetricType = Field(
        ..., description="Semantic type (histogram/counter/gauge)"
    )
    unit: str = Field(default="", description="Human-readable unit")
    description: str = Field(default="", description="Short description of the metric")

    model_config = ConfigDict(frozen=True, str_strip_whitespace=True)


#: Registry mapping every ``OperationalMetricKey`` to its definition.
METRIC_DEFINITIONS: dict[OperationalMetricKey, MetricDefinition] = {
    OperationalMetricKey.PLAN_DURATION_MS: MetricDefinition(
        key=OperationalMetricKey.PLAN_DURATION_MS,
        metric_type=MetricType.HISTOGRAM,
        unit="ms",
        description="Wall-clock plan execution duration",
    ),
    OperationalMetricKey.PLAN_TOTAL_COST_USD: MetricDefinition(
        key=OperationalMetricKey.PLAN_TOTAL_COST_USD,
        metric_type=MetricType.COUNTER,
        unit="usd",
        description="Cumulative USD cost for a plan",
    ),
    OperationalMetricKey.PLAN_DECISION_COUNT: MetricDefinition(
        key=OperationalMetricKey.PLAN_DECISION_COUNT,
        metric_type=MetricType.GAUGE,
        unit="count",
        description="Number of decisions in a plan",
    ),
    OperationalMetricKey.SUBPLAN_COUNT: MetricDefinition(
        key=OperationalMetricKey.SUBPLAN_COUNT,
        metric_type=MetricType.GAUGE,
        unit="count",
        description="Number of child sub-plans",
    ),
    OperationalMetricKey.ACTOR_INVOCATION_COUNT: MetricDefinition(
        key=OperationalMetricKey.ACTOR_INVOCATION_COUNT,
        metric_type=MetricType.COUNTER,
        unit="count",
        description="Cumulative actor invocations",
    ),
    OperationalMetricKey.ACTOR_LATENCY_MS: MetricDefinition(
        key=OperationalMetricKey.ACTOR_LATENCY_MS,
        metric_type=MetricType.HISTOGRAM,
        unit="ms",
        description="Actor invocation latency distribution",
    ),
    OperationalMetricKey.TOOL_INVOCATION_COUNT: MetricDefinition(
        key=OperationalMetricKey.TOOL_INVOCATION_COUNT,
        metric_type=MetricType.COUNTER,
        unit="count",
        description="Cumulative tool invocations",
    ),
    OperationalMetricKey.TOOL_ERROR_RATE: MetricDefinition(
        key=OperationalMetricKey.TOOL_ERROR_RATE,
        metric_type=MetricType.GAUGE,
        unit="ratio",
        description="Tool error rate (0.0 to 1.0)",
    ),
    OperationalMetricKey.CONTEXT_BUILD_TIME_MS: MetricDefinition(
        key=OperationalMetricKey.CONTEXT_BUILD_TIME_MS,
        metric_type=MetricType.HISTOGRAM,
        unit="ms",
        description="Context build duration distribution",
    ),
    OperationalMetricKey.CONTEXT_TOKEN_COUNT: MetricDefinition(
        key=OperationalMetricKey.CONTEXT_TOKEN_COUNT,
        metric_type=MetricType.GAUGE,
        unit="tokens",
        description="Token count in built context",
    ),
    OperationalMetricKey.LLM_CALL_COUNT: MetricDefinition(
        key=OperationalMetricKey.LLM_CALL_COUNT,
        metric_type=MetricType.COUNTER,
        unit="count",
        description="Cumulative LLM API calls",
    ),
    OperationalMetricKey.LLM_TOTAL_TOKENS: MetricDefinition(
        key=OperationalMetricKey.LLM_TOTAL_TOKENS,
        metric_type=MetricType.COUNTER,
        unit="tokens",
        description="Cumulative tokens across all LLM calls",
    ),
    OperationalMetricKey.LLM_TOTAL_COST_USD: MetricDefinition(
        key=OperationalMetricKey.LLM_TOTAL_COST_USD,
        metric_type=MetricType.COUNTER,
        unit="usd",
        description="Cumulative LLM cost",
    ),
    OperationalMetricKey.LLM_AVG_LATENCY_MS: MetricDefinition(
        key=OperationalMetricKey.LLM_AVG_LATENCY_MS,
        metric_type=MetricType.GAUGE,
        unit="ms",
        description="Average LLM call latency",
    ),
}


class MetricEntry(BaseModel):
    """A single metric observation.

    Each entry records the key, numeric value, plan context, timestamp,
    and optional labels for dimensional filtering.

    Attributes:
        key: The operational metric being recorded.
        value: Numeric value of the observation.
        plan_id: ULID of the plan this metric belongs to.
        timestamp: When the observation was recorded.
        labels: Arbitrary dimensional labels for drill-down
            (e.g. ``{"actor": "planner", "provider": "openai"}``).
        metric_type: Resolved metric type from the definition registry.
    """

    key: OperationalMetricKey = Field(
        ...,
        description="The operational metric being recorded",
    )
    value: float = Field(
        ...,
        description="Numeric value of the observation",
    )
    plan_id: str = Field(
        ...,
        min_length=26,
        max_length=26,
        description="ULID of the plan this metric belongs to",
    )
    timestamp: datetime = Field(
        default_factory=datetime.utcnow,
        description="When the observation was recorded",
    )
    labels: dict[str, Any] = Field(
        default_factory=dict,
        description="Dimensional labels for drill-down filtering",
    )
    metric_type: MetricType | None = Field(
        default=None,
        description="Resolved semantic type (histogram/counter/gauge)",
    )

    model_config = ConfigDict(
        frozen=True,
        str_strip_whitespace=True,
    )


class MetricCollector:
    """Stateless factory for creating ``MetricEntry`` objects.

    Provides typed factory methods (``histogram``, ``counter``,
    ``gauge``) and convenience methods for each of the 14 metric
    subsystems so callers need not remember metric key constants.
    """

    @staticmethod
    def record(
        key: OperationalMetricKey,
        value: float,
        plan_id: str,
        *,
        labels: dict[str, Any] | None = None,
        timestamp: datetime | None = None,
    ) -> MetricEntry:
        """Create a ``MetricEntry`` for the given key and value.

        Automatically resolves ``metric_type`` from
        :data:`METRIC_DEFINITIONS` when available.

        Args:
            key: The operational metric key.
            value: Numeric observation value.
            plan_id: ULID of the owning plan.
            labels: Optional dimensional labels.
            timestamp: Override timestamp (defaults to ``utcnow``).

        Returns:
            A new ``MetricEntry`` instance.
        """
        defn = METRIC_DEFINITIONS.get(key)
        resolved_type = defn.metric_type if defn is not None else None
        kwargs: dict[str, Any] = {
            "key": key,
            "value": value,
            "plan_id": plan_id,
            "labels": labels or {},
            "metric_type": resolved_type,
        }
        if timestamp is not None:
            kwargs["timestamp"] = timestamp
        return MetricEntry(**kwargs)

    # ----- typed factory methods -----

    @staticmethod
    def histogram(
        key: OperationalMetricKey,
        value: float,
        plan_id: str,
        **labels: Any,
    ) -> MetricEntry:
        """Create a histogram metric entry.

        Args:
            key: The operational metric key.
            value: Observed value for the histogram bucket.
            plan_id: ULID of the owning plan.
            **labels: Dimensional labels.

        Returns:
            A ``MetricEntry`` with ``metric_type=HISTOGRAM``.
        """
        return MetricCollector.record(key, value, plan_id, labels=dict(labels))

    @staticmethod
    def counter(
        key: OperationalMetricKey,
        value: float,
        plan_id: str,
        **labels: Any,
    ) -> MetricEntry:
        """Create a counter metric entry.

        Args:
            key: The operational metric key.
            value: Increment value for the counter.
            plan_id: ULID of the owning plan.
            **labels: Dimensional labels.

        Returns:
            A ``MetricEntry`` with ``metric_type=COUNTER``.
        """
        return MetricCollector.record(key, value, plan_id, labels=dict(labels))

    @staticmethod
    def gauge(
        key: OperationalMetricKey,
        value: float,
        plan_id: str,
        **labels: Any,
    ) -> MetricEntry:
        """Create a gauge metric entry.

        Args:
            key: The operational metric key.
            value: Current gauge value.
            plan_id: ULID of the owning plan.
            **labels: Dimensional labels.

        Returns:
            A ``MetricEntry`` with ``metric_type=GAUGE``.
        """
        return MetricCollector.record(key, value, plan_id, labels=dict(labels))

    # ----- convenience helpers (all 14 metrics) -----

    @staticmethod
    def plan_duration(plan_id: str, duration_ms: float, **labels: Any) -> MetricEntry:
        """Record ``PLAN_DURATION_MS``."""
        return MetricCollector.record(
            OperationalMetricKey.PLAN_DURATION_MS,
            duration_ms,
            plan_id,
            labels=dict(labels),
        )

    @staticmethod
    def plan_cost(plan_id: str, cost_usd: float, **labels: Any) -> MetricEntry:
        """Record ``PLAN_TOTAL_COST_USD``."""
        return MetricCollector.record(
            OperationalMetricKey.PLAN_TOTAL_COST_USD,
            cost_usd,
            plan_id,
            labels=dict(labels),
        )

    @staticmethod
    def plan_decision_count(plan_id: str, count: float, **labels: Any) -> MetricEntry:
        """Record ``PLAN_DECISION_COUNT``."""
        return MetricCollector.record(
            OperationalMetricKey.PLAN_DECISION_COUNT,
            count,
            plan_id,
            labels=dict(labels),
        )

    @staticmethod
    def subplan_count(plan_id: str, count: float, **labels: Any) -> MetricEntry:
        """Record ``SUBPLAN_COUNT``."""
        return MetricCollector.record(
            OperationalMetricKey.SUBPLAN_COUNT,
            count,
            plan_id,
            labels=dict(labels),
        )

    @staticmethod
    def actor_invocation_count(
        plan_id: str, count: float, **labels: Any
    ) -> MetricEntry:
        """Record ``ACTOR_INVOCATION_COUNT``."""
        return MetricCollector.record(
            OperationalMetricKey.ACTOR_INVOCATION_COUNT,
            count,
            plan_id,
            labels=dict(labels),
        )

    @staticmethod
    def actor_latency(plan_id: str, latency_ms: float, **labels: Any) -> MetricEntry:
        """Record ``ACTOR_LATENCY_MS``."""
        return MetricCollector.record(
            OperationalMetricKey.ACTOR_LATENCY_MS,
            latency_ms,
            plan_id,
            labels=dict(labels),
        )

    @staticmethod
    def tool_invocation_count(plan_id: str, count: float, **labels: Any) -> MetricEntry:
        """Record ``TOOL_INVOCATION_COUNT``."""
        return MetricCollector.record(
            OperationalMetricKey.TOOL_INVOCATION_COUNT,
            count,
            plan_id,
            labels=dict(labels),
        )

    @staticmethod
    def tool_error_rate(plan_id: str, rate: float, **labels: Any) -> MetricEntry:
        """Record ``TOOL_ERROR_RATE``."""
        return MetricCollector.record(
            OperationalMetricKey.TOOL_ERROR_RATE,
            rate,
            plan_id,
            labels=dict(labels),
        )

    @staticmethod
    def context_build_time(
        plan_id: str, duration_ms: float, **labels: Any
    ) -> MetricEntry:
        """Record ``CONTEXT_BUILD_TIME_MS``."""
        return MetricCollector.record(
            OperationalMetricKey.CONTEXT_BUILD_TIME_MS,
            duration_ms,
            plan_id,
            labels=dict(labels),
        )

    @staticmethod
    def context_token_count(plan_id: str, count: float, **labels: Any) -> MetricEntry:
        """Record ``CONTEXT_TOKEN_COUNT``."""
        return MetricCollector.record(
            OperationalMetricKey.CONTEXT_TOKEN_COUNT,
            count,
            plan_id,
            labels=dict(labels),
        )

    @staticmethod
    def llm_call_count(plan_id: str, count: float, **labels: Any) -> MetricEntry:
        """Record ``LLM_CALL_COUNT``."""
        return MetricCollector.record(
            OperationalMetricKey.LLM_CALL_COUNT,
            count,
            plan_id,
            labels=dict(labels),
        )

    @staticmethod
    def llm_total_tokens(plan_id: str, count: float, **labels: Any) -> MetricEntry:
        """Record ``LLM_TOTAL_TOKENS``."""
        return MetricCollector.record(
            OperationalMetricKey.LLM_TOTAL_TOKENS,
            count,
            plan_id,
            labels=dict(labels),
        )

    @staticmethod
    def llm_total_cost(plan_id: str, cost_usd: float, **labels: Any) -> MetricEntry:
        """Record ``LLM_TOTAL_COST_USD``."""
        return MetricCollector.record(
            OperationalMetricKey.LLM_TOTAL_COST_USD,
            cost_usd,
            plan_id,
            labels=dict(labels),
        )

    @staticmethod
    def llm_avg_latency(plan_id: str, latency_ms: float, **labels: Any) -> MetricEntry:
        """Record ``LLM_AVG_LATENCY_MS``."""
        return MetricCollector.record(
            OperationalMetricKey.LLM_AVG_LATENCY_MS,
            latency_ms,
            plan_id,
            labels=dict(labels),
        )
