use crate::{
    core::{
        circuits::key_recovery::{
            finalize::{key_decryption_finalize, key_recovery_finalize},
            init::hash_keys,
            MXE_KEYS_ENC_COUNT,
        },
        compile::*,
        expressions::{
            curve_expr::{self, CurveExpr},
            field_expr::{FieldExpr, InputInfo},
            InputKind,
        },
        global_value::{
            curve_array::CurveArray,
            curve_value::CurveValue,
            field_array::FieldArray,
            global_expr_store::with_local_expr_store_as_global,
            value::FieldValue,
        },
        ir_builder::{ExprStore, IRBuilder},
    },
    utils::{
        crypto::key::{
            X25519PrivateKey,
            X25519PublicKey,
            ED25519_VERIFYING_KEY_COUNT,
            RESCUE_KEY_COUNT,
        },
        field::{BaseField, ScalarField},
        used_field::UsedField,
    },
};
use arcis_interface::{CircuitInterface, ScalarKind, Value};
use core_utils::key_recovery::{MXE_KEY_RECOVERY_D, MXE_KEY_RECOVERY_K, MXE_KEY_RECOVERY_N};
use ff::PrimeField;
use std::rc::Rc;

pub fn mxe_key_recovery_finalize_circuit() -> Vec<u8> {
    let mut expr_store = IRBuilder::new(true);
    let inp_ids = (0..MXE_KEY_RECOVERY_N)
        .flat_map(|i| {
            let mut ids = (0..RESCUE_KEY_COUNT)
                .map(|j| {
                    // the encrypted key shares
                    expr_store.push_field(FieldExpr::Input(
                        i * (RESCUE_KEY_COUNT + 1) + j,
                        Rc::new(InputInfo::<BaseField>::from(InputKind::Plaintext)),
                    ))
                })
                .collect::<Vec<usize>>();
            // the public key
            ids.push(<IRBuilder as ExprStore<BaseField>>::push_curve(
                &mut expr_store,
                CurveExpr::Input(
                    i * (RESCUE_KEY_COUNT + 1) + RESCUE_KEY_COUNT,
                    Rc::new(curve_expr::InputInfo::from(InputKind::Plaintext)),
                ),
            ));
            ids
        })
        .collect::<Vec<usize>>();
    let mut offset = MXE_KEY_RECOVERY_N * (RESCUE_KEY_COUNT + 1);
    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 alpha_pow_ids = (offset..offset + MXE_KEY_RECOVERY_K)
        .map(|i| {
            expr_store.push_field(FieldExpr::Input(
                i,
                Rc::new(InputInfo::<BaseField>::from(InputKind::Plaintext)),
            ))
        })
        .collect::<Vec<usize>>();
    offset += MXE_KEY_RECOVERY_K;
    let scaled_polynomial_ids = (offset..offset + MXE_KEY_RECOVERY_K)
        .map(|i| {
            expr_store.push_field(FieldExpr::Input(
                i,
                Rc::new(InputInfo::<BaseField>::from(InputKind::Plaintext)),
            ))
        })
        .collect::<Vec<usize>>();
    offset += MXE_KEY_RECOVERY_K;
    let mxe_key_enc_ids = (offset..offset + MXE_KEYS_ENC_COUNT)
        .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 rescue_base_field_key_shares_enc = (0..RESCUE_KEY_COUNT)
                .map(|i| {
                    FieldArray::from(
                        TryInto::<[FieldValue<BaseField>; MXE_KEY_RECOVERY_N]>::try_into(
                            inp_ids
                                .iter()
                                .skip(i)
                                .step_by(RESCUE_KEY_COUNT + 1)
                                .map(|id| FieldValue::<BaseField>::from_id(*id))
                                .collect::<Vec<FieldValue<BaseField>>>(),
                        )
                        .unwrap_or_else(|v: Vec<FieldValue<BaseField>>| {
                            panic!(
                                "Expected a Vec of length {} (found {})",
                                MXE_KEY_RECOVERY_N,
                                v.len()
                            )
                        }),
                    )
                })
                .collect::<Vec<FieldArray<MXE_KEY_RECOVERY_N, BaseField>>>()
                .try_into()
                .unwrap_or_else(|v: Vec<FieldArray<MXE_KEY_RECOVERY_N, BaseField>>| {
                    panic!(
                        "Expected a Vec of length {} (found {})",
                        RESCUE_KEY_COUNT,
                        v.len()
                    )
                });
            let peer_x25519_pubkeys = X25519PublicKey::new(
                CurveArray::from(
                    TryInto::<[CurveValue; MXE_KEY_RECOVERY_N]>::try_into(
                        inp_ids
                            .into_iter()
                            .skip(RESCUE_KEY_COUNT)
                            .step_by(RESCUE_KEY_COUNT + 1)
                            .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 alpha_pows = FieldArray::from(
                TryInto::<[FieldValue<BaseField>; MXE_KEY_RECOVERY_K]>::try_into(
                    alpha_pow_ids
                        .into_iter()
                        .map(FieldValue::<BaseField>::from_id)
                        .collect::<Vec<FieldValue<BaseField>>>(),
                )
                .unwrap_or_else(|v: Vec<FieldValue<BaseField>>| {
                    panic!(
                        "Expected a Vec of length {} (found {})",
                        MXE_KEY_RECOVERY_K,
                        v.len()
                    )
                }),
            );
            let scaled_polynomials = FieldArray::from(
                TryInto::<[FieldValue<BaseField>; MXE_KEY_RECOVERY_K]>::try_into(
                    scaled_polynomial_ids
                        .into_iter()
                        .map(FieldValue::<BaseField>::from_id)
                        .collect::<Vec<FieldValue<BaseField>>>(),
                )
                .unwrap_or_else(|v: Vec<FieldValue<BaseField>>| {
                    panic!(
                        "Expected a Vec of length {} (found {})",
                        MXE_KEY_RECOVERY_K,
                        v.len()
                    )
                }),
            );
            let mxe_keys_enc = mxe_key_enc_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_KEYS_ENC_COUNT,
                        v.len()
                    )
                });

            let backup_mxe_x25519_private_key =
                X25519PrivateKey::<FieldValue<ScalarField>>::mxe_private_key();

            let (mxe_rescue_base_field_key, errors) = key_recovery_finalize::<
                MXE_KEY_RECOVERY_N,
                MXE_KEY_RECOVERY_K,
                { MXE_KEY_RECOVERY_D - 1 },
            >(
                backup_mxe_x25519_private_key,
                peer_x25519_pubkeys,
                nonce,
                rescue_base_field_key_shares_enc,
                n,
                alpha_pows,
                scaled_polynomials,
            );

            let (
                mxe_x25519_private_key,
                mxe_x25519_pubkey,
                mxe_rescue_scalar_field_key,
                mxe_aes_128_key,
                mxe_aes_192_key,
                mxe_aes_256_key,
                mxe_ed25519_secret_key,
                mxe_ed25519_signing_key,
                mxe_elgamal_keypair,
            ) = key_decryption_finalize(mxe_rescue_base_field_key, nonce, mxe_keys_enc);

            let mxe_keys_digest = hash_keys(
                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_keypair.secret(),
            );

            // x25519
            let mut output_ids = vec![
                mxe_x25519_private_key.inner().get_id(),
                mxe_x25519_pubkey.inner().get_id(),
            ];
            // Rescue
            output_ids.extend(
                mxe_rescue_base_field_key
                    .inner()
                    .into_iter()
                    .map(|val| val.get_id()),
            );
            output_ids.extend(
                mxe_rescue_scalar_field_key
                    .inner()
                    .into_iter()
                    .map(|val| val.get_id()),
            );
            // aes
            output_ids.extend(
                mxe_aes_128_key
                    .inner()
                    .into_iter()
                    .flat_map(|byte| byte.to_vec())
                    .map(|bit| bit.get_id()),
            );
            output_ids.extend(
                mxe_aes_192_key
                    .inner()
                    .into_iter()
                    .flat_map(|byte| byte.to_vec())
                    .map(|bit| bit.get_id()),
            );
            output_ids.extend(
                mxe_aes_256_key
                    .inner()
                    .into_iter()
                    .flat_map(|byte| byte.to_vec())
                    .map(|bit| bit.get_id()),
            );
            // ed25519
            output_ids.extend(
                mxe_ed25519_secret_key
                    .inner()
                    .into_iter()
                    .flat_map(|byte| byte.to_vec())
                    .map(|bit| bit.get_id()),
            );
            output_ids.push(mxe_ed25519_signing_key.s.get_id());
            output_ids.extend(
                mxe_ed25519_signing_key
                    .hash_prefix
                    .into_iter()
                    .flat_map(|byte| byte.to_vec())
                    .map(|bit| bit.get_id()),
            );
            output_ids.extend(
                mxe_ed25519_signing_key
                    .verifying_key
                    .public_key_encoded
                    .into_iter()
                    .map(|byte| FieldValue::<BaseField>::from(byte).get_id()),
            );
            // ElGamal
            output_ids.push(mxe_elgamal_keypair.secret().get_scalar().get_id());
            output_ids.push(mxe_elgamal_keypair.pubkey().get_point().get_id());

            // digest of all MXE keys
            output_ids.extend(mxe_keys_digest.map(|byte| byte.get_id()));

            // errors
            output_ids.extend(
                (0..MXE_KEY_RECOVERY_N).flat_map(|i| errors.map(|error| error[i].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_finalize_interface() -> String {
    let base_field_size = BaseField::NUM_BITS as usize;

    // input types
    let inp_types = (0..MXE_KEY_RECOVERY_N)
        .flat_map(|_| {
            let mut types = vec![
                Value::Ciphertext {
                    size_in_bits: base_field_size,
                };
                RESCUE_KEY_COUNT
            ];
            types.push(Value::ArcisX25519Pubkey);
            types
        })
        .collect::<Vec<Value>>();
    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 alpha_pow_types = vec![
        Value::Scalar {
            size_in_bits: base_field_size,
            kind: ScalarKind::Unsigned,
        };
        MXE_KEY_RECOVERY_K
    ];
    let scaled_polynomial_types = vec![
        Value::Scalar {
            size_in_bits: base_field_size,
            kind: ScalarKind::Unsigned,
        };
        MXE_KEY_RECOVERY_K
    ];
    let mxe_key_enc_types = vec![
        Value::Ciphertext {
            size_in_bits: base_field_size,
        };
        MXE_KEYS_ENC_COUNT
    ];

    let mut input_types = inp_types;
    input_types.extend(vec![nonce_type.clone(), n_type.clone()]);
    input_types.extend(alpha_pow_types);
    input_types.extend(scaled_polynomial_types);
    input_types.extend(mxe_key_enc_types);

    // output types
    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 keys_digest_type = vec![
        Value::Scalar {
            size_in_bits: 8,
            kind: ScalarKind::Unsigned
        };
        32
    ];
    let error_types = vec![
        Value::Scalar {
            size_in_bits: base_field_size,
            kind: ScalarKind::Unsigned
        };
        MXE_KEY_RECOVERY_N * RESCUE_KEY_COUNT
    ];

    let mut output_types = vec![x25519_pubkey_type];
    output_types.extend(ed25519_verifying_key);
    output_types.push(elgamal_pubkey_type);
    output_types.extend(keys_digest_type);
    output_types.extend(error_types);

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        core::circuits::{
            key_recovery::MXE_KEYS_ENC_COUNT,
            pre_compiled::constants::ARTIFACTS_DIR,
        },
        utils::{crypto::key::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_finalize_circuit() {
        fs::create_dir_all(ARTIFACTS_DIR).expect("Failed to create artifacts directory");

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

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

    #[test]
    fn test_mxe_key_recovery_finalize_circuit() {
        let rng = &mut crate::utils::test_rng::get();
        let circuit = mxe_key_recovery_finalize_circuit();
        let expected_circuit = fs::read(format!(
            "{}/mxe_key_recovery_finalize/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");
        let mut inputs = Vec::new();
        // for the key shares we can set 0, but for the public key we need to provide valid
        // curve points
        for _ in 0..MXE_KEY_RECOVERY_N {
            let mut shares = vec![BigUint::from(0u8); RESCUE_KEY_COUNT];
            shares.push(BigUint::from_le_bytes(
                CurvePoint::generator().to_bytes().as_ref(),
            ));
            inputs.extend(shares);
        }
        // for the nonce we can set 0
        inputs.push(BigUint::from(0u8));
        // for n we must set a positive number
        inputs.push(BigUint::from(1u8));
        // for alpha_pows, scaled_polynomials, mxe_keys_enc and the x25519 private key we can again
        // set 0
        inputs.extend(vec![
            BigUint::from(0u8);
            MXE_KEY_RECOVERY_K
                + MXE_KEY_RECOVERY_K
                + MXE_KEYS_ENC_COUNT
                + X25519_PRIVATE_KEY_COUNT
        ]);
        deserialized.mock_eval_big_uint(inputs, rng);

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