use crate::{
    core::{
        circuits::boolean::{
            boolean_value::BooleanValue,
            ed25519::{Ed25519SecretKey, Ed25519SigningKey},
        },
        compile::*,
        global_value::{
            curve_value::CurveValue,
            global_expr_store::with_local_expr_store_as_global,
            value::FieldValue,
        },
        ir_builder::IRBuilder,
    },
    traits::{Random, Reveal, ToLeBytes},
    utils::{
        crypto::key::{
            AES128Key,
            AES192Key,
            AES256Key,
            RescueKey,
            X25519PrivateKey,
            X25519PublicKey,
            ED25519_VERIFYING_KEY_COUNT,
        },
        field::{BaseField, ScalarField},
        zkp::{
            elgamal::ElGamalKeypair,
            pubkey_validity::{
                PubkeyValidityProof,
                PubkeyValidityProofData,
                PUBKEY_VALIDITY_PROOF_LEN,
            },
        },
    },
};
use arcis_interface::{CircuitInterface, ScalarKind, Value};

#[allow(clippy::type_complexity)]
fn mxe_keygen() -> (
    X25519PrivateKey<FieldValue<ScalarField>>,
    X25519PublicKey<CurveValue>,
    RescueKey<FieldValue<BaseField>>,
    RescueKey<FieldValue<ScalarField>>,
    AES128Key<BooleanValue>,
    AES192Key<BooleanValue>,
    AES256Key<BooleanValue>,
    Ed25519SecretKey,
    Ed25519SigningKey,
    ElGamalKeypair,
    PubkeyValidityProof,
) {
    // x25519
    let x25519_private_key = X25519PrivateKey::random();
    let x25519_pubkey = X25519PublicKey::new_from_private_key(x25519_private_key).reveal();
    // rescue
    let rescue_base_field_key = RescueKey::<FieldValue<BaseField>>::random();
    let rescue_scalar_field_key = RescueKey::<FieldValue<ScalarField>>::random();
    // aes
    let aes_128_key = AES128Key::<BooleanValue>::random();
    let aes_192_key = AES192Key::<BooleanValue>::random();
    let aes_256_key = AES256Key::<BooleanValue>::random();
    // ed25519
    let ed25519_secret_key = Ed25519SecretKey::random();
    let ed25519_signing_key = Ed25519SigningKey::from(ed25519_secret_key);
    let ed25519_signing_key = Ed25519SigningKey::new(
        ed25519_signing_key.s,
        ed25519_signing_key.hash_prefix,
        ed25519_signing_key.verifying_key.reveal(),
    );
    // ElGamal
    let elgamal_keypair = ElGamalKeypair::new_rand();
    let pubkey_validity_proof = PubkeyValidityProofData::new(&elgamal_keypair).proof;

    (
        x25519_private_key,
        x25519_pubkey,
        rescue_base_field_key,
        rescue_scalar_field_key,
        aes_128_key,
        aes_192_key,
        aes_256_key,
        ed25519_secret_key,
        ed25519_signing_key,
        elgamal_keypair,
        pubkey_validity_proof,
    )
}

pub fn mxe_keygen_circuit() -> Vec<u8> {
    let mut expr_store = IRBuilder::new(true);
    let output_ids = with_local_expr_store_as_global(
        || {
            let (
                x25519_private_key,
                x25519_pubkey,
                rescue_base_field_key,
                rescue_scalar_field_key,
                aes_128_key,
                aes_192_key,
                aes_256_key,
                ed25519_secret_key,
                ed25519_signing_key,
                elgamal_keypair,
                pubkey_validity_proof,
            ) = mxe_keygen();

            // x25519
            let mut output_ids = vec![
                x25519_private_key.inner().get_id(),
                x25519_pubkey.inner().get_id(),
            ];
            // Rescue
            output_ids.extend(
                rescue_base_field_key
                    .inner()
                    .into_iter()
                    .map(|val| val.get_id()),
            );
            output_ids.extend(
                rescue_scalar_field_key
                    .inner()
                    .into_iter()
                    .map(|val| val.get_id()),
            );
            // aes
            output_ids.extend(
                aes_128_key
                    .inner()
                    .into_iter()
                    .flat_map(|byte| byte.to_vec())
                    .map(|bit| bit.get_id()),
            );
            output_ids.extend(
                aes_192_key
                    .inner()
                    .into_iter()
                    .flat_map(|byte| byte.to_vec())
                    .map(|bit| bit.get_id()),
            );
            output_ids.extend(
                aes_256_key
                    .inner()
                    .into_iter()
                    .flat_map(|byte| byte.to_vec())
                    .map(|bit| bit.get_id()),
            );
            // ed25519
            output_ids.extend(
                ed25519_secret_key
                    .inner()
                    .into_iter()
                    .flat_map(|byte| byte.to_vec())
                    .map(|bit| bit.get_id()),
            );
            output_ids.push(ed25519_signing_key.s.get_id());
            output_ids.extend(
                ed25519_signing_key
                    .hash_prefix
                    .into_iter()
                    .flat_map(|byte| byte.to_vec())
                    .map(|bit| bit.get_id()),
            );
            output_ids.extend(
                ed25519_signing_key
                    .verifying_key
                    .public_key_encoded
                    .into_iter()
                    .map(|byte| FieldValue::<BaseField>::from(byte).get_id()),
            );
            // ElGamal
            output_ids.push(elgamal_keypair.secret().get_scalar().get_id());
            output_ids.push(elgamal_keypair.pubkey().get_point().get_id());
            output_ids.extend(
                pubkey_validity_proof
                    .Y
                    .to_bytes()
                    .into_iter()
                    .map(|byte| FieldValue::<BaseField>::from(byte).get_id()),
            );
            output_ids.extend(
                pubkey_validity_proof
                    .z
                    .to_le_bytes()
                    .into_iter()
                    .map(|byte| FieldValue::<BaseField>::from(byte).get_id()),
            );

            output_ids
        },
        &mut expr_store,
    );
    let ir = expr_store.into_ir(output_ids);
    let compiled_circuit = Compiler::optimize_into_circuitable(ir).to_async_mpc_circuit();
    bincode::serialize(&compiled_circuit).unwrap()
}

pub fn mxe_keygen_interface() -> String {
    let x25519_pubkey_type = Value::ArcisX25519Pubkey;
    let ed25519_verifying_key = vec![
        Value::Scalar {
            size_in_bits: 8,
            kind: ScalarKind::Unsigned
        };
        ED25519_VERIFYING_KEY_COUNT
    ];
    let elgamal_pubkey_type = Value::Point;
    let pubkey_validity_proof_type = vec![
        Value::Scalar {
            size_in_bits: 8,
            kind: ScalarKind::Unsigned
        };
        PUBKEY_VALIDITY_PROOF_LEN
    ];

    let mut key_types = vec![x25519_pubkey_type];
    key_types.extend(ed25519_verifying_key);
    key_types.push(elgamal_pubkey_type);
    key_types.extend(pubkey_validity_proof_type);

    let interface = CircuitInterface::new("mxe_keygen".to_string(), Vec::new(), key_types);
    interface.serialize().unwrap()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{core::circuits::pre_compiled::constants::ARTIFACTS_DIR, AsyncMPCCircuit};
    use std::fs;

    // This test overwrites the current circuit and interface.
    // ⚠️ Only run intentionally, using: cargo test --all-features -- --ignored
    #[test]
    #[ignore]
    fn write_keygen_circuit() {
        fs::create_dir_all(ARTIFACTS_DIR).expect("Failed to create artifacts directory");

        let circuit = mxe_keygen_circuit();
        fs::write(
            format!("{}/mxe_keygen/circuit.arcis", ARTIFACTS_DIR),
            &circuit,
        )
        .expect("Failed to write circuit artifact");

        let interface = mxe_keygen_interface();
        fs::write(
            format!("{}/mxe_keygen/circuit.idarc", ARTIFACTS_DIR),
            interface.as_bytes(),
        )
        .expect("Failed to write interface artifact");
    }

    #[test]
    fn test_keygen_circuit() {
        let rng = &mut crate::utils::test_rng::get();
        let circuit = mxe_keygen_circuit();
        let expected_circuit = fs::read(format!("{}/mxe_keygen/circuit.arcis", ARTIFACTS_DIR))
            .expect("Failed to read stored circuit");
        // `assert_eq!` displays the objects in case of error.
        // Here the objects are unreadable and long.
        // The length causes problems in RustRover.
        assert!(circuit == expected_circuit, "Circuit mismatch");

        let deserialized: AsyncMPCCircuit =
            bincode::deserialize(&circuit).expect("Deserialization failed");
        deserialized.mock_eval_big_uint(Vec::new(), rng);

        let interface = mxe_keygen_interface();
        let expected_interface =
            fs::read_to_string(format!("{}/mxe_keygen/circuit.idarc", ARTIFACTS_DIR))
                .expect("Failed to read stored interface");
        assert_eq!(interface, expected_interface, "Interface mismatch");
    }
}
