use std::{
    sync::atomic::{AtomicU64, Ordering},
    time::{SystemTime, UNIX_EPOCH},
};

use base64::Engine;
use jsonwebtoken::{Algorithm, DecodingKey, EncodingKey, Header};
use uuid::Uuid;

use crate::{
    error::KeeperError,
    types::{OpsClaims, PendingOp, ScopeEntry, SingleOrVec},
};

/// Stateful JWT signer: holds the loaded Ed25519 key pair and a monotonic sequence counter.
/// The `seq` counter is per-process; across restarts it resets, which is acceptable because
/// the ops-controller validates `jti` uniqueness (not `seq` ordering) as its idempotency key.
pub struct Signer {
    encoding_key: EncodingKey,
    decoding_key: DecodingKey,
    pub issuer_id: String,
    workspace: String,
    token_validity_secs: u64,
    seq: AtomicU64,
}

impl Signer {
    /// Load an Ed25519 signer from PKCS#8 PEM files.
    /// The private key must be in PKCS#8 PEM format (e.g., output of
    /// `openssl genpkey -algorithm ED25519 -out keeper.pem`).
    pub fn from_pem_files(
        private_pem: &[u8],
        public_pem: &[u8],
        issuer_id: String,
        workspace: String,
        token_validity_secs: u64,
    ) -> Result<Self, KeeperError> {
        let priv_der = pem_to_der(private_pem)?;
        let pub_der = pem_to_der(public_pem)?;
        let encoding_key = EncodingKey::from_ed_der(&priv_der);
        let decoding_key = DecodingKey::from_ed_der(&pub_der);
        Ok(Self {
            encoding_key,
            decoding_key,
            issuer_id,
            workspace,
            token_validity_secs,
            seq: AtomicU64::new(1),
        })
    }

    /// Build and sign a JWT for the given pending op.
    /// Returns the compact JWT string ready for publication to ops.cmd.*.
    pub fn sign_op(&self, op: &PendingOp) -> Result<String, KeeperError> {
        let now = unix_now();
        let claims = OpsClaims {
            iss: self.issuer_id.clone(),
            sub: op.sub.clone(),
            aud: SingleOrVec::single(&self.workspace),
            scopes: vec![ScopeEntry {
                op_type: op.op_type.clone(),
                target_pattern: op.target.clone(),
            }],
            seq: self.seq.fetch_add(1, Ordering::SeqCst),
            jti: Uuid::new_v4().to_string(),
            expected_state_version: op.expected_state_version.clone(),
            exp: now + self.token_validity_secs,
            nbf: now,
            op: op.clone(),
        };
        jsonwebtoken::encode(&Header::new(Algorithm::EdDSA), &claims, &self.encoding_key)
            .map_err(|e| KeeperError::Sign(e.to_string()))
    }

    /// Verify a JWT signed by this signer. Used in tests and CLI describe command.
    pub fn verify(&self, token: &str) -> Result<OpsClaims, KeeperError> {
        let mut validation = jsonwebtoken::Validation::new(Algorithm::EdDSA);
        validation.set_audience(&[&self.workspace]);
        validation.validate_nbf = true;
        jsonwebtoken::decode::<OpsClaims>(token, &self.decoding_key, &validation)
            .map(|td| td.claims)
            .map_err(|e| KeeperError::Sign(format!("JWT verification failed: {e}")))
    }

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

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

/// Strip PEM headers/footers and base64-decode to DER bytes.
fn pem_to_der(pem: &[u8]) -> Result<Vec<u8>, KeeperError> {
    let pem_str = std::str::from_utf8(pem)
        .map_err(|e| KeeperError::Key(format!("PEM is not UTF-8: {e}")))?;
    let body: String = pem_str
        .lines()
        .filter(|l| !l.starts_with("-----"))
        .collect();
    base64::engine::general_purpose::STANDARD
        .decode(body.as_bytes())
        .map_err(|e| KeeperError::Key(format!("PEM base64 decode: {e}")))
}

fn unix_now() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .expect("system clock is before Unix epoch")
        .as_secs()
}

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

    fn make_test_signer() -> Signer {
        // Minimal HMAC-secret signer for structural tests — not EdDSA, only used to
        // verify seq counter and struct layout without requiring real Ed25519 keys.
        Signer {
            encoding_key: EncodingKey::from_secret(b"test-secret-at-least-32-bytes-long!!"),
            decoding_key: DecodingKey::from_secret(b"test-secret-at-least-32-bytes-long!!"),
            issuer_id: "test-keeper".to_string(),
            workspace: "test-ws".to_string(),
            token_validity_secs: 3600,
            seq: AtomicU64::new(1),
        }
    }

    #[test]
    fn seq_increments_on_each_sign_attempt() {
        let signer = make_test_signer();
        let initial = signer.seq.load(Ordering::SeqCst);
        let op = PendingOp {
            op_type: "scale".into(),
            target: "test".into(),
            sub: "test-principal".into(),
            expected_state_version: "v1".into(),
            image: String::new(),
            params: serde_json::Value::Null,
        };
        // sign_op will fail (HMAC key + EdDSA header) but seq is incremented before encode
        let _ = signer.sign_op(&op);
        assert!(signer.seq.load(Ordering::SeqCst) > initial);
    }

    #[test]
    fn pem_to_der_strips_headers() {
        let pem = b"-----BEGIN PUBLIC KEY-----\nYWJj\n-----END PUBLIC KEY-----\n";
        // "YWJj" decodes to b"abc"
        assert_eq!(pem_to_der(pem).unwrap(), b"abc");
    }

    #[test]
    fn pem_to_der_rejects_invalid_base64() {
        let pem = b"-----BEGIN PUBLIC KEY-----\n!!!notbase64!!!\n-----END PUBLIC KEY-----\n";
        assert!(pem_to_der(pem).is_err());
    }
}
