// Copyright 2026-Present Datadog, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fmt::Display;

use aws_lc_rs::digest::{SHA256, SHA256_OUTPUT_LEN};
use thiserror::Error;

use crate::{cached_string_repr::CachedStringRepr, hex::colon_string, keys::PublicKey};

/// Constructing a [`KeyId`] from a byte slice failed due to incorrect length.
#[derive(Debug, Error)]
#[error("invalid key ID length {0}, expected {SHA256_OUTPUT_LEN}")]
pub struct KeyIdParseError(usize);

/// A [`KeyId`] uniquely identifies a [`PublicKey`].
///
/// This [`KeyId`] is suitable for use as an X509 Subject Key Identifier defined
/// in [RFC 5280 § 4.2.1.2]. The SKI is generated by constructing the
/// `SubjectPublicKeyInfo` ASN.1 message as defined in [RFC 5280 § 4.1] and
/// hashing the serialised DER bytes with SHA-256.
///
/// [RFC 5280 § 4.1]: https://tools.ietf.org/html/rfc5280#section-4.1
/// [RFC 5280 § 4.2.1.2]:
///     https://datatracker.ietf.org/doc/html/rfc5280#section-4.2.1.2
#[derive(Debug, PartialEq, Eq, Hash, Clone)]
pub struct KeyId {
    digest: [u8; SHA256_OUTPUT_LEN],

    /// A lazily-rendered string representation of `digest`.
    ///
    /// See [`Self::as_hex_str()`] for initialisation.
    rendered: CachedStringRepr,
}

impl KeyId {
    /// Render the [`KeyId`] as a lowercase hex string delimited by colons in
    /// the style of OpenSSL.
    ///
    /// Example: `cc:cb:0f:63:f1:63:5e:f1:0e:26:e8:82:f7:7a:6e:f9`
    ///
    /// This value is lazily rendered and cached for reuse.
    pub fn as_hex_str(&self) -> &str {
        self.rendered.get_or_init(|| colon_string(self.as_ref()))
    }

    /// Return this [`KeyId`] as a raw byte slice.
    pub fn as_bytes(&self) -> &[u8] {
        &self.digest
    }
}

impl std::ops::Deref for KeyId {
    type Target = [u8; 32];

    fn deref(&self) -> &Self::Target {
        &self.digest
    }
}

impl From<&PublicKey<'_>> for KeyId {
    fn from(key: &PublicKey) -> Self {
        // Combine the algorithm identifier and raw key bytes into a DER-encoded
        // SubjectPublicKeyInfo message.
        let info = rcgen::PublicKeyData::subject_public_key_info(key);

        Self {
            digest: aws_lc_rs::digest::digest(&SHA256, &info)
                .as_ref()
                .try_into()
                .expect("sha256 digest is 32 bytes"),
            rendered: Default::default(),
        }
    }
}

impl TryFrom<&[u8]> for KeyId {
    type Error = KeyIdParseError;

    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
        Ok(Self {
            digest: value.try_into().map_err(|_| KeyIdParseError(value.len()))?,
            rendered: Default::default(),
        })
    }
}

impl Display for KeyId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.as_hex_str())
    }
}

#[cfg(test)]
mod tests {
    use std::hash::{DefaultHasher, Hash, Hasher};

    use crate::keys::{PrivateKey, tests::fixture_key};

    use super::*;

    /// Ensure a fixed public key always returns the same Subject Key
    /// Identifier.
    #[test]
    fn test_ski_fixture() {
        const WANT: &[u8] = &[
            242, 141, 210, 92, 111, 76, 250, 141, 48, 196, 108, 210, 4, 182, 182, 128, 17, 12, 24,
            54, 159, 16, 208, 42, 122, 158, 205, 152, 190, 76, 82, 160,
        ];

        let key = fixture_key();
        let ski = KeyId::from(&key.public_key());

        assert_eq!(*ski, WANT);
        assert_eq!(
            ski.to_string(),
            "f2:8d:d2:5c:6f:4c:fa:8d:30:c4:6c:d2:04:b6:b6:80:11:0c:18:36:9f:10:d0:2a:7a:9e:cd:98:be:4c:52:a0"
        );
    }

    #[test]
    fn test_deterministic_ski() {
        let key = PrivateKey::new();
        let public = key.public_key();

        let ski = KeyId::from(&public);
        assert_eq!(ski, KeyId::from(&public));
    }

    #[test]
    fn test_eq() {
        let key = fixture_key();
        let a = KeyId::from(&key.public_key());
        let b = KeyId::from(&key.public_key());

        assert_eq!(a, b);

        // Drive the population of the cached rendered repr.
        let _ = b.to_string();
        assert_eq!(a, b);
    }

    #[test]
    fn test_hash() {
        let key = fixture_key();
        let a = KeyId::from(&key.public_key());
        let b = KeyId::from(&key.public_key());

        fn do_hash<T: Hash>(t: &T) -> u64 {
            let mut s = DefaultHasher::new();
            t.hash(&mut s);
            s.finish()
        }

        assert_eq!(do_hash(&a), do_hash(&b));

        // Drive the population of the cached rendered repr.
        let _ = b.to_string();
        assert_eq!(do_hash(&a), do_hash(&b));
    }
}
