use miden_protocol::Word;
use miden_protocol::account::auth::{AuthScheme, PublicKey, PublicKeyCommitment};
use miden_protocol::account::component::{
    AccountComponentCode,
    AccountComponentMetadata,
    SchemaType,
    StorageSchema,
    StorageSlotSchema,
};
use miden_protocol::account::{
    AccountComponent,
    AccountComponentName,
    StorageSlot,
    StorageSlotName,
};
use miden_protocol::crypto::dsa::{ecdsa_k256_keccak, falcon512_poseidon2};
use miden_protocol::utils::sync::LazyLock;

use crate::account::account_component_code;

account_component_code!(SINGLESIG_CODE, "auth/singlesig.masl");

// CONSTANTS
// ================================================================================================

static PUBKEY_SLOT_NAME: LazyLock<StorageSlotName> = LazyLock::new(|| {
    StorageSlotName::new("miden::standards::auth::singlesig::pub_key")
        .expect("storage slot name should be valid")
});

static SCHEME_ID_SLOT_NAME: LazyLock<StorageSlotName> = LazyLock::new(|| {
    StorageSlotName::new("miden::standards::auth::singlesig::scheme")
        .expect("storage slot name should be valid")
});

/// An [`AccountComponent`] implementing the signature scheme for authentication
/// of transactions.
///
/// This component exports `auth_tx`, which loads the public key and signature scheme id from
/// storage and delegates transaction authentication to
/// `miden::standards::auth::signature::authenticate_transaction`.
///
/// When linking against this component, the `miden::standards` library must be available to the
/// assembler (which also implies availability of `miden::protocol`). This is the case when using
/// [`CodeBuilder`][builder].
///
/// [builder]: crate::code_builder::CodeBuilder
pub struct AuthSingleSig {
    pub_key: PublicKeyCommitment,
    auth_scheme: AuthScheme,
}

impl AuthSingleSig {
    /// The name of the component.
    pub const NAME: &'static str = "miden::standards::components::auth::singlesig";

    /// Returns the canonical [`AccountComponentName`] of this component.
    pub const fn name() -> AccountComponentName {
        AccountComponentName::from_static_str(Self::NAME)
    }

    /// Returns the [`AccountComponentCode`] of this component.
    pub fn code() -> &'static AccountComponentCode {
        &SINGLESIG_CODE
    }

    /// Creates a new [`AuthSingleSig`] component with the given `public_key`.
    pub fn new(pub_key: PublicKeyCommitment, auth_scheme: AuthScheme) -> Self {
        Self { pub_key, auth_scheme }
    }

    /// Creates a new [`AuthSingleSig`] component using the Falcon512Poseidon2 signature scheme.
    ///
    /// The public key commitment is derived from the provided Falcon512 public key.
    pub fn falcon512_poseidon2(pub_key: falcon512_poseidon2::PublicKey) -> Self {
        Self {
            pub_key: pub_key.into(),
            auth_scheme: AuthScheme::Falcon512Poseidon2,
        }
    }

    /// Creates a new [`AuthSingleSig`] component using the EcdsaK256Keccak signature scheme.
    ///
    /// The public key commitment is derived from the provided ECDSA K256 public key.
    pub fn ecdsa_k256_keccak(pub_key: ecdsa_k256_keccak::PublicKey) -> Self {
        Self {
            pub_key: pub_key.into(),
            auth_scheme: AuthScheme::EcdsaK256Keccak,
        }
    }

    /// Creates a new [`AuthSingleSig`] component from a [`PublicKey`].
    ///
    /// The authentication scheme and public key commitment are derived from the provided key.
    pub fn from_public_key(pub_key: PublicKey) -> Self {
        Self {
            auth_scheme: pub_key.auth_scheme(),
            pub_key: pub_key.to_commitment(),
        }
    }

    /// Returns the [`StorageSlotName`] where the public key is stored.
    pub fn public_key_slot() -> &'static StorageSlotName {
        &PUBKEY_SLOT_NAME
    }

    // Returns the [`StorageSlotName`] where the scheme ID is stored.
    pub fn scheme_id_slot() -> &'static StorageSlotName {
        &SCHEME_ID_SLOT_NAME
    }

    /// Returns the storage slot schema for the public key slot.
    pub fn public_key_slot_schema() -> (StorageSlotName, StorageSlotSchema) {
        (
            Self::public_key_slot().clone(),
            StorageSlotSchema::value("Public key commitment", SchemaType::pub_key()),
        )
    }
    /// Returns the storage slot schema for the scheme ID slot.
    pub fn auth_scheme_slot_schema() -> (StorageSlotName, StorageSlotSchema) {
        (
            Self::scheme_id_slot().clone(),
            StorageSlotSchema::value("Scheme ID", SchemaType::auth_scheme()),
        )
    }

    /// Returns the [`AccountComponentMetadata`] for this component.
    pub fn component_metadata() -> AccountComponentMetadata {
        let storage_schema = StorageSchema::new(vec![
            Self::public_key_slot_schema(),
            Self::auth_scheme_slot_schema(),
        ])
        .expect("storage schema should be valid");

        AccountComponentMetadata::new(Self::NAME)
            .with_description(
                "Authentication component using ECDSA K256 Keccak or Falcon512 Poseidon2 signature scheme",
            )
            .with_storage_schema(storage_schema)
    }
}

impl From<AuthSingleSig> for AccountComponent {
    fn from(basic_signature: AuthSingleSig) -> Self {
        let metadata = AuthSingleSig::component_metadata();

        let storage_slots = vec![
            StorageSlot::with_value(
                AuthSingleSig::public_key_slot().clone(),
                basic_signature.pub_key.into(),
            ),
            StorageSlot::with_value(
                AuthSingleSig::scheme_id_slot().clone(),
                Word::from([basic_signature.auth_scheme.as_u8(), 0, 0, 0]),
            ),
        ];

        AccountComponent::new(AuthSingleSig::code().clone(), storage_slots, metadata).expect(
            "singlesig component should satisfy the requirements of a valid account component",
        )
    }
}
