Monitoring & Observability

use axum::{extract::Extension, http::StatusCode, response::Response, routing::get, Router};
use prometheus::{Counter, Histogram, Gauge, Registry, Encoder, TextEncoder, opts};
use scoutquest_rust::ServiceDiscoveryClient;
use std::sync::Arc;
use std::time::Instant;

#[derive(Clone)]
struct Metrics {
    http_requests_total: Counter,
    http_request_duration: Histogram,
    service_discovery_duration: Histogram,
    active_connections: Gauge,
    registry: Registry,
}

impl Metrics {
    fn new() -> Result> {
        let registry = Registry::new();

        let http_requests_total = Counter::with_opts(
            opts!("http_requests_total", "Total number of HTTP requests")
                .variable_labels(vec!["method".to_string(), "route".to_string(), "status_code".to_string()])
        )?;

        let http_request_duration = Histogram::with_opts(
            prometheus::HistogramOpts::new("http_request_duration_seconds", "Duration of HTTP requests")
                .variable_labels(vec!["method".to_string(), "route".to_string()])
                .buckets(vec![0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0])
        )?;

        let service_discovery_duration = Histogram::with_opts(
            prometheus::HistogramOpts::new("service_discovery_duration_seconds", "Service discovery request duration")
                .variable_labels(vec!["service_name".to_string(), "operation".to_string()])
        )?;

        let active_connections = Gauge::with_opts(
            opts!("active_connections", "Number of active connections")
        )?;

        registry.register(Box::new(http_requests_total.clone()))?;
        registry.register(Box::new(http_request_duration.clone()))?;
        registry.register(Box::new(service_discovery_duration.clone()))?;
        registry.register(Box::new(active_connections.clone()))?;

        Ok(Metrics {
            http_requests_total,
            http_request_duration,
            service_discovery_duration,
            active_connections,
            registry,
        })
    }

    async fn instrumented_service_call(
        &self,
        client: &ServiceDiscoveryClient,
        service_name: &str,
        path: &str
    ) -> Result> {
        let start = Instant::now();

        match client.get_service(service_name, path).await {
            Ok(result) => {
                self.service_discovery_duration
                    .with_label_values(&[service_name, "get_service"])
                    .observe(start.elapsed().as_secs_f64());
                Ok(result)
            }
            Err(e) => {
                self.service_discovery_duration
                    .with_label_values(&[service_name, "get_service_error"])
                    .observe(start.elapsed().as_secs_f64());
                Err(e.into())
            }
        }
    }
}

async fn metrics_handler(Extension(metrics): Extension) -> Response {
    let encoder = TextEncoder::new();
    let metric_families = metrics.registry.gather();

    let mut buffer = Vec::new();
    encoder.encode(&metric_families, &mut buffer).unwrap();

    Response::builder()
        .status(StatusCode::OK)
        .header("Content-Type", encoder.format_type())
        .body(buffer.into())
        .unwrap()
}

#[tokio::main]
async fn main() -> Result<(), Box> {
    let metrics = Arc::new(Metrics::new()?);
    let client = Arc::new(ServiceDiscoveryClient::new("http://localhost:8080")?);

    // Background task to update connection metrics
    let metrics_clone = metrics.clone();
    tokio::spawn(async move {
        let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(5));
        loop {
            interval.tick().await;
            let connections = get_active_connection_count().await;
            metrics_clone.active_connections.set(connections as f64);
        }
    });

    let app = Router::new()
        .route("/metrics", get(metrics_handler))
        .layer(Extension(metrics))
        .layer(Extension(client));

    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await?;
    axum::serve(listener, app).await?;

    Ok(())
}

async fn get_active_connection_count() -> u64 {
    // Implementation to count active connections
    42
}

use serde_json::json;
use tracing::{info, error, warn, debug, instrument};
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
use uuid::Uuid;
use scoutquest_rust::ServiceDiscoveryClient;

// Initialize structured logging
pub fn init_logging() -> Result<(), Box> {
    tracing_subscriber::registry()
        .with(EnvFilter::from_default_env().add_directive("scoutquest=info".parse()?))
        .with(
            tracing_subscriber::fmt::layer()
                .json()
                .with_target(true)
                .with_current_span(false)
                .with_span_list(true)
        )
        .init();

    Ok(())
}

#[derive(Clone)]
pub struct ServiceLogger {
    service_name: String,
    version: String,
}

impl ServiceLogger {
    pub fn new(service_name: &str, version: &str) -> Self {
        Self {
            service_name: service_name.to_string(),
            version: version.to_string(),
        }
    }

    #[instrument(skip(self), fields(service = %self.service_name, version = %self.version))]
    pub async fn discover_service_with_logging(
        &self,
        client: &ServiceDiscoveryClient,
        service_name: &str,
        trace_id: &str
    ) -> Result> {
        info!(
            service_name = %service_name,
            trace_id = %trace_id,
            operation = "service_discovery",
            "Discovering service"
        );

        let start_time = std::time::Instant::now();

        match client.discover_service(service_name).await {
            Ok(instance) => {
                let duration = start_time.elapsed();

                info!(
                    service_name = %service_name,
                    trace_id = %trace_id,
                    instance_id = %instance.id,
                    host = %instance.host,
                    port = instance.port,
                    duration_ms = duration.as_millis(),
                    operation = "service_discovery",
                    "Service discovered successfully"
                );

                Ok(instance)
            }
            Err(e) => {
                let duration = start_time.elapsed();

                error!(
                    service_name = %service_name,
                    trace_id = %trace_id,
                    error = %e,
                    duration_ms = duration.as_millis(),
                    operation = "service_discovery",
                    "Service discovery failed"
                );

                Err(e.into())
            }
        }
    }

    #[instrument(skip(self), fields(service = %self.service_name))]
    pub fn log_service_registration(&self, instance_id: &str, host: &str, port: u16) {
        info!(
            instance_id = %instance_id,
            host = %host,
            port = port,
            operation = "service_registration",
            "Service registered successfully"
        );
    }

    #[instrument(skip(self), fields(service = %self.service_name))]
    pub fn log_health_check_result(&self, status: &str, response_time_ms: u64) {
        info!(
            status = %status,
            response_time_ms = response_time_ms,
            operation = "health_check",
            "Health check completed"
        );
    }
}

// Request tracing middleware for Axum
use axum::{extract::Request, middleware::Next, response::Response};

pub async fn request_tracing_middleware(
    mut request: Request,
    next: Next,
) -> Response {
    let trace_id = request.headers()
        .get("x-trace-id")
        .and_then(|h| h.to_str().ok())
        .unwrap_or_else(|| {
            let id = Uuid::new_v4().to_string();
            request.headers_mut().insert(
                "x-trace-id",
                id.parse().unwrap()
            );
            request.headers().get("x-trace-id").unwrap().to_str().unwrap()
        });

    let method = request.method().clone();
    let uri = request.uri().clone();
    let start_time = std::time::Instant::now();

    info!(
        method = %method,
        uri = %uri,
        trace_id = %trace_id,
        operation = "http_request",
        "HTTP request started"
    );

    let response = next.run(request).await;

    let duration = start_time.elapsed();
    let status = response.status();

    info!(
        method = %method,
        uri = %uri,
        status_code = status.as_u16(),
        duration_ms = duration.as_millis(),
        trace_id = %trace_id,
        operation = "http_request",
        "HTTP request completed"
    );

    response
}

#[tokio::main]
async fn main() -> Result<(), Box> {
    init_logging()?;

    let service_logger = ServiceLogger::new("user-service", "1.0.0");
    let client = ServiceDiscoveryClient::new("http://localhost:8080")?;

    // Example usage
    let trace_id = Uuid::new_v4().to_string();
    let instance = service_logger
        .discover_service_with_logging(&client, "auth-service", &trace_id)
        .await?;

    Ok(())
}

use opentelemetry::{global, trace::TraceError, KeyValue};
use opentelemetry_jaeger::JaegerTraceExporter;
use opentelemetry_sdk::{trace as sdktrace, Resource};
use tracing::{info, instrument, Span};
use tracing_opentelemetry::OpenTelemetryLayer;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
use scoutquest_rust::ServiceDiscoveryClient;

// Initialize OpenTelemetry with Jaeger
fn init_tracer() -> Result {
    let jaeger_exporter = JaegerTraceExporter::builder()
        .with_agent_endpoint("http://localhost:14268/api/traces")
        .build()?;

    let tracer = opentelemetry_sdk::trace::TracerProvider::builder()
        .with_batch_exporter(jaeger_exporter, opentelemetry_sdk::runtime::Tokio)
        .with_resource(Resource::new(vec![
            KeyValue::new("service.name", "user-service"),
            KeyValue::new("service.version", "1.0.0"),
            KeyValue::new("deployment.environment", "production"),
        ]))
        .build()
        .tracer("user-service");

    Ok(tracer)
}

// Initialize tracing with OpenTelemetry
pub fn init_tracing() -> Result<(), Box> {
    let tracer = init_tracer()?;

    tracing_subscriber::registry()
        .with(tracing_subscriber::EnvFilter::from_default_env())
        .with(tracing_subscriber::fmt::layer())
        .with(OpenTelemetryLayer::new(tracer))
        .init();

    Ok(())
}

// Enhanced service client with tracing
pub struct TracingServiceDiscoveryClient {
    client: ServiceDiscoveryClient,
}

impl TracingServiceDiscoveryClient {
    pub fn new(server_url: &str) -> Result> {
        Ok(Self {
            client: ServiceDiscoveryClient::new(server_url)?,
        })
    }

    #[instrument(
        skip(self),
        fields(
            service.name = %service_name,
            service.discovery.method = "discover_service"
        )
    )]
    pub async fn discover_service(
        &self,
        service_name: &str,
    ) -> Result> {
        let span = Span::current();

        span.record("service.discovery.service_name", service_name);

        match self.client.discover_service(service_name).await {
            Ok(instance) => {
                span.record("service.discovery.instance_id", &instance.id);
                span.record("service.discovery.host", &instance.host);
                span.record("service.discovery.port", instance.port);
                span.record("service.discovery.status", "success");

                info!(
                    service_name = %service_name,
                    instance_id = %instance.id,
                    "Service discovered successfully"
                );

                Ok(instance)
            }
            Err(e) => {
                span.record("service.discovery.status", "error");
                span.record("error.message", &e.to_string());

                tracing::error!(
                    service_name = %service_name,
                    error = %e,
                    "Service discovery failed"
                );

                Err(e.into())
            }
        }
    }

    #[instrument(
        skip(self),
        fields(
            http.method = "GET",
            http.url = %path,
            service.name = %service_name,
            service.discovery.method = "get_service"
        )
    )]
    pub async fn get_service(
        &self,
        service_name: &str,
        path: &str,
    ) -> Result> {
        let span = Span::current();

        match self.client.get_service(service_name, path).await {
            Ok(response) => {
                span.record("http.status_code", 200);
                span.record("service.discovery.status", "success");

                info!(
                    service_name = %service_name,
                    path = %path,
                    "Service call successful"
                );

                Ok(response)
            }
            Err(e) => {
                span.record("http.status_code", 500);
                span.record("service.discovery.status", "error");
                span.record("error.message", &e.to_string());

                tracing::error!(
                    service_name = %service_name,
                    path = %path,
                    error = %e,
                    "Service call failed"
                );

                Err(e.into())
            }
        }
    }
}

// Business logic with custom spans
#[instrument(
    fields(
        user.id = %user_id,
        order.id = %order_id,
        business.operation = "process_order"
    )
)]
pub async fn process_user_order(
    client: &TracingServiceDiscoveryClient,
    user_id: &str,
    order_id: &str,
) -> Result> {
    let span = Span::current();

    info!("Processing user order");

    // These calls will be automatically traced
    let user_response = client
        .get_service("user-service", &format!("/users/{}", user_id))
        .await?;

    let order_response = client
        .get_service("order-service", &format!("/orders/{}", order_id))
        .await?;

    // Add custom attributes to span
    span.record("user.data", &user_response[..100.min(user_response.len())]);
    span.record("order.data", &order_response[..100.min(order_response.len())]);

    info!("Order processed successfully");

    Ok("Order processed".to_string())
}

#[tokio::main]
async fn main() -> Result<(), Box> {
    init_tracing()?;

    let client = TracingServiceDiscoveryClient::new("http://localhost:8080")?;

    // Example usage
    let result = process_user_order(&client, "user123", "order456").await?;
    println!("Result: {}", result);

    // Ensure all spans are exported
    global::shutdown_tracer_provider();

    Ok(())
}