use crate::{
    core::{
        circuits::{
            boolean::{boolean_value::BooleanValue, ed25519::Ed25519SecretKey},
            key_recovery::{
                init::{hash_keys, key_encryption_init, key_recovery_init},
                MXE_KEYS_ENC_COUNT,
            },
        },
        compile::*,
        expressions::{
            curve_expr::{self, CurveExpr},
            field_expr::{FieldExpr, InputInfo},
            InputKind,
        },
        global_value::{
            curve_array::CurveArray,
            curve_value::CurveValue,
            global_expr_store::with_local_expr_store_as_global,
            value::FieldValue,
        },
        ir_builder::{ExprStore, IRBuilder},
    },
    utils::{
        crypto::key::{
            AES128Key,
            AES192Key,
            AES256Key,
            RescueKey,
            X25519PrivateKey,
            X25519PublicKey,
            RESCUE_KEY_COUNT,
        },
        field::{BaseField, ScalarField},
        used_field::UsedField,
        zkp::elgamal::ElGamalKeypair,
    },
};
use arcis_interface::{CircuitInterface, ScalarKind, Value};
use core_utils::key_recovery::{MXE_KEY_RECOVERY_D, MXE_KEY_RECOVERY_N};
use ff::PrimeField;
use std::rc::Rc;

pub fn mxe_key_recovery_init_circuit() -> Vec<u8> {
    let mut expr_store = IRBuilder::new(true);
    let peer_x25519_pubkey_ids = (0..MXE_KEY_RECOVERY_N)
        .map(|i| {
            <IRBuilder as ExprStore<BaseField>>::push_curve(
                &mut expr_store,
                CurveExpr::Input(
                    i,
                    Rc::new(curve_expr::InputInfo::from(InputKind::Plaintext)),
                ),
            )
        })
        .collect::<Vec<usize>>();
    let mut offset = MXE_KEY_RECOVERY_N;
    let nonce_id = expr_store.push_field(FieldExpr::Input(
        offset,
        Rc::new(InputInfo::<BaseField> {
            kind: InputKind::Plaintext,
            min: BaseField::from(0),
            max: BaseField::power_of_two(128) - BaseField::from(1),
            ..InputInfo::default()
        }),
    ));
    offset += 1;
    let n_id = expr_store.push_field(FieldExpr::Input(
        offset,
        Rc::new(InputInfo::<BaseField> {
            kind: InputKind::Plaintext,
            min: BaseField::from(0),
            max: BaseField::from(MXE_KEY_RECOVERY_N as u64),
            ..InputInfo::default()
        }),
    ));
    offset += 1;
    let g_ids = (offset..offset + MXE_KEY_RECOVERY_D)
        .map(|i| {
            expr_store.push_field(FieldExpr::Input(
                i,
                Rc::new(InputInfo::<BaseField>::from(InputKind::Plaintext)),
            ))
        })
        .collect::<Vec<usize>>();

    let output_ids = with_local_expr_store_as_global(
        || {
            let peer_x25519_pubkeys = X25519PublicKey::new(
                CurveArray::from(
                    TryInto::<[CurveValue; MXE_KEY_RECOVERY_N]>::try_into(
                        peer_x25519_pubkey_ids
                            .into_iter()
                            .map(CurveValue::new)
                            .collect::<Vec<CurveValue>>(),
                    )
                    .unwrap_or_else(|v: Vec<CurveValue>| {
                        panic!(
                            "Expected a Vec of length {} (found {})",
                            MXE_KEY_RECOVERY_N,
                            v.len()
                        )
                    }),
                ),
                true,
            );
            let nonce = FieldValue::<BaseField>::from_id(nonce_id);
            let n = FieldValue::<BaseField>::from_id(n_id);
            let g = g_ids
                .into_iter()
                .map(FieldValue::<BaseField>::from_id)
                .collect::<Vec<FieldValue<BaseField>>>()
                .try_into()
                .unwrap_or_else(|v: Vec<FieldValue<BaseField>>| {
                    panic!(
                        "Expected a Vec of length {} (found {})",
                        MXE_KEY_RECOVERY_D,
                        v.len()
                    )
                });

            let base_mxe_x25519_private_key =
                X25519PrivateKey::<FieldValue<ScalarField>>::mxe_private_key();
            let mxe_rescue_base_field_key = RescueKey::<FieldValue<BaseField>>::mxe_rescue_key();
            let mxe_rescue_scalar_field_key =
                RescueKey::<FieldValue<ScalarField>>::mxe_rescue_key();
            let mxe_aes_128_key = AES128Key::<BooleanValue>::mxe_aes_key();
            let mxe_aes_192_key = AES192Key::<BooleanValue>::mxe_aes_key();
            let mxe_aes_256_key = AES256Key::<BooleanValue>::mxe_aes_key();
            let mxe_ed25519_secret_key = Ed25519SecretKey::mxe_secret_key();
            let mxe_elgamal_keypair = ElGamalKeypair::mxe_keypair();
            let mxe_elgamal_secret_key = *mxe_elgamal_keypair.secret();

            let rescue_key_shares_enc = key_recovery_init::<MXE_KEY_RECOVERY_N, MXE_KEY_RECOVERY_D>(
                base_mxe_x25519_private_key,
                mxe_rescue_base_field_key,
                peer_x25519_pubkeys,
                nonce,
                n,
                g,
            );

            let mxe_keys_enc = key_encryption_init(
                mxe_rescue_base_field_key,
                nonce,
                base_mxe_x25519_private_key,
                mxe_rescue_scalar_field_key,
                mxe_aes_128_key,
                mxe_aes_192_key,
                mxe_aes_256_key,
                mxe_ed25519_secret_key,
                mxe_elgamal_secret_key,
            );

            let mxe_keys_digest = hash_keys(
                base_mxe_x25519_private_key,
                mxe_rescue_base_field_key,
                mxe_rescue_scalar_field_key,
                mxe_aes_128_key,
                mxe_aes_192_key,
                mxe_aes_256_key,
                mxe_ed25519_secret_key,
                mxe_elgamal_secret_key,
            );

            let mut output_ids = (0..MXE_KEY_RECOVERY_N)
                .flat_map(|i| {
                    rescue_key_shares_enc
                        .iter()
                        .map(|share| share[i].get_id())
                        .collect::<Vec<usize>>()
                })
                .collect::<Vec<usize>>();
            output_ids.extend(mxe_keys_enc.map(|key| key.get_id()));
            output_ids.extend(mxe_keys_digest.map(|byte| byte.get_id()));
            output_ids.push(nonce.get_id());
            output_ids.push(n.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_key_recovery_init_interface() -> String {
    let base_field_size = BaseField::NUM_BITS as usize;

    // input types
    let pubkey_types = vec![Value::ArcisX25519Pubkey; MXE_KEY_RECOVERY_N];
    let nonce_type = Value::Scalar {
        size_in_bits: 128,
        kind: ScalarKind::Unsigned,
    };
    let n_type = Value::Scalar {
        size_in_bits: 16,
        kind: ScalarKind::Unsigned,
    };
    let g_types = vec![
        Value::Scalar {
            size_in_bits: base_field_size,
            kind: ScalarKind::Unsigned,
        };
        MXE_KEY_RECOVERY_D
    ];

    let mut input_types = pubkey_types;
    input_types.extend(vec![nonce_type.clone(), n_type.clone()]);
    input_types.extend(g_types);

    // output types
    // the encrypted Rescue key shares
    let mut output_types = vec![
        Value::Ciphertext {
            size_in_bits: base_field_size,
        };
        MXE_KEY_RECOVERY_N * RESCUE_KEY_COUNT
    ];
    // the encrypted MXE keys (x25519 private key, Rescue scalar field key, AES keys, ed25519 secret
    // key and ElGamal secret key)
    output_types.extend(vec![
        Value::Ciphertext {
            size_in_bits: base_field_size,
        };
        MXE_KEYS_ENC_COUNT
    ]);
    // the SHA3-256 digest of all MXE keys
    output_types.extend(vec![
        Value::Scalar {
            size_in_bits: 8,
            kind: ScalarKind::Unsigned
        };
        32
    ]);
    // we also output the nonce and n
    input_types.extend(vec![nonce_type, n_type]);

    let interface = CircuitInterface::new(
        "mxe_key_recovery_init".to_string(),
        input_types,
        output_types,
    );
    interface.serialize().unwrap()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        core::circuits::pre_compiled::constants::ARTIFACTS_DIR,
        utils::{
            crypto::key::{
                AES_128_KEY_COUNT,
                AES_192_KEY_COUNT,
                AES_256_KEY_COUNT,
                ED25519_SECRET_KEY_COUNT,
                ELGAMAL_SECRET_KEY_COUNT,
                X25519_PRIVATE_KEY_COUNT,
            },
            curve_point::CurvePoint,
        },
        AsyncMPCCircuit,
    };
    use num_bigint::BigUint;
    use num_traits::FromBytes;
    use std::fs;

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

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

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

    #[test]
    fn test_mxe_key_recovery_init_circuit() {
        let rng = &mut crate::utils::test_rng::get();
        let circuit = mxe_key_recovery_init_circuit();
        let expected_circuit = fs::read(format!(
            "{}/mxe_key_recovery_init/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");
        // for the pubkeys we need to provide valid points
        let mut inputs = vec![
            BigUint::from_le_bytes(CurvePoint::generator().to_bytes().as_ref());
            MXE_KEY_RECOVERY_N
        ];
        // for the nonce, n, g, x25519 private key, Rescue base field key, Rescue
        // scalar field key, AES keys, ed25519 secret key and ElGamal secret key can simply
        // set 0
        inputs.extend(vec![
            BigUint::from(0u8);
            1 + 1
                + MXE_KEY_RECOVERY_D
                + X25519_PRIVATE_KEY_COUNT
                + 2 * RESCUE_KEY_COUNT
                + AES_128_KEY_COUNT
                + AES_192_KEY_COUNT
                + AES_256_KEY_COUNT
                + ED25519_SECRET_KEY_COUNT
                + ELGAMAL_SECRET_KEY_COUNT
        ]);
        deserialized.mock_eval_big_uint(inputs, rng);

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