use serde::{Deserialize, Deserializer, Serialize, de};
use uuid::Uuid;

const TRACEPARENT_VERSION: &str = "00";
const ZERO_TRACE_ID: &str = "00000000000000000000000000000000";
const ZERO_PARENT_ID: &str = "0000000000000000";

#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize)]
#[serde(transparent)]
pub struct RequestId(String);

impl RequestId {
    pub const MAX_LEN: usize = 128;

    pub fn generate() -> Self {
        Self(Uuid::now_v7().to_string())
    }

    pub fn resolve(candidate: Option<&str>) -> Self {
        candidate
            .filter(|value| Self::is_valid(value))
            .map(|value| Self(value.to_owned()))
            .unwrap_or_else(Self::generate)
    }

    pub fn parse(value: &str) -> Result<Self, CorrelationError> {
        Self::is_valid(value)
            .then(|| Self(value.to_owned()))
            .ok_or(CorrelationError::InvalidRequestId)
    }

    pub fn is_valid(value: &str) -> bool {
        !value.is_empty()
            && value.len() <= Self::MAX_LEN
            && value
                .bytes()
                .all(|byte| byte.is_ascii_graphic() && byte != b',' && byte != b';')
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl std::fmt::Display for RequestId {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.write_str(&self.0)
    }
}

impl<'de> Deserialize<'de> for RequestId {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let value = String::deserialize(deserializer)?;
        Self::parse(&value).map_err(de::Error::custom)
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize)]
#[serde(transparent)]
pub struct TraceId(String);

impl TraceId {
    pub const LEN: usize = 32;

    pub fn generate() -> Self {
        let value = Uuid::now_v7().simple().to_string();
        debug_assert_ne!(value, ZERO_TRACE_ID);
        Self(value)
    }

    pub fn parse(value: &str) -> Result<Self, CorrelationError> {
        if is_lower_hex(value, Self::LEN) && value != ZERO_TRACE_ID {
            Ok(Self(value.to_owned()))
        } else {
            Err(CorrelationError::InvalidTraceId)
        }
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl std::fmt::Display for TraceId {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.write_str(&self.0)
    }
}

impl<'de> Deserialize<'de> for TraceId {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let value = String::deserialize(deserializer)?;
        Self::parse(&value).map_err(de::Error::custom)
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
pub struct TraceContext {
    trace_id: TraceId,
    traceparent: String,
}

impl TraceContext {
    pub const TRACEPARENT_LEN: usize = 55;
    pub const TRACESTATE_MAX_BYTES: usize = 512;
    pub const TRACESTATE_MAX_MEMBERS: usize = 32;
    pub const BAGGAGE_MAX_BYTES: usize = 8_192;
    pub const BAGGAGE_MAX_MEMBERS: usize = 64;

    pub fn generate() -> Self {
        let trace_id = TraceId::generate();
        let mut parent_id = Uuid::now_v7().simple().to_string()[..16].to_owned();
        if parent_id == ZERO_PARENT_ID {
            parent_id.replace_range(15..16, "1");
        }
        // A context generated outside an SDK span must not claim that a sampler
        // selected it. Ingress replaces this seed with the actual local span
        // context before application code runs.
        let traceparent = format!("{TRACEPARENT_VERSION}-{trace_id}-{parent_id}-00");
        Self {
            trace_id,
            traceparent,
        }
    }

    pub fn parse(value: &str) -> Result<Self, CorrelationError> {
        if value.len() != Self::TRACEPARENT_LEN {
            return Err(CorrelationError::InvalidTraceparent);
        }
        let bytes = value.as_bytes();
        if bytes[2] != b'-' || bytes[35] != b'-' || bytes[52] != b'-' {
            return Err(CorrelationError::InvalidTraceparent);
        }
        let version = &value[0..2];
        let trace_id = &value[3..35];
        let parent_id = &value[36..52];
        let flags = &value[53..55];
        if version != TRACEPARENT_VERSION
            || !is_lower_hex(parent_id, 16)
            || parent_id == ZERO_PARENT_ID
            || !matches!(flags, "00" | "01")
        {
            return Err(CorrelationError::InvalidTraceparent);
        }
        Ok(Self {
            trace_id: TraceId::parse(trace_id).map_err(|_| CorrelationError::InvalidTraceparent)?,
            traceparent: value.to_owned(),
        })
    }

    pub fn from_span_parts(
        trace_id: &str,
        span_id: &str,
        sampled: bool,
    ) -> Result<Self, CorrelationError> {
        let flags = if sampled { "01" } else { "00" };
        Self::parse(&format!(
            "{TRACEPARENT_VERSION}-{trace_id}-{span_id}-{flags}"
        ))
    }

    pub fn continue_local(&self) -> Self {
        let mut span_id = Uuid::now_v7().simple().to_string()[..16].to_owned();
        if span_id == ZERO_PARENT_ID {
            span_id.replace_range(15..16, "1");
        }
        let sampled = self.traceparent.ends_with("-01");
        Self::from_span_parts(self.trace_id.as_str(), &span_id, sampled)
            .expect("generated span identity is canonical")
    }

    pub fn trace_id(&self) -> &TraceId {
        &self.trace_id
    }

    pub fn traceparent(&self) -> &str {
        &self.traceparent
    }

    /// Whether this context was selected for recording by the upstream/local sampler.
    pub fn is_sampled(&self) -> bool {
        self.traceparent.ends_with("-01")
    }

    pub fn tracestate_within_budget(value: &str) -> bool {
        header_list_within_budget(
            value,
            Self::TRACESTATE_MAX_BYTES,
            Self::TRACESTATE_MAX_MEMBERS,
        )
    }

    pub fn baggage_within_budget(value: &str) -> bool {
        header_list_within_budget(value, Self::BAGGAGE_MAX_BYTES, Self::BAGGAGE_MAX_MEMBERS)
    }
}

impl<'de> Deserialize<'de> for TraceContext {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        #[derive(Deserialize)]
        #[serde(deny_unknown_fields)]
        struct WireTraceContext {
            trace_id: TraceId,
            traceparent: String,
        }

        let wire = WireTraceContext::deserialize(deserializer)?;
        let context = Self::parse(&wire.traceparent).map_err(de::Error::custom)?;
        if context.trace_id != wire.trace_id {
            return Err(de::Error::custom(CorrelationError::InvalidTraceparent));
        }
        Ok(context)
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CorrelationContext {
    request_id: RequestId,
    trace_context: TraceContext,
}

impl CorrelationContext {
    pub fn new(request_id: RequestId, trace_context: TraceContext) -> Self {
        Self {
            request_id,
            trace_context,
        }
    }

    pub fn generate() -> Self {
        Self::new(RequestId::generate(), TraceContext::generate())
    }

    pub fn request_id(&self) -> &RequestId {
        &self.request_id
    }

    pub fn trace_context(&self) -> &TraceContext {
        &self.trace_context
    }

    pub fn trace_id(&self) -> &TraceId {
        self.trace_context.trace_id()
    }
}

#[derive(Clone, Copy, Debug, thiserror::Error, PartialEq, Eq)]
pub enum CorrelationError {
    #[error("invalid request identity")]
    InvalidRequestId,
    #[error("invalid trace identity")]
    InvalidTraceId,
    #[error("invalid trace parent")]
    InvalidTraceparent,
}

fn is_lower_hex(value: &str, expected_len: usize) -> bool {
    value.len() == expected_len
        && value
            .bytes()
            .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
}

fn header_list_within_budget(value: &str, max_bytes: usize, max_members: usize) -> bool {
    !value.is_empty()
        && value.len() <= max_bytes
        && value.is_ascii()
        && !value.bytes().any(|byte| byte.is_ascii_control())
        && value.split(',').count() <= max_members
        && value.split(',').all(|member| !member.trim().is_empty())
}