//! IAM Token gRPC service implementation

use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
use sha2::{Digest, Sha256};
use tonic::{Request, Response, Status};

use iam_authn::InternalTokenService;
use iam_store::{PrincipalStore, TokenStore};
use iam_types::{
    InternalTokenClaims, PrincipalKind as TypesPrincipalKind, PrincipalRef, Scope as TypesScope,
    TokenMetadata, TokenType,
};

use crate::proto::{
    self, iam_token_server::IamToken, scope, IssueTokenRequest, IssueTokenResponse, PrincipalKind,
    RefreshTokenRequest, RefreshTokenResponse, RevokeTokenRequest, RevokeTokenResponse, Scope,
    ValidateTokenRequest, ValidateTokenResponse,
};

/// IAM Token service implementation
pub struct IamTokenService {
    token_service: Arc<InternalTokenService>,
    principal_store: Arc<PrincipalStore>,
    token_store: Arc<TokenStore>,
    admin_token: Option<String>,
}

impl IamTokenService {
    /// Create a new token service
    pub fn new(
        token_service: Arc<InternalTokenService>,
        principal_store: Arc<PrincipalStore>,
        token_store: Arc<TokenStore>,
        admin_token: Option<String>,
    ) -> Self {
        Self {
            token_service,
            principal_store,
            token_store,
            admin_token,
        }
    }

    fn convert_scope(proto_scope: &Option<Scope>) -> TypesScope {
        match proto_scope {
            Some(s) => match &s.scope {
                Some(scope::Scope::System(true)) => TypesScope::System,
                Some(scope::Scope::Org(org)) => TypesScope::org(&org.id),
                Some(scope::Scope::Project(proj)) => TypesScope::project(&proj.id, &proj.org_id),
                Some(scope::Scope::Resource(res)) => {
                    TypesScope::resource(&res.id, &res.project_id, &res.org_id)
                }
                _ => TypesScope::System,
            },
            None => TypesScope::System,
        }
    }

    fn convert_scope_to_proto(scope: &TypesScope) -> Scope {
        Scope {
            scope: Some(match scope {
                TypesScope::System => scope::Scope::System(true),
                TypesScope::Org { id } => scope::Scope::Org(proto::OrgScope { id: id.clone() }),
                TypesScope::Project { id, org_id } => scope::Scope::Project(proto::ProjectScope {
                    id: id.clone(),
                    org_id: org_id.clone(),
                }),
                TypesScope::Resource {
                    id,
                    project_id,
                    org_id,
                } => scope::Scope::Resource(proto::ResourceScope {
                    id: id.clone(),
                    project_id: project_id.clone(),
                    org_id: org_id.clone(),
                }),
            }),
        }
    }

    fn compute_token_id(token: &str) -> String {
        let mut hasher = Sha256::new();
        hasher.update(token.as_bytes());
        let digest = hasher.finalize();
        URL_SAFE_NO_PAD.encode(digest)
    }

    fn now_ts() -> u64 {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs()
    }

    async fn persist_metadata(
        &self,
        claims: &InternalTokenClaims,
        token: &str,
    ) -> Result<String, Status> {
        let token_id = Self::compute_token_id(token);
        let meta = TokenMetadata::new(
            &token_id,
            &claims.principal_id,
            TokenType::Access,
            claims.iat,
            claims.exp,
        );
        if let Err(e) = self.token_store.put(&meta).await {
            return Err(Status::internal(format!(
                "failed to persist token metadata: {}",
                e
            )));
        }
        // ensure we mirror revocation flags if claims already mark them
        if meta.revoked {
            return Ok(token_id);
        }
        Ok(token_id)
    }

    fn admin_token_valid(metadata: &tonic::metadata::MetadataMap, token: &str) -> bool {
        if let Some(value) = metadata.get("x-iam-admin-token") {
            if let Ok(raw) = value.to_str() {
                if raw.trim() == token {
                    return true;
                }
            }
        }

        if let Some(value) = metadata.get("authorization") {
            if let Ok(raw) = value.to_str() {
                let raw = raw.trim();
                if let Some(rest) = raw.strip_prefix("Bearer ") {
                    return rest.trim() == token;
                }
                if let Some(rest) = raw.strip_prefix("bearer ") {
                    return rest.trim() == token;
                }
            }
        }

        false
    }

    fn require_admin_token<T>(&self, request: &Request<T>) -> Result<(), Status> {
        if let Some(token) = self.admin_token.as_deref() {
            if !Self::admin_token_valid(request.metadata(), token) {
                return Err(Status::unauthenticated(
                    "missing or invalid IAM admin token",
                ));
            }
        }
        Ok(())
    }

    fn validate_scope_for_principal(
        principal: &iam_types::Principal,
        scope: &TypesScope,
    ) -> Result<(), Status> {
        let principal_org = principal.org_id.as_deref();
        let principal_project = principal.project_id.as_deref();

        match principal.kind {
            TypesPrincipalKind::ServiceAccount => match scope {
                TypesScope::System => Err(Status::permission_denied(
                    "service accounts cannot mint system-scoped tokens",
                )),
                TypesScope::Org { id } => {
                    if principal_org == Some(id.as_str()) {
                        Ok(())
                    } else {
                        Err(Status::permission_denied(
                            "service account token scope must match principal tenant",
                        ))
                    }
                }
                TypesScope::Project { id, org_id } => {
                    if principal_org == Some(org_id.as_str())
                        && principal_project == Some(id.as_str())
                    {
                        Ok(())
                    } else {
                        Err(Status::permission_denied(
                            "service account token scope must match principal tenant",
                        ))
                    }
                }
                TypesScope::Resource {
                    project_id, org_id, ..
                } => {
                    if principal_org == Some(org_id.as_str())
                        && principal_project == Some(project_id.as_str())
                    {
                        Ok(())
                    } else {
                        Err(Status::permission_denied(
                            "service account token scope must match principal tenant",
                        ))
                    }
                }
            },
            _ => {
                if let Some(org_id) = principal_org {
                    match scope {
                        TypesScope::System => {}
                        TypesScope::Org { id } if id == org_id => {}
                        TypesScope::Project {
                            org_id: scope_org, ..
                        }
                        | TypesScope::Resource {
                            org_id: scope_org, ..
                        } if scope_org == org_id => {}
                        _ => {
                            return Err(Status::permission_denied(
                                "token scope must match principal tenant",
                            ))
                        }
                    }
                }
                Ok(())
            }
        }
    }
}

#[tonic::async_trait]
impl IamToken for IamTokenService {
    async fn issue_token(
        &self,
        request: Request<IssueTokenRequest>,
    ) -> Result<Response<IssueTokenResponse>, Status> {
        self.require_admin_token(&request)?;
        let req = request.into_inner();

        // Get principal kind
        let principal_kind = match PrincipalKind::try_from(req.principal_kind) {
            Ok(PrincipalKind::User) => TypesPrincipalKind::User,
            Ok(PrincipalKind::ServiceAccount) => TypesPrincipalKind::ServiceAccount,
            Ok(PrincipalKind::Group) => TypesPrincipalKind::Group,
            _ => return Err(Status::invalid_argument("invalid principal kind")),
        };

        // Get principal from store
        let principal_ref = PrincipalRef::new(principal_kind, &req.principal_id);
        let principal = self
            .principal_store
            .get(&principal_ref)
            .await
            .map_err(|e| Status::internal(format!("Failed to get principal: {}", e)))?
            .ok_or_else(|| Status::not_found("Principal not found"))?;
        if !principal.enabled {
            return Err(Status::permission_denied("Principal is disabled"));
        }

        // Convert scope
        let scope = Self::convert_scope(&req.scope);
        Self::validate_scope_for_principal(&principal, &scope)?;

        // Determine TTL
        let ttl = if req.ttl_seconds > 0 {
            Some(Duration::from_secs(req.ttl_seconds))
        } else {
            None
        };

        // Issue token
        let issued = self
            .token_service
            .issue(&principal, req.roles.clone(), scope, ttl)
            .await
            .map_err(|e| Status::internal(format!("Failed to issue token: {}", e)))?;

        // Persist metadata for revocation tracking
        self.persist_metadata(&issued.claims, &issued.token).await?;

        Ok(Response::new(IssueTokenResponse {
            token: issued.token,
            expires_at: issued.expires_at,
            session_id: issued.claims.session_id,
        }))
    }

    async fn validate_token(
        &self,
        request: Request<ValidateTokenRequest>,
    ) -> Result<Response<ValidateTokenResponse>, Status> {
        let req = request.into_inner();

        match self.token_service.verify(&req.token).await {
            Ok(claims) => {
                // Check revocation list
                let token_id = Self::compute_token_id(&req.token);
                if let Some((meta, _)) = self
                    .token_store
                    .get(&claims.principal_id, &token_id)
                    .await
                    .map_err(|e| {
                        Status::internal(format!("Failed to read token metadata: {}", e))
                    })?
                {
                    if meta.revoked {
                        return Ok(Response::new(ValidateTokenResponse {
                            valid: false,
                            claims: None,
                            reason: "token revoked".into(),
                        }));
                    }
                }

                // Ensure principal still exists and is enabled
                let principal_ref =
                    PrincipalRef::new(claims.principal_kind.clone(), &claims.principal_id);
                let principal = self
                    .principal_store
                    .get(&principal_ref)
                    .await
                    .map_err(|e| Status::internal(format!("Failed to read principal: {}", e)))?;
                let Some(principal) = principal else {
                    return Ok(Response::new(ValidateTokenResponse {
                        valid: false,
                        claims: None,
                        reason: "principal not found".into(),
                    }));
                };
                if !principal.enabled {
                    return Ok(Response::new(ValidateTokenResponse {
                        valid: false,
                        claims: None,
                        reason: "principal disabled".into(),
                    }));
                }

                let proto_claims = crate::proto::InternalTokenClaims {
                    principal_id: claims.principal_id.clone(),
                    principal_kind: match claims.principal_kind {
                        TypesPrincipalKind::User => PrincipalKind::User as i32,
                        TypesPrincipalKind::ServiceAccount => PrincipalKind::ServiceAccount as i32,
                        TypesPrincipalKind::Group => PrincipalKind::Group as i32,
                    },
                    principal_name: claims.principal_name.clone(),
                    roles: claims.roles.clone(),
                    scope: Some(Self::convert_scope_to_proto(&claims.scope)),
                    org_id: claims.org_id.clone(),
                    project_id: claims.project_id.clone(),
                    node_id: claims.node_id.clone(),
                    iat: claims.iat,
                    exp: claims.exp,
                    session_id: claims.session_id.clone(),
                    auth_method: claims.auth_method.to_string(),
                };

                Ok(Response::new(ValidateTokenResponse {
                    valid: true,
                    claims: Some(proto_claims),
                    reason: String::new(),
                }))
            }
            Err(e) => Ok(Response::new(ValidateTokenResponse {
                valid: false,
                claims: None,
                reason: e.to_string(),
            })),
        }
    }

    async fn revoke_token(
        &self,
        request: Request<RevokeTokenRequest>,
    ) -> Result<Response<RevokeTokenResponse>, Status> {
        let req = request.into_inner();

        let claims = self
            .token_service
            .verify(&req.token)
            .await
            .map_err(|e| Status::invalid_argument(format!("Invalid token: {}", e)))?;

        let token_id = Self::compute_token_id(&req.token);
        let now = Self::now_ts();

        let existing = self
            .token_store
            .get(&claims.principal_id, &token_id)
            .await
            .map_err(|e| Status::internal(format!("Failed to read token metadata: {}", e)))?;

        // Try revoking existing metadata; if it does not exist, create a revoked record
        let revoked = if let Some((meta, _)) = existing {
            if meta.revoked {
                true
            } else {
                self.token_store
                    .revoke(&claims.principal_id, &token_id, &req.reason, now)
                    .await
                    .map_err(|e| Status::internal(format!("Failed to revoke token: {}", e)))?
            }
        } else {
            let mut meta = TokenMetadata::new(
                &token_id,
                &claims.principal_id,
                TokenType::Access,
                claims.iat,
                claims.exp,
            );
            meta.revoke(now, &req.reason);
            self.token_store
                .put(&meta)
                .await
                .map_err(|e| Status::internal(format!("Failed to write revocation: {}", e)))?;
            true
        };

        Ok(Response::new(RevokeTokenResponse { success: revoked }))
    }

    async fn refresh_token(
        &self,
        request: Request<RefreshTokenRequest>,
    ) -> Result<Response<RefreshTokenResponse>, Status> {
        let req = request.into_inner();

        // Validate current token
        let claims = self
            .token_service
            .verify(&req.token)
            .await
            .map_err(|e| Status::unauthenticated(format!("Invalid token: {}", e)))?;

        // Check revocation list
        let token_id = Self::compute_token_id(&req.token);
        if let Some((meta, _)) = self
            .token_store
            .get(&claims.principal_id, &token_id)
            .await
            .map_err(|e| Status::internal(format!("Failed to read token metadata: {}", e)))?
        {
            if meta.revoked {
                return Err(Status::permission_denied("token revoked"));
            }
        }

        // Get principal
        let principal_kind = claims.principal_kind.clone();
        let principal_ref = PrincipalRef::new(principal_kind, &claims.principal_id);
        let principal = self
            .principal_store
            .get(&principal_ref)
            .await
            .map_err(|e| Status::internal(format!("Failed to get principal: {}", e)))?
            .ok_or_else(|| Status::not_found("Principal not found"))?;
        if !principal.enabled {
            return Err(Status::permission_denied("Principal is disabled"));
        }

        // Determine new TTL
        let ttl = if req.ttl_seconds > 0 {
            Some(Duration::from_secs(req.ttl_seconds))
        } else {
            None
        };

        // Issue new token
        let issued = self
            .token_service
            .issue(&principal, claims.roles.clone(), claims.scope.clone(), ttl)
            .await
            .map_err(|e| Status::internal(format!("Failed to issue token: {}", e)))?;

        self.persist_metadata(&issued.claims, &issued.token).await?;

        Ok(Response::new(RefreshTokenResponse {
            token: issued.token,
            expires_at: issued.expires_at,
        }))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use iam_authn::{InternalTokenConfig, SigningKey};
    use iam_store::{Backend, TokenStore};
    use iam_types::Principal;

    fn test_setup() -> (
        Arc<InternalTokenService>,
        Arc<PrincipalStore>,
        Arc<TokenStore>,
    ) {
        let backend = Arc::new(Backend::memory());
        let principal_store = Arc::new(PrincipalStore::new(backend.clone()));
        let token_store = Arc::new(TokenStore::new(backend));

        let signing_key = SigningKey::generate("test-key");
        let config = InternalTokenConfig::new(signing_key, "iam-test");
        let token_service = Arc::new(InternalTokenService::new(config));

        (token_service, principal_store, token_store)
    }

    #[tokio::test]
    async fn test_issue_token_principal_not_found() {
        let (token_service, principal_store, token_store) = test_setup();
        let service = IamTokenService::new(token_service, principal_store, token_store, None);

        let req = IssueTokenRequest {
            principal_id: "nonexistent".into(),
            principal_kind: PrincipalKind::User as i32,
            roles: vec![],
            scope: None,
            ttl_seconds: 3600,
        };

        let result = service.issue_token(Request::new(req)).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_revoke_and_validate_blocklist() {
        let (token_service, principal_store, token_store) = test_setup();
        let service = IamTokenService::new(
            token_service,
            principal_store.clone(),
            token_store.clone(),
            None,
        );

        // create principal
        let principal = Principal::new_user("alice", "Alice");
        principal_store.create(&principal).await.unwrap();

        // issue token
        let issue_resp = service
            .issue_token(Request::new(IssueTokenRequest {
                principal_id: "alice".into(),
                principal_kind: PrincipalKind::User as i32,
                roles: vec!["roles/ProjectAdmin".into()],
                scope: None,
                ttl_seconds: 3600,
            }))
            .await
            .unwrap()
            .into_inner();

        // validate OK
        let valid_resp = service
            .validate_token(Request::new(ValidateTokenRequest {
                token: issue_resp.token.clone(),
            }))
            .await
            .unwrap()
            .into_inner();
        assert!(valid_resp.valid);

        // revoke
        let revoke_resp = service
            .revoke_token(Request::new(RevokeTokenRequest {
                token: issue_resp.token.clone(),
                reason: "test".into(),
            }))
            .await
            .unwrap()
            .into_inner();
        assert!(revoke_resp.success);

        // validate should be rejected
        let invalid_resp = service
            .validate_token(Request::new(ValidateTokenRequest {
                token: issue_resp.token.clone(),
            }))
            .await
            .unwrap()
            .into_inner();
        assert!(!invalid_resp.valid);
        assert_eq!(invalid_resp.reason, "token revoked");
    }

    #[tokio::test]
    async fn test_validate_token_principal_disabled() {
        let (token_service, principal_store, token_store) = test_setup();
        let service = IamTokenService::new(
            token_service,
            principal_store.clone(),
            token_store.clone(),
            None,
        );

        let principal = Principal::new_user("alice", "Alice");
        principal_store.create(&principal).await.unwrap();

        let issue_resp = service
            .issue_token(Request::new(IssueTokenRequest {
                principal_id: "alice".into(),
                principal_kind: PrincipalKind::User as i32,
                roles: vec!["roles/ProjectAdmin".into()],
                scope: None,
                ttl_seconds: 3600,
            }))
            .await
            .unwrap()
            .into_inner();

        let (mut stored, version) = principal_store
            .get_with_version(&principal.to_ref())
            .await
            .unwrap()
            .unwrap();
        stored.enabled = false;
        principal_store.update(&stored, version).await.unwrap();

        let valid_resp = service
            .validate_token(Request::new(ValidateTokenRequest {
                token: issue_resp.token,
            }))
            .await
            .unwrap()
            .into_inner();

        assert!(!valid_resp.valid);
        assert!(valid_resp.reason.contains("disabled"));
    }

    #[tokio::test]
    async fn test_service_account_system_scope_is_rejected() {
        let (token_service, principal_store, token_store) = test_setup();
        let service = IamTokenService::new(
            token_service,
            principal_store.clone(),
            token_store.clone(),
            None,
        );

        let principal = Principal::new_service_account("svc-1", "Service 1", "org-1", "proj-1");
        principal_store.create(&principal).await.unwrap();

        let result = service
            .issue_token(Request::new(IssueTokenRequest {
                principal_id: "svc-1".into(),
                principal_kind: PrincipalKind::ServiceAccount as i32,
                roles: vec![],
                scope: Some(IamTokenService::convert_scope_to_proto(&TypesScope::System)),
                ttl_seconds: 3600,
            }))
            .await;

        assert!(result.is_err());
        assert_eq!(result.unwrap_err().code(), tonic::Code::PermissionDenied);
    }
}