from prometheus_client import Counter, Histogram, Gauge, generate_latest, CONTENT_TYPE_LATEST
from fastapi import Response
import time

# Define Prometheus metrics
prediction_counter = Counter(
    'predictions_total',
    'Total number of predictions made',
    ['model_version', 'status']
)

prediction_duration = Histogram(
    'prediction_duration_seconds',
    'Time spent processing prediction',
    buckets=[0.1, 0.5, 1.0, 2.0, 5.0]
)

model_accuracy = Gauge(
    'model_accuracy',
    'Current model accuracy',
    ['model_version']
)

data_drift_detected = Gauge(
    'data_drift_detected',
    'Whether data drift has been detected (1=yes, 0=no)'
)

active_requests = Gauge(
    'active_requests',
    'Number of active requests'
)


class MetricsMiddleware:
    """Middleware to track request metrics"""
    
    async def __call__(self, request, call_next):
        active_requests.inc()
        start_time = time.time()
        
        try:
            response = await call_next(request)
            return response
        finally:
            active_requests.dec()
            duration = time.time() - start_time
            if request.url.path == "/predict/":
                prediction_duration.observe(duration)


def get_metrics() -> Response:
    """Endpoint to expose Prometheus metrics"""
    return Response(content=generate_latest(), media_type=CONTENT_TYPE_LATEST)