"""
FoundryMatch SaaS - Model Monitoring
====================================
Provides a wrapper class to monitor machine learning models using
Prometheus metrics. This allows for real-time tracking of model
performance, latency, and usage.
"""

import time
from typing import Any

from prometheus_client import Counter, Histogram, Gauge

# --- Prometheus Metrics Definitions ---

# 1. Counter for total predictions, labeled by model, version, and outcome.
MODEL_PREDICTIONS_TOTAL = Counter(
    "fmatch_model_predictions_total",
    "Total number of predictions made by a model.",
    ["model_name", "model_version", "outcome"],
)

# 2. Histogram for prediction latency.
MODEL_PREDICTION_LATENCY_SECONDS = Histogram(
    "fmatch_model_prediction_latency_seconds",
    "Latency of model predictions in seconds.",
    ["model_name", "model_version"],
    buckets=(0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0),
)

# 3. Histogram for prediction confidence scores.
MODEL_PREDICTION_CONFIDENCE = Histogram(
    "fmatch_model_prediction_confidence",
    "Distribution of model prediction confidence scores.",
    ["model_name", "model_version"],
    buckets=(0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0),
)

# 4. Gauge for the last time a model was used.
MODEL_LAST_USED_TIMESTAMP = Gauge(
    "fmatch_model_last_used_timestamp_seconds",
    "Timestamp of the last time a model was used for a prediction.",
    ["model_name", "model_version"],
)


class MonitoredModel:
    """
    A wrapper class for machine learning models that adds Prometheus monitoring
    for predictions, latency, and confidence scores.
    """

    def __init__(self, model: Any, model_name: str, model_version: str):
        """
        Initializes the monitored model wrapper.

        Args:
            model: The actual machine learning model object (e.g., a scikit-learn classifier).
                   The model must have a `predict` method and optionally a `predict_proba` method.
            model_name: The name of the model for labeling metrics.
            model_version: The version of the model for labeling metrics.
        """
        if not hasattr(model, "predict"):
            raise TypeError(
                f"The provided model of type '{type(model).__name__}' must have a `predict` method."
            )

        self._model = model
        self.model_name = model_name
        self.model_version = model_version
        self._labels = {
            "model_name": self.model_name,
            "model_version": self.model_version,
        }

    def predict(self, *args, **kwargs) -> Any:
        """
        Wraps the model's `predict` method to record performance metrics.
        """
        start_time = time.perf_counter()
        try:
            result = self._model.predict(*args, **kwargs)
            MODEL_PREDICTIONS_TOTAL.labels(**self._labels, outcome="success").inc()
            MODEL_LAST_USED_TIMESTAMP.labels(**self._labels).set_to_current_time()
            return result
        except Exception as e:
            MODEL_PREDICTIONS_TOTAL.labels(**self._labels, outcome="failure").inc()
            raise e
        finally:
            latency = time.perf_counter() - start_time
            MODEL_PREDICTION_LATENCY_SECONDS.labels(**self._labels).observe(latency)

    def predict_proba(self, *args, **kwargs) -> Any:
        """
        Wraps the model's `predict_proba` method to record confidence scores.
        """
        if not hasattr(self._model, "predict_proba"):
            raise AttributeError(
                f"The wrapped model '{type(self._model).__name__}' has no 'predict_proba' method."
            )

        start_time = time.perf_counter()
        try:
            probabilities = self._model.predict_proba(*args, **kwargs)
            if hasattr(probabilities, "max"):
                max_probs = probabilities.max(axis=1)
                for prob in max_probs:
                    MODEL_PREDICTION_CONFIDENCE.labels(**self._labels).observe(prob)
            MODEL_PREDICTIONS_TOTAL.labels(**self._labels, outcome="success").inc()
            MODEL_LAST_USED_TIMESTAMP.labels(**self._labels).set_to_current_time()
            return probabilities
        except Exception as e:
            MODEL_PREDICTIONS_TOTAL.labels(**self._labels, outcome="failure").inc()
            raise e
        finally:
            latency = time.perf_counter() - start_time
            MODEL_PREDICTION_LATENCY_SECONDS.labels(**self._labels).observe(latency)

    def __getattr__(self, name: str) -> Any:
        """
        Proxy any other attribute access to the wrapped model, making the
        wrapper transparent for other model methods or attributes.
        """
        return getattr(self._model, name)
