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

use aes_gcm::{aead::Aead, Aes256Gcm, Key, KeyInit, Nonce};
use argon2::{
    password_hash::{PasswordHasher, SaltString},
    Argon2,
};
use base64::{engine::general_purpose::STANDARD, Engine};
use rand_core::{OsRng, RngCore};
use tonic::{Request, Response, Status};

use iam_store::{CredentialStore, PrincipalStore};
use iam_types::{
    Argon2Params, CredentialRecord, PrincipalKind as TypesPrincipalKind, PrincipalRef,
};

use crate::proto::{
    iam_credential_server::IamCredential, CreateS3CredentialRequest, CreateS3CredentialResponse,
    Credential, GetSecretKeyRequest, GetSecretKeyResponse, ListCredentialsRequest,
    ListCredentialsResponse, PrincipalKind, RevokeCredentialRequest, RevokeCredentialResponse,
};

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

pub struct IamCredentialService {
    store: Arc<CredentialStore>,
    principal_store: Arc<PrincipalStore>,
    cipher: Aes256Gcm,
    key_id: String,
    admin_token: Option<String>,
}

impl IamCredentialService {
    pub fn new(
        store: Arc<CredentialStore>,
        principal_store: Arc<PrincipalStore>,
        master_key: &[u8],
        key_id: &str,
        admin_token: Option<String>,
    ) -> Result<Self, Status> {
        if master_key.len() != 32 {
            return Err(Status::failed_precondition(
                "IAM_CRED_MASTER_KEY must be 32 bytes",
            ));
        }
        let cipher = Aes256Gcm::new(Key::<Aes256Gcm>::from_slice(master_key));
        Ok(Self {
            store,
            principal_store,
            cipher,
            key_id: key_id.to_string(),
            admin_token,
        })
    }

    fn generate_secret() -> (String, Vec<u8>) {
        let raw = uuid::Uuid::new_v4().as_bytes().to_vec();
        let secret_b64 = STANDARD.encode(&raw);
        (secret_b64, raw)
    }

    fn hash_secret(raw: &[u8]) -> (String, Argon2Params) {
        let salt = SaltString::generate(&mut OsRng);
        let argon2 = Argon2::default();
        let hash = argon2
            .hash_password(raw, &salt)
            .expect("argon2 hash")
            .to_string();
        let params = Argon2Params {
            m_cost_kib: argon2.params().m_cost(),
            t_cost: argon2.params().t_cost(),
            p_cost: argon2.params().p_cost(),
            salt_b64: salt.to_string(),
        };
        (hash, params)
    }

    fn encrypt_secret(&self, raw: &[u8]) -> Result<String, Status> {
        let mut nonce_bytes = [0u8; 12];
        OsRng.fill_bytes(&mut nonce_bytes);
        let nonce = Nonce::from_slice(&nonce_bytes);
        let ciphertext = self
            .cipher
            .encrypt(nonce, raw)
            .map_err(|e| Status::internal(format!("encrypt secret: {}", e)))?;
        let mut combined = nonce_bytes.to_vec();
        combined.extend_from_slice(&ciphertext);
        Ok(STANDARD.encode(combined))
    }

    fn decrypt_secret(&self, enc_b64: &str) -> Result<Vec<u8>, Status> {
        let data = STANDARD
            .decode(enc_b64)
            .map_err(|e| Status::internal(format!("invalid b64: {}", e)))?;
        if data.len() < 12 {
            return Err(Status::internal("ciphertext too short"));
        }
        let (nonce_bytes, ct) = data.split_at(12);
        let nonce = Nonce::from_slice(nonce_bytes);
        self.cipher
            .decrypt(nonce, ct)
            .map_err(|e| Status::internal(format!("decrypt failed: {}", e)))
    }

    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 map_principal_kind(kind: i32) -> Result<TypesPrincipalKind, Status> {
    match PrincipalKind::try_from(kind).unwrap_or(PrincipalKind::Unspecified) {
        PrincipalKind::User => Ok(TypesPrincipalKind::User),
        PrincipalKind::ServiceAccount => Ok(TypesPrincipalKind::ServiceAccount),
        PrincipalKind::Group => Ok(TypesPrincipalKind::Group),
        PrincipalKind::Unspecified => Err(Status::invalid_argument("principal_kind is required")),
    }
}

#[tonic::async_trait]
impl IamCredential for IamCredentialService {
    async fn create_s3_credential(
        &self,
        request: Request<CreateS3CredentialRequest>,
    ) -> Result<Response<CreateS3CredentialResponse>, Status> {
        self.require_admin_token(&request)?;
        let req = request.into_inner();
        let now = now_ts();
        let principal_kind = map_principal_kind(req.principal_kind)?;
        let principal_ref = PrincipalRef::new(principal_kind.clone(), &req.principal_id);
        let principal = self
            .principal_store
            .get(&principal_ref)
            .await
            .map_err(|e| Status::internal(format!("principal lookup failed: {}", e)))?
            .ok_or_else(|| Status::not_found("principal not found"))?;
        if principal.org_id != req.org_id || principal.project_id != req.project_id {
            return Err(Status::invalid_argument(
                "credential tenant does not match principal tenant",
            ));
        }
        let (secret_b64, raw_secret) = Self::generate_secret();
        let (hash, kdf) = Self::hash_secret(&raw_secret);
        let secret_enc = self.encrypt_secret(&raw_secret)?;

        let access_key_id = format!("ak_{}", uuid::Uuid::new_v4());
        let record = CredentialRecord {
            access_key_id: access_key_id.clone(),
            principal_id: req.principal_id.clone(),
            principal_kind,
            org_id: req.org_id.clone(),
            project_id: req.project_id.clone(),
            created_at: now,
            expires_at: req.expires_at,
            revoked: false,
            description: if req.description.is_empty() {
                None
            } else {
                Some(req.description)
            },
            secret_hash: hash,
            secret_enc,
            key_id: self.key_id.clone(),
            version: 1,
            kdf,
        };

        self.store
            .put(&record)
            .await
            .map_err(|e| Status::internal(format!("store credential: {}", e)))?;

        Ok(Response::new(CreateS3CredentialResponse {
            access_key_id,
            secret_key: secret_b64,
            created_at: now,
            expires_at: req.expires_at,
        }))
    }

    async fn get_secret_key(
        &self,
        request: Request<GetSecretKeyRequest>,
    ) -> Result<Response<GetSecretKeyResponse>, Status> {
        self.require_admin_token(&request)?;
        let req = request.into_inner();
        let record = match self.store.get(&req.access_key_id).await {
            Ok(Some((rec, _))) => rec,
            Ok(None) => return Err(Status::not_found("access key not found")),
            Err(e) => {
                return Err(Status::internal(format!(
                    "failed to load credential: {}",
                    e
                )))
            }
        };
        if record.revoked {
            return Err(Status::permission_denied("access key revoked"));
        }
        if let Some(exp) = record.expires_at {
            if now_ts() > exp {
                return Err(Status::permission_denied("access key expired"));
            }
        }
        let secret = self.decrypt_secret(&record.secret_enc)?;

        Ok(Response::new(GetSecretKeyResponse {
            secret_key: STANDARD.encode(secret),
            principal_id: record.principal_id,
            expires_at: record.expires_at,
            org_id: record.org_id,
            project_id: record.project_id,
            principal_kind: match record.principal_kind {
                TypesPrincipalKind::User => PrincipalKind::User as i32,
                TypesPrincipalKind::ServiceAccount => PrincipalKind::ServiceAccount as i32,
                TypesPrincipalKind::Group => PrincipalKind::Group as i32,
            },
        }))
    }

    async fn list_credentials(
        &self,
        request: Request<ListCredentialsRequest>,
    ) -> Result<Response<ListCredentialsResponse>, Status> {
        self.require_admin_token(&request)?;
        let req = request.into_inner();
        let items = self
            .store
            .list_for_principal(&req.principal_id, 1000)
            .await
            .map_err(|e| Status::internal(format!("list credentials: {}", e)))?;
        let creds: Vec<Credential> = items
            .into_iter()
            .map(|c| Credential {
                access_key_id: c.access_key_id,
                principal_id: c.principal_id,
                created_at: c.created_at,
                expires_at: c.expires_at,
                revoked: c.revoked,
                description: c.description.unwrap_or_default(),
                org_id: c.org_id,
                project_id: c.project_id,
                principal_kind: match c.principal_kind {
                    TypesPrincipalKind::User => PrincipalKind::User as i32,
                    TypesPrincipalKind::ServiceAccount => PrincipalKind::ServiceAccount as i32,
                    TypesPrincipalKind::Group => PrincipalKind::Group as i32,
                },
            })
            .collect();
        Ok(Response::new(ListCredentialsResponse {
            credentials: creds,
        }))
    }

    async fn revoke_credential(
        &self,
        request: Request<RevokeCredentialRequest>,
    ) -> Result<Response<RevokeCredentialResponse>, Status> {
        self.require_admin_token(&request)?;
        let req = request.into_inner();
        let revoked = self
            .store
            .revoke(&req.access_key_id)
            .await
            .map_err(|e| Status::internal(format!("revoke: {}", e)))?;
        Ok(Response::new(RevokeCredentialResponse { success: revoked }))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use base64::engine::general_purpose::STANDARD;
    use iam_store::Backend;
    use iam_types::Principal;

    fn test_service() -> (IamCredentialService, Arc<PrincipalStore>) {
        let backend = Arc::new(Backend::memory());
        let store = Arc::new(CredentialStore::new(backend.clone()));
        let principal_store = Arc::new(PrincipalStore::new(backend));
        let master_key = [0x42u8; 32];
        (
            IamCredentialService::new(
                store,
                principal_store.clone(),
                &master_key,
                "test-key",
                None,
            )
            .unwrap(),
            principal_store,
        )
    }

    async fn seed_service_account(
        principal_store: &PrincipalStore,
        principal_id: &str,
        org_id: &str,
        project_id: &str,
    ) {
        let principal =
            Principal::new_service_account(principal_id, principal_id, org_id, project_id);
        principal_store.create(&principal).await.unwrap();
    }

    #[tokio::test]
    async fn create_and_get_roundtrip() {
        let (svc, principal_store) = test_service();
        seed_service_account(&principal_store, "p1", "org-a", "project-a").await;
        let create = svc
            .create_s3_credential(Request::new(CreateS3CredentialRequest {
                principal_id: "p1".into(),
                description: "".into(),
                expires_at: None,
                org_id: Some("org-a".into()),
                project_id: Some("project-a".into()),
                principal_kind: PrincipalKind::ServiceAccount as i32,
            }))
            .await
            .unwrap()
            .into_inner();

        let get = svc
            .get_secret_key(Request::new(GetSecretKeyRequest {
                access_key_id: create.access_key_id.clone(),
            }))
            .await
            .unwrap()
            .into_inner();

        let orig = STANDARD.decode(create.secret_key).unwrap();
        let fetched = STANDARD.decode(get.secret_key).unwrap();
        assert_eq!(orig, fetched);
        assert_eq!(get.principal_id, "p1");
        assert_eq!(get.org_id.as_deref(), Some("org-a"));
        assert_eq!(get.project_id.as_deref(), Some("project-a"));
        assert_eq!(get.principal_kind, PrincipalKind::ServiceAccount as i32);
    }

    #[tokio::test]
    async fn list_filters_by_principal() {
        let (svc, principal_store) = test_service();
        seed_service_account(&principal_store, "pA", "org-a", "project-a").await;
        seed_service_account(&principal_store, "pB", "org-b", "project-b").await;
        let a = svc
            .create_s3_credential(Request::new(CreateS3CredentialRequest {
                principal_id: "pA".into(),
                description: "".into(),
                expires_at: None,
                org_id: Some("org-a".into()),
                project_id: Some("project-a".into()),
                principal_kind: PrincipalKind::ServiceAccount as i32,
            }))
            .await
            .unwrap()
            .into_inner();
        let _b = svc
            .create_s3_credential(Request::new(CreateS3CredentialRequest {
                principal_id: "pB".into(),
                description: "".into(),
                expires_at: None,
                org_id: Some("org-b".into()),
                project_id: Some("project-b".into()),
                principal_kind: PrincipalKind::ServiceAccount as i32,
            }))
            .await
            .unwrap();

        let list_a = svc
            .list_credentials(Request::new(ListCredentialsRequest {
                principal_id: "pA".into(),
            }))
            .await
            .unwrap()
            .into_inner();
        assert_eq!(list_a.credentials.len(), 1);
        assert_eq!(list_a.credentials[0].access_key_id, a.access_key_id);
    }

    #[tokio::test]
    async fn revoke_blocks_get() {
        let (svc, principal_store) = test_service();
        seed_service_account(&principal_store, "p1", "org-a", "project-a").await;
        let created = svc
            .create_s3_credential(Request::new(CreateS3CredentialRequest {
                principal_id: "p1".into(),
                description: "".into(),
                expires_at: None,
                org_id: Some("org-a".into()),
                project_id: Some("project-a".into()),
                principal_kind: PrincipalKind::ServiceAccount as i32,
            }))
            .await
            .unwrap()
            .into_inner();

        let revoke1 = svc
            .revoke_credential(Request::new(RevokeCredentialRequest {
                access_key_id: created.access_key_id.clone(),
            }))
            .await
            .unwrap()
            .into_inner();
        assert!(revoke1.success);

        let revoke2 = svc
            .revoke_credential(Request::new(RevokeCredentialRequest {
                access_key_id: created.access_key_id.clone(),
            }))
            .await
            .unwrap()
            .into_inner();
        assert!(!revoke2.success);

        let err = svc
            .get_secret_key(Request::new(GetSecretKeyRequest {
                access_key_id: created.access_key_id,
            }))
            .await
            .unwrap_err();
        assert_eq!(err.code(), Status::permission_denied("").code());
    }

    #[tokio::test]
    async fn credential_rotation_cutover_keeps_new_key_live() {
        let (svc, principal_store) = test_service();
        seed_service_account(&principal_store, "p1", "org-a", "project-a").await;

        let old_credential = svc
            .create_s3_credential(Request::new(CreateS3CredentialRequest {
                principal_id: "p1".into(),
                description: "old".into(),
                expires_at: None,
                org_id: Some("org-a".into()),
                project_id: Some("project-a".into()),
                principal_kind: PrincipalKind::ServiceAccount as i32,
            }))
            .await
            .unwrap()
            .into_inner();

        let new_credential = svc
            .create_s3_credential(Request::new(CreateS3CredentialRequest {
                principal_id: "p1".into(),
                description: "new".into(),
                expires_at: None,
                org_id: Some("org-a".into()),
                project_id: Some("project-a".into()),
                principal_kind: PrincipalKind::ServiceAccount as i32,
            }))
            .await
            .unwrap()
            .into_inner();

        assert_ne!(old_credential.access_key_id, new_credential.access_key_id);

        let listed = svc
            .list_credentials(Request::new(ListCredentialsRequest {
                principal_id: "p1".into(),
            }))
            .await
            .unwrap()
            .into_inner();
        assert_eq!(listed.credentials.len(), 2);

        let revoke_old = svc
            .revoke_credential(Request::new(RevokeCredentialRequest {
                access_key_id: old_credential.access_key_id.clone(),
            }))
            .await
            .unwrap()
            .into_inner();
        assert!(revoke_old.success);

        let old_err = svc
            .get_secret_key(Request::new(GetSecretKeyRequest {
                access_key_id: old_credential.access_key_id,
            }))
            .await
            .unwrap_err();
        assert_eq!(old_err.code(), Status::permission_denied("").code());

        let new_secret = svc
            .get_secret_key(Request::new(GetSecretKeyRequest {
                access_key_id: new_credential.access_key_id,
            }))
            .await
            .unwrap()
            .into_inner();
        assert_eq!(new_secret.principal_id, "p1");
        assert_eq!(new_secret.org_id.as_deref(), Some("org-a"));
        assert_eq!(new_secret.project_id.as_deref(), Some("project-a"));
    }

    #[tokio::test]
    async fn expired_key_is_denied() {
        let (svc, principal_store) = test_service();
        seed_service_account(&principal_store, "p1", "org-a", "project-a").await;
        // Manually insert an expired record
        let expired = CredentialRecord {
            access_key_id: "expired-ak".into(),
            principal_id: "p1".into(),
            principal_kind: TypesPrincipalKind::ServiceAccount,
            org_id: Some("org-a".into()),
            project_id: Some("project-a".into()),
            created_at: now_ts(),
            expires_at: Some(now_ts() - 10),
            revoked: false,
            description: None,
            secret_hash: "hash".into(),
            secret_enc: STANDARD.encode(b"dead"),
            key_id: "k".into(),
            version: 1,
            kdf: Argon2Params {
                m_cost_kib: 19456,
                t_cost: 2,
                p_cost: 1,
                salt_b64: "c2FsdA==".into(),
            },
        };
        svc.store.put(&expired).await.unwrap();
        let err = svc
            .get_secret_key(Request::new(GetSecretKeyRequest {
                access_key_id: "expired-ak".into(),
            }))
            .await
            .unwrap_err();
        assert_eq!(err.code(), Status::permission_denied("").code());
    }

    #[test]
    fn master_key_length_enforced() {
        let backend = Arc::new(Backend::memory());
        let store = Arc::new(CredentialStore::new(backend.clone()));
        let principal_store = Arc::new(PrincipalStore::new(backend));
        let bad = IamCredentialService::new(store.clone(), principal_store, &[0u8; 16], "k", None);
        assert!(bad.is_err());
    }
}