use actix_web::http::header::{HeaderMap, HeaderName, HeaderValue};

pub const X_ATHENA_TIME_HEADER: &str = "x-athena-time";
pub const X_ATHENA_REQUEST_ID_HEADER: &str = "x-athena-request-id";
pub const X_ATHENA_CACHED_HEADER: &str = "x-athena-cached";
pub const X_ATHENA_TRACE_ID_HEADER: &str = "x-athena-trace-id";
pub const X_ATHENA_CACHE_ID_HEADER: &str = "x-athena-cache-id";
pub const X_ATHENA_CACHE_KEY_HEADER: &str = "x-athena-cache-key";
pub const X_ATHENA_CACHE_OUTCOME_HEADER: &str = "x-athena-cache-outcome";
pub const X_ATHENA_CACHE_SOURCE_HEADER: &str = "x-athena-cache-source";
pub const X_ATHENA_CLIENT_HEADER: &str = "x-athena-client";

fn insert_header(headers: &mut HeaderMap, name: &'static str, value: &str) {
    if let Ok(header_value) = HeaderValue::from_str(value) {
        headers.insert(HeaderName::from_static(name), header_value);
    }
}

fn insert_header_if_missing(headers: &mut HeaderMap, name: &'static str, value: &str) {
    if !headers.contains_key(name) {
        insert_header(headers, name, value);
    }
}

fn synchronize_cache_identifier_headers(headers: &mut HeaderMap) {
    if !headers.contains_key(X_ATHENA_CACHE_ID_HEADER)
        && let Some(value) = headers.get(X_ATHENA_CACHE_KEY_HEADER).cloned()
    {
        headers.insert(HeaderName::from_static(X_ATHENA_CACHE_ID_HEADER), value);
    }

    if !headers.contains_key(X_ATHENA_CACHE_KEY_HEADER)
        && let Some(value) = headers.get(X_ATHENA_CACHE_ID_HEADER).cloned()
    {
        headers.insert(HeaderName::from_static(X_ATHENA_CACHE_KEY_HEADER), value);
    }
}

pub fn set_response_trace_id(headers: &mut HeaderMap, trace_id: &str) {
    insert_header(headers, X_ATHENA_TRACE_ID_HEADER, trace_id);
}

pub fn set_cache_headers(
    headers: &mut HeaderMap,
    cached: bool,
    cache_id: Option<&str>,
    cache_outcome: Option<&str>,
    cache_source: Option<&str>,
) {
    insert_header(
        headers,
        X_ATHENA_CACHED_HEADER,
        if cached { "true" } else { "false" },
    );

    if let Some(cache_id) = cache_id.filter(|value| !value.trim().is_empty()) {
        insert_header(headers, X_ATHENA_CACHE_ID_HEADER, cache_id);
        insert_header(headers, X_ATHENA_CACHE_KEY_HEADER, cache_id);
    }

    if let Some(cache_outcome) = cache_outcome.filter(|value| !value.trim().is_empty()) {
        insert_header(headers, X_ATHENA_CACHE_OUTCOME_HEADER, cache_outcome);
    }

    if let Some(cache_source) = cache_source.filter(|value| !value.trim().is_empty()) {
        insert_header(headers, X_ATHENA_CACHE_SOURCE_HEADER, cache_source);
    }

    synchronize_cache_identifier_headers(headers);
}

pub fn apply_common_response_headers(
    headers: &mut HeaderMap,
    elapsed_ms: u128,
    request_id: Option<&str>,
    trace_id: Option<&str>,
    client_name: Option<&str>,
) {
    insert_header(headers, X_ATHENA_TIME_HEADER, &elapsed_ms.to_string());

    if let Some(request_id) = request_id.filter(|value| !value.trim().is_empty()) {
        insert_header_if_missing(headers, X_ATHENA_REQUEST_ID_HEADER, request_id);
        insert_header_if_missing(headers, X_ATHENA_TRACE_ID_HEADER, request_id);
    }

    if let Some(trace_id) = trace_id.filter(|value| !value.trim().is_empty()) {
        insert_header(headers, X_ATHENA_TRACE_ID_HEADER, trace_id);
    }

    if let Some(client_name) = client_name.filter(|value| !value.trim().is_empty()) {
        insert_header_if_missing(headers, X_ATHENA_CLIENT_HEADER, client_name);
    }

    synchronize_cache_identifier_headers(headers);
    insert_header_if_missing(headers, X_ATHENA_CACHED_HEADER, "false");
}

#[cfg(test)]
mod tests {
    use super::{
        X_ATHENA_CACHE_ID_HEADER, X_ATHENA_CACHE_KEY_HEADER, X_ATHENA_CACHED_HEADER,
        X_ATHENA_CLIENT_HEADER, X_ATHENA_REQUEST_ID_HEADER, X_ATHENA_TIME_HEADER,
        X_ATHENA_TRACE_ID_HEADER, apply_common_response_headers, set_cache_headers,
    };
    use actix_web::HttpResponse;

    #[test]
    fn common_headers_add_time_trace_client_and_default_cached() {
        let mut response = HttpResponse::Ok().finish();

        apply_common_response_headers(
            response.headers_mut(),
            42,
            Some("req-123"),
            Some("trace-123"),
            Some("tenant_alpha"),
        );

        assert_eq!(response.headers().get(X_ATHENA_TIME_HEADER).unwrap(), "42");
        assert_eq!(
            response.headers().get(X_ATHENA_REQUEST_ID_HEADER).unwrap(),
            "req-123"
        );
        assert_eq!(
            response.headers().get(X_ATHENA_TRACE_ID_HEADER).unwrap(),
            "trace-123"
        );
        assert_eq!(
            response.headers().get(X_ATHENA_CLIENT_HEADER).unwrap(),
            "tenant_alpha"
        );
        assert_eq!(
            response.headers().get(X_ATHENA_CACHED_HEADER).unwrap(),
            "false"
        );
    }

    #[test]
    fn common_headers_alias_legacy_cache_key_to_cache_id() {
        let mut response = HttpResponse::Ok().finish();
        response.headers_mut().insert(
            X_ATHENA_CACHE_KEY_HEADER.parse().unwrap(),
            "cache-123".parse().unwrap(),
        );

        apply_common_response_headers(response.headers_mut(), 1, None, None, None);

        assert_eq!(
            response.headers().get(X_ATHENA_CACHE_ID_HEADER).unwrap(),
            "cache-123"
        );
    }

    #[test]
    fn common_headers_fallback_trace_id_to_request_id() {
        let mut response = HttpResponse::Ok().finish();

        apply_common_response_headers(response.headers_mut(), 7, Some("req-789"), None, None);

        assert_eq!(
            response.headers().get(X_ATHENA_TRACE_ID_HEADER).unwrap(),
            "req-789"
        );
    }

    #[test]
    fn set_cache_headers_write_both_cache_id_and_legacy_key() {
        let mut response = HttpResponse::Ok().finish();

        set_cache_headers(
            response.headers_mut(),
            true,
            Some("cache-456"),
            Some("hit_local"),
            Some("local"),
        );

        assert_eq!(
            response.headers().get(X_ATHENA_CACHED_HEADER).unwrap(),
            "true"
        );
        assert_eq!(
            response.headers().get(X_ATHENA_CACHE_ID_HEADER).unwrap(),
            "cache-456"
        );
        assert_eq!(
            response.headers().get(X_ATHENA_CACHE_KEY_HEADER).unwrap(),
            "cache-456"
        );
    }
}
