//! Correlation ID generation and propagation for distributed tracing.
//!
//! This module provides utilities for generating unique correlation IDs and
//! propagating them across service boundaries via gRPC metadata. Correlation IDs
//! enable end-to-end request tracking in a microservices architecture.
//!
//! # Architecture
//!
//! - Correlation IDs are generated as UUIDv4 strings at the API Gateway
//! - IDs are propagated via gRPC metadata using the `x-correlation-id` header
//! - Thread-local storage is used to maintain correlation context in async tasks
//! - All log statements automatically include the current correlation ID
//!
//! # Example
//!
//! ```rust
//! use common::observability::CorrelationId;
//! use tracing::info;
//!
//! let correlation_id = CorrelationId::new();
//!
//! info!(
//!     correlation_id = %correlation_id,
//!     "Processing request"
//! );
//! ```

use std::fmt;
use std::sync::Arc;
use tokio::sync::RwLock;
use uuid::Uuid;

/// The gRPC metadata header key for correlation IDs.
pub const CORRELATION_ID_HEADER: &str = "x-correlation-id";

// Thread-local storage for the current correlation ID.
//
// This allows correlation IDs to be automatically available in any function
// within the current async task context without explicit parameter passing.
tokio::task_local! {
    pub static CURRENT_CORRELATION_ID: Arc<RwLock<Option<CorrelationId>>>;
}

/// A unique identifier that tracks a request across service boundaries.
///
/// Correlation IDs are generated as UUIDv4 strings and propagated via gRPC
/// metadata. They enable end-to-end request tracking and are automatically
/// included in all log statements.
///
/// # Example
///
/// ```rust
/// use common::observability::CorrelationId;
///
/// // Generate a new correlation ID
/// let id = CorrelationId::new();
///
/// // Parse from string (e.g., from gRPC metadata)
/// let id = CorrelationId::from_string("550e8400-e29b-41d4-a716-446655440000").unwrap();
///
/// // Convert to string for logging
/// println!("Correlation ID: {}", id);
/// ```
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct CorrelationId(String);

impl CorrelationId {
    /// Generate a new random correlation ID using UUIDv4.
    ///
    /// This should typically be called only at the API Gateway when a new
    /// request enters the system. Downstream services should extract the
    /// correlation ID from gRPC metadata.
    ///
    /// # Example
    ///
    /// ```rust
    /// use common::observability::CorrelationId;
    ///
    /// let id = CorrelationId::new();
    /// println!("Generated ID: {}", id);
    /// ```
    pub fn new() -> Self {
        Self(Uuid::new_v4().to_string())
    }

    /// Create a correlation ID from a string.
    ///
    /// This is used when extracting correlation IDs from gRPC metadata or
    /// other external sources. The string should be a valid UUIDv4.
    ///
    /// # Arguments
    ///
    /// * `s` - A string representation of a UUID
    ///
    /// # Returns
    ///
    /// Returns `Some(CorrelationId)` if the string is a valid UUID, or `None`
    /// if the string is invalid.
    ///
    /// # Example
    ///
    /// ```rust
    /// use common::observability::CorrelationId;
    ///
    /// let id = CorrelationId::from_string("550e8400-e29b-41d4-a716-446655440000");
    /// assert!(id.is_some());
    ///
    /// let invalid = CorrelationId::from_string("not-a-uuid");
    /// assert!(invalid.is_none());
    /// ```
    pub fn from_string<S: AsRef<str>>(s: S) -> Option<Self> {
        let s = s.as_ref();
        // Validate that it's a proper UUID format
        Uuid::parse_str(s).ok().map(|_| Self(s.to_string()))
    }

    /// Get the correlation ID as a string slice.
    ///
    /// # Example
    ///
    /// ```rust
    /// use common::observability::CorrelationId;
    ///
    /// let id = CorrelationId::new();
    /// let id_str: &str = id.as_str();
    /// ```
    pub fn as_str(&self) -> &str {
        &self.0
    }

    /// Convert the correlation ID into an owned String.
    ///
    /// # Example
    ///
    /// ```rust
    /// use common::observability::CorrelationId;
    ///
    /// let id = CorrelationId::new();
    /// let id_string: String = id.into_string();
    /// ```
    pub fn into_string(self) -> String {
        self.0
    }
}

impl Default for CorrelationId {
    fn default() -> Self {
        Self::new()
    }
}

impl fmt::Display for CorrelationId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0)
    }
}

impl From<CorrelationId> for String {
    fn from(id: CorrelationId) -> Self {
        id.0
    }
}

impl AsRef<str> for CorrelationId {
    fn as_ref(&self) -> &str {
        &self.0
    }
}

/// Extension trait for gRPC request types to extract correlation IDs.
///
/// This trait provides convenient methods for extracting correlation IDs from
/// gRPC metadata in both server and client contexts.
pub trait CorrelationIdExt {
    /// Extract the correlation ID from gRPC metadata.
    ///
    /// If no correlation ID is present in the metadata, a new one is generated
    /// and returned. This ensures that all requests have a correlation ID for
    /// tracking purposes.
    ///
    /// # Returns
    ///
    /// Returns a `CorrelationId`, either extracted from metadata or newly generated.
    fn get_or_generate_correlation_id(&self) -> CorrelationId;

    /// Extract the correlation ID from gRPC metadata without generating a new one.
    ///
    /// # Returns
    ///
    /// Returns `Some(CorrelationId)` if present in metadata, or `None` if not found.
    fn get_correlation_id(&self) -> Option<CorrelationId>;
}

impl CorrelationIdExt for tonic::metadata::MetadataMap {
    fn get_or_generate_correlation_id(&self) -> CorrelationId {
        self.get_correlation_id().unwrap_or_default()
    }

    fn get_correlation_id(&self) -> Option<CorrelationId> {
        self.get(CORRELATION_ID_HEADER)
            .and_then(|v| v.to_str().ok())
            .and_then(CorrelationId::from_string)
    }
}

/// Set the current correlation ID for the async task context.
///
/// This function stores the correlation ID in thread-local storage, making it
/// available to all code executing within the current async task without
/// explicit parameter passing.
///
/// # Arguments
///
/// * `correlation_id` - The correlation ID to set
///
/// # Example
///
/// ```rust,no_run
/// use common::observability::{CorrelationId, set_correlation_id};
///
/// async fn handle_request() {
///     let id = CorrelationId::new();
///     set_correlation_id(id).await;
///
///     // The correlation ID is now available in all downstream calls
///     process_request().await;
/// }
///
/// async fn process_request() {
///     // Can access the correlation ID set by parent
/// }
/// ```
pub async fn set_correlation_id(correlation_id: CorrelationId) {
    let maybe_lock = CURRENT_CORRELATION_ID
        .try_with(Arc::clone)
        .ok();
    if let Some(lock) = maybe_lock {
        let mut guard = lock.write().await;
        *guard = Some(correlation_id);
    }
}

/// Get the current correlation ID from the async task context.
///
/// # Returns
///
/// Returns `Some(CorrelationId)` if a correlation ID has been set in the current
/// context, or `None` if no ID is available.
///
/// # Example
///
/// ```rust,no_run
/// use common::observability::get_correlation_id;
/// use tracing::info;
///
/// async fn log_with_correlation() {
///     if let Some(id) = get_correlation_id().await {
///         info!(correlation_id = %id, "Processing with correlation ID");
///     }
/// }
/// ```
pub async fn get_correlation_id() -> Option<CorrelationId> {
    let lock = CURRENT_CORRELATION_ID
        .try_with(Arc::clone)
        .ok()?;
    let guard = lock.read().await;
    guard.clone()
}

/// Create a new metadata map with the correlation ID header set.
///
/// This is useful when making outbound gRPC calls to propagate the correlation
/// ID to downstream services.
///
/// # Arguments
///
/// * `correlation_id` - The correlation ID to include in the metadata
///
/// # Returns
///
/// Returns a new `MetadataMap` with the correlation ID header set.
///
/// # Example
///
/// ```rust,no_run
/// use common::observability::{CorrelationId, create_metadata_with_correlation_id};
/// use tonic::Request;
///
/// async fn make_grpc_call() {
///     let correlation_id = CorrelationId::new();
///     let metadata = create_metadata_with_correlation_id(&correlation_id);
///
///     let mut request = Request::new(MyRequest {});
///     *request.metadata_mut() = metadata;
///
///     // Make gRPC call with correlation ID
/// }
/// ```
pub fn create_metadata_with_correlation_id(
    correlation_id: &CorrelationId,
) -> tonic::metadata::MetadataMap {
    let mut metadata = tonic::metadata::MetadataMap::new();

    if let Ok(value) = tonic::metadata::MetadataValue::try_from(correlation_id.as_str()) {
        metadata.insert(CORRELATION_ID_HEADER, value);
    }

    metadata
}

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

    #[test]
    fn test_correlation_id_new() {
        let id1 = CorrelationId::new();
        let id2 = CorrelationId::new();

        // IDs should be unique
        assert_ne!(id1, id2);

        // IDs should be valid UUID format
        assert!(Uuid::parse_str(id1.as_str()).is_ok());
        assert!(Uuid::parse_str(id2.as_str()).is_ok());
    }

    #[test]
    fn test_correlation_id_from_string() {
        let valid_uuid = "550e8400-e29b-41d4-a716-446655440000";
        let id = CorrelationId::from_string(valid_uuid);

        assert!(id.is_some());
        assert_eq!(id.unwrap().as_str(), valid_uuid);

        // Invalid UUID should return None
        let invalid = CorrelationId::from_string("not-a-uuid");
        assert!(invalid.is_none());
    }

    #[test]
    fn test_correlation_id_display() {
        let uuid_str = "550e8400-e29b-41d4-a716-446655440000";
        let id = CorrelationId::from_string(uuid_str).unwrap();

        assert_eq!(format!("{}", id), uuid_str);
    }

    #[test]
    fn test_correlation_id_into_string() {
        let uuid_str = "550e8400-e29b-41d4-a716-446655440000";
        let id = CorrelationId::from_string(uuid_str).unwrap();

        let string: String = id.into_string();
        assert_eq!(string, uuid_str);
    }

    #[test]
    fn test_metadata_extraction() {
        let mut metadata = tonic::metadata::MetadataMap::new();
        let uuid_str = "550e8400-e29b-41d4-a716-446655440000";

        let value = tonic::metadata::MetadataValue::try_from(uuid_str).unwrap();
        metadata.insert(CORRELATION_ID_HEADER, value);

        // Extract correlation ID from metadata
        let id = metadata.get_correlation_id();
        assert!(id.is_some());
        assert_eq!(id.unwrap().as_str(), uuid_str);
    }

    #[test]
    fn test_metadata_extraction_missing() {
        let metadata = tonic::metadata::MetadataMap::new();

        // Should return None when header is missing
        let id = metadata.get_correlation_id();
        assert!(id.is_none());
    }

    #[test]
    fn test_get_or_generate_correlation_id() {
        let mut metadata = tonic::metadata::MetadataMap::new();

        // Should generate new ID when missing
        let id1 = metadata.get_or_generate_correlation_id();
        assert!(Uuid::parse_str(id1.as_str()).is_ok());

        // Add correlation ID to metadata
        let uuid_str = "550e8400-e29b-41d4-a716-446655440000";
        let value = tonic::metadata::MetadataValue::try_from(uuid_str).unwrap();
        metadata.insert(CORRELATION_ID_HEADER, value);

        // Should extract existing ID
        let id2 = metadata.get_or_generate_correlation_id();
        assert_eq!(id2.as_str(), uuid_str);
    }

    #[test]
    fn test_create_metadata_with_correlation_id() {
        let id = CorrelationId::new();
        let metadata = create_metadata_with_correlation_id(&id);

        // Verify metadata contains correlation ID
        let extracted = metadata.get_correlation_id();
        assert!(extracted.is_some());
        assert_eq!(extracted.unwrap(), id);
    }

    #[tokio::test]
    async fn test_set_and_get_correlation_id() {
        let id = CorrelationId::new();
        let id_clone = id.clone();

        CURRENT_CORRELATION_ID
            .scope(Arc::new(RwLock::new(None)), async move {
                // Set correlation ID
                set_correlation_id(id_clone).await;

                // Retrieve correlation ID
                let retrieved = get_correlation_id().await;
                assert!(retrieved.is_some());
                assert_eq!(retrieved.unwrap(), id);
            })
            .await;
    }

    #[tokio::test]
    async fn test_get_correlation_id_not_set() {
        CURRENT_CORRELATION_ID
            .scope(Arc::new(RwLock::new(None)), async move {
                // Should return None when not set
                let retrieved = get_correlation_id().await;
                assert!(retrieved.is_none());
            })
            .await;
    }
}
