mod actor_id;
mod aud;
mod azp;
mod common;
mod org;
mod role;
mod scope;
mod services;

use crate::WorkspaceId;
use serde::{Deserialize, Serialize};
use uuid::Uuid;

#[doc(inline)]
pub use actor_id::{ActorId, ActorIdentifier, ActorKind};
#[doc(inline)]
pub use aud::Audience;
#[doc(inline)]
pub use azp::Azp;
#[doc(inline)]
pub use org::Org;
#[doc(inline)]
pub use role::{Role, RoleError, RoleSet};
#[doc(inline)]
pub use scope::*;
#[doc(inline)]
pub use services::{ServiceType, Services};

/// Defines the standard claims on a CipherStash "Service Token".
/// Service Tokens are signed by CTS and used to authenticate requests to CipherStash services like ZeroKMS.
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Claims {
    #[serde(deserialize_with = "deserialize_workspace_id")]
    pub workspace: WorkspaceId,
    pub iss: String,
    pub sub: String,
    pub aud: Audience,
    pub iat: u64,
    pub exp: u64,
    // TODO: Should this be always present?
    pub azp: Option<Azp>,
    pub scope: Scope,

    #[serde(default, alias = "role", alias = "roles")]
    pub role_set: RoleSet,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub device_id: Option<Uuid>,

    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub actor_id: Option<ActorIdentifier>,

    #[serde(default, skip_serializing_if = "Services::is_empty")]
    pub services: Services,

    /// The organisation id (from the upstream IDP / workspace).
    ///
    /// Producer-required: every CTS-issued token MUST contain this claim
    /// (enforced by `cts_domain::trust::ClaimsBuilder::build`).
    ///
    /// Consumer-optional: declared `Option<String>` with serde-default and
    /// `skip_serializing_if = "Option::is_none"` so tokens issued before this
    /// field existed still deserialize for the duration of one token TTL.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub org_id: Option<String>,
}

#[cfg(feature = "cached")]
const TOKEN_EXPIRY_LEEWAY_SECONDS: u64 = 60;

#[cfg(feature = "cached")]
impl cached::CanExpire for Claims {
    fn is_expired(&self) -> bool {
        (self.exp + TOKEN_EXPIRY_LEEWAY_SECONDS) < chrono::offset::Utc::now().timestamp() as u64
    }
}

// NOTE: This is a shim to maintain backward compatibility with the old WorkspaceId format.
// i.e. "ws:<workspace_id>".
// This can be removed 24 hours after we release once all old tokens have expired.
fn deserialize_workspace_id<'de, D>(deserializer: D) -> Result<WorkspaceId, D::Error>
where
    D: serde::Deserializer<'de>,
{
    use std::str::FromStr;
    let id = String::deserialize(deserializer)?;
    let parts: Vec<&str> = id.split(":").collect();
    match parts.len() {
        1 => WorkspaceId::from_str(&id).map_err(serde::de::Error::custom),
        2 => {
            if parts[0] != "ws" {
                return Err(serde::de::Error::custom(format!(
                    "Invalid workspace ID prefix: {}",
                    parts[0]
                )));
            }
            WorkspaceId::from_str(parts[1]).map_err(serde::de::Error::custom)
        }
        _ => Err(serde::de::Error::custom(format!(
            "Invalid workspace ID: {id}"
        ))),
    }
}

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

    #[test]
    fn test_deserialize_legacy_workspace_id() {
        let json_data = json!({
            "workspace": "ws:7366ITCXSAPCH5TN",
            "iss": "http://example.com",
            "sub": "user123",
            "aud": "example_audience",
            "iat": 1622547800,
            "exp": 1622547900,
            "azp": "OIDC",
            "scope": "read write",
            "role": "admin"
        });
        let claims: Claims =
            serde_json::from_value(json_data).expect("Failed to deserialize Claims");
        assert_eq!(
            claims.workspace,
            WorkspaceId::try_from("7366ITCXSAPCH5TN").unwrap()
        );
    }

    #[test]
    fn test_deserialize_workspace_id() {
        let json_data = json!({
            "workspace": "7366ITCXSAPCH5TN",
            "iss": "http://example.com",
            "sub": "user123",
            "aud": "example_audience",
            "iat": 1622547800,
            "exp": 1622547900,
            "azp": "OIDC",
            "scope": "read write",
            "role": "admin"
        });
        let claims: Claims =
            serde_json::from_value(json_data).expect("Failed to deserialize Claims");
        assert_eq!(
            claims.workspace,
            WorkspaceId::try_from("7366ITCXSAPCH5TN").unwrap()
        );
        assert_eq!(claims.iss, "http://example.com");
        assert_eq!(claims.sub, "user123");
        assert_eq!(claims.aud, Audience::new("example_audience"));
        assert_eq!(claims.iat, 1622547800);
        assert_eq!(claims.exp, 1622547900);
        assert_eq!(claims.azp, Some(Azp::RootProvider));
        assert_eq!(claims.scope, Scope::parse("read write"));
        assert!(claims.role_set.has_role(Role::Admin));
    }

    #[test]
    fn test_deserialize_with_org_id() {
        let json_data = json!({
            "workspace": "7366ITCXSAPCH5TN",
            "iss": "http://example.com",
            "sub": "user123",
            "aud": "example_audience",
            "iat": 1622547800,
            "exp": 1622547900,
            "azp": "OIDC",
            "scope": "read write",
            "org_id": "org_2abc"
        });
        let claims: Claims =
            serde_json::from_value(json_data).expect("Failed to deserialize Claims");
        assert_eq!(claims.org_id.as_deref(), Some("org_2abc"));
    }

    #[test]
    fn test_deserialize_without_org_id_is_tolerated() {
        let json_data = json!({
            "workspace": "7366ITCXSAPCH5TN",
            "iss": "http://example.com",
            "sub": "user123",
            "aud": "example_audience",
            "iat": 1622547800,
            "exp": 1622547900,
            "azp": "OIDC",
            "scope": "read write"
        });
        let claims: Claims =
            serde_json::from_value(json_data).expect("Failed to deserialize Claims");
        assert!(claims.org_id.is_none());
    }

    #[test]
    fn test_serialize_skips_org_id_when_none() {
        let claims = Claims {
            workspace: WorkspaceId::try_from("7366ITCXSAPCH5TN").unwrap(),
            iss: "http://example.com".into(),
            sub: "user123".into(),
            aud: Audience::new("example_audience"),
            iat: 1,
            exp: 2,
            azp: None,
            scope: Scope::parse(""),
            role_set: RoleSet::default(),
            device_id: None,
            actor_id: None,
            services: Services::new(),
            org_id: None,
        };
        let json = serde_json::to_value(&claims).unwrap();
        assert!(
            json.get("org_id").is_none(),
            "expected org_id to be omitted when None, got {json}"
        );
    }

    mod role_claim {
        use serde_json::Value;

        use super::*;

        fn build_claims(field: &str, role_set: Option<Value>) -> Value {
            let mut json_data = json!({
                "workspace": "7366ITCXSAPCH5TN",
                "iss": "http://example.com",
                "sub": "user123",
                "aud": "example_audience",
                "iat": 1622547800,
                "exp": 1622547900,
                "azp": "OIDC",
                "scope": "read write",
            });

            if let Some(role_set) = role_set {
                if let Some(obj) = json_data.as_object_mut() {
                    obj.insert(field.into(), role_set);
                }
            };

            json_data
        }

        #[test]
        fn test_deserialize_no_role() {
            fn check(claim_name: &str) {
                let json_data = build_claims(claim_name, None);
                let claims: Claims =
                    serde_json::from_value(json_data).expect("Failed to deserialize Claims");
                assert!(claims.role_set.is_empty())
            }

            check("role");
            check("roles");
            check("role_set");
        }

        #[test]
        fn test_deserialize_one_role() {
            fn check(claim_name: &str) {
                let json_data = build_claims(claim_name, Some(json!("admin")));
                let claims: Claims =
                    serde_json::from_value(json_data).expect("Failed to deserialize Claims");
                assert!(claims.role_set.has_role(Role::Admin));
                assert_eq!(claims.role_set.len(), 1);
            }

            check("role");
            check("roles");
            check("role_set");
        }

        #[test]
        fn test_deserialize_array_of_one_role() {
            fn check(claim_name: &str) {
                let json_data = build_claims(claim_name, Some(json!(["admin"])));
                let claims: Claims =
                    serde_json::from_value(json_data).expect("Failed to deserialize Claims");
                assert!(claims.role_set.has_role(Role::Admin));
                assert_eq!(claims.role_set.len(), 1);
            }

            check("role");
            check("roles");
            check("role_set");
        }

        #[test]
        fn test_deserialize_muliple_roles() {
            fn check(claim_name: &str) {
                let json_data = build_claims(claim_name, Some(json!(["admin", "member"])));
                let claims: Claims =
                    serde_json::from_value(json_data).expect("Failed to deserialize Claims");
                assert!(claims.role_set.has_role(Role::Admin));
                assert!(claims.role_set.has_role(Role::Member));
                assert_eq!(claims.role_set.len(), 2);
            }

            check("role");
            check("roles");
            check("role_set");
        }

        #[test]
        fn test_deserialize_empty_roles_array() {
            fn check(claim_name: &str) {
                let json_data = build_claims(claim_name, Some(json!([])));
                let claims: Claims =
                    serde_json::from_value(json_data).expect("Failed to deserialize Claims");
                assert!(claims.role_set.is_empty());
            }

            check("role");
            check("roles");
            check("role_set");
        }
    }
}
