//! gRPC metrics Tower layer for automatic request instrumentation.
//!
//! Provides `grpc_server_started_total`, `grpc_server_handled_total`,
//! and `grpc_server_handling_seconds` metrics for all gRPC handlers.
//!
//! # Usage
//!
//! ```rust,ignore
//! use common::metrics::grpc_metrics::GrpcMetricsLayer;
//! use tonic::transport::Server;
//!
//! Server::builder()
//!     .layer(GrpcMetricsLayer::new("trading_service"))
//!     .add_service(my_service)
//!     .serve(addr)
//!     .await?;
//! ```

use once_cell::sync::Lazy;
use prometheus::{register_counter_vec, register_histogram_vec, CounterVec, HistogramVec};
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use tower_layer::Layer;
use tower_service::Service;

/// Counter for gRPC requests started.
static GRPC_STARTED: Lazy<CounterVec> = Lazy::new(|| {
    register_counter_vec!(
        "grpc_server_started_total",
        "Total gRPC RPCs started on the server",
        &["service", "grpc_method"]
    )
    .unwrap_or_else(|e| {
        // Metric registration failure at startup is unrecoverable.
        // The process cannot operate without metrics instrumentation.
        eprintln!("FATAL: failed to register grpc_server_started_total: {e}");
        std::process::abort()
    })
});

/// Counter for gRPC requests completed with status code.
static GRPC_HANDLED: Lazy<CounterVec> = Lazy::new(|| {
    register_counter_vec!(
        "grpc_server_handled_total",
        "Total gRPC RPCs completed on the server with status code",
        &["service", "grpc_method", "grpc_code"]
    )
    .unwrap_or_else(|e| {
        eprintln!("FATAL: failed to register grpc_server_handled_total: {e}");
        std::process::abort()
    })
});

/// Histogram for gRPC request handling duration.
static GRPC_HANDLING_SECONDS: Lazy<HistogramVec> = Lazy::new(|| {
    register_histogram_vec!(
        "grpc_server_handling_seconds",
        "Histogram of gRPC request handling duration in seconds",
        &["service", "grpc_method"],
        vec![0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 5.0]
    )
    .unwrap_or_else(|e| {
        eprintln!("FATAL: failed to register grpc_server_handling_seconds: {e}");
        std::process::abort()
    })
});

/// Initialize gRPC metrics (force lazy registration).
///
/// Call this early in service startup to ensure metrics are registered
/// before any requests arrive. It is safe to call multiple times.
pub fn init_grpc_metrics() {
    let _ = &*GRPC_STARTED;
    let _ = &*GRPC_HANDLED;
    let _ = &*GRPC_HANDLING_SECONDS;
}

/// Extract the method name from a gRPC URI path.
///
/// gRPC paths follow the format `/package.Service/Method`.
/// This extracts the last segment (the method name).
fn extract_method(path: &str) -> &str {
    match path.rsplit('/').next() {
        Some(m) if !m.is_empty() => m,
        _ => "unknown",
    }
}

/// Tower Layer that instruments gRPC handlers with Prometheus metrics.
///
/// Records three metric families:
/// - `grpc_server_started_total` — incremented when a request begins
/// - `grpc_server_handled_total` — incremented when a request completes (with status code)
/// - `grpc_server_handling_seconds` — histogram of request duration
#[derive(Debug, Clone)]
pub struct GrpcMetricsLayer {
    service_name: &'static str,
}

impl GrpcMetricsLayer {
    /// Create a new gRPC metrics layer for the given service.
    ///
    /// `service_name` is used as the `service` label on all emitted metrics.
    pub fn new(service_name: &'static str) -> Self {
        init_grpc_metrics();
        Self { service_name }
    }
}

impl<S> Layer<S> for GrpcMetricsLayer {
    type Service = GrpcMetricsService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        GrpcMetricsService {
            inner,
            service_name: self.service_name,
        }
    }
}

/// Tower Service that records gRPC metrics around the inner service.
#[derive(Debug, Clone)]
pub struct GrpcMetricsService<S> {
    inner: S,
    service_name: &'static str,
}

impl<S, ReqBody, ResBody> Service<http::Request<ReqBody>> for GrpcMetricsService<S>
where
    S: Service<http::Request<ReqBody>, Response = http::Response<ResBody>>
        + Clone
        + Send
        + 'static,
    S::Future: Send + 'static,
    S::Error: Send + 'static,
    ReqBody: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: http::Request<ReqBody>) -> Self::Future {
        let service_name = self.service_name;
        let method = extract_method(req.uri().path()).to_string();

        GRPC_STARTED
            .with_label_values(&[service_name, &method])
            .inc();

        let start = std::time::Instant::now();
        // Clone inner before calling to satisfy tower's Service contract:
        // poll_ready was called on `self`, so we must call on `self`
        // (not the clone). We swap so the clone becomes `self` for
        // the next call, and we consume the ready instance.
        let mut inner = self.inner.clone();
        std::mem::swap(&mut self.inner, &mut inner);

        Box::pin(async move {
            let result = inner.call(req).await;
            let elapsed = start.elapsed().as_secs_f64();

            let code = match &result {
                Ok(resp) => resp
                    .headers()
                    .get("grpc-status")
                    .and_then(|v| v.to_str().ok())
                    .map_or_else(|| "0".to_string(), |s| s.to_string()),
                Err(_) => "13".to_string(), // INTERNAL
            };

            GRPC_HANDLED
                .with_label_values(&[service_name, &method, &code])
                .inc();

            GRPC_HANDLING_SECONDS
                .with_label_values(&[service_name, &method])
                .observe(elapsed);

            result
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_extract_method_normal_path() {
        assert_eq!(extract_method("/package.Service/MethodName"), "MethodName");
    }

    #[test]
    fn test_extract_method_simple_path() {
        assert_eq!(extract_method("/Method"), "Method");
    }

    #[test]
    fn test_extract_method_empty_path() {
        assert_eq!(extract_method(""), "unknown");
    }

    #[test]
    fn test_extract_method_trailing_slash() {
        // rsplit('/').next() on "foo/" yields "" (empty), so fallback
        assert_eq!(extract_method("/Service/"), "unknown");
    }

    #[test]
    fn test_extract_method_root_slash() {
        assert_eq!(extract_method("/"), "unknown");
    }

    #[test]
    fn test_metrics_layer_creates_service() {
        let layer = GrpcMetricsLayer::new("test_svc");
        assert_eq!(layer.service_name, "test_svc");
    }

    #[test]
    fn test_init_grpc_metrics_is_idempotent() {
        // Calling init multiple times should not panic
        init_grpc_metrics();
        init_grpc_metrics();
        init_grpc_metrics();
    }

    #[test]
    fn test_metrics_registered() {
        init_grpc_metrics();
        // Verify metrics are accessible via the global statics
        GRPC_STARTED
            .with_label_values(&["test_svc", "TestMethod"])
            .inc();
        let val = GRPC_STARTED
            .with_label_values(&["test_svc", "TestMethod"])
            .get();
        assert!(val >= 1.0, "Counter should have been incremented");
    }
}
