use crate::{
    error::{Error, HasSpan},
    interpreter::{
        data::{ArcisValue, CollectionKind, CurveData, Data, ValData},
        pointer::Pointer,
        supported_type::{CurveType, RustType, SupportedType},
        Interpreter,
    },
};
use arcis_compiler::{
    core::global_value::{curve_value::CurveValue, value::FieldValue},
    traits::ToMontgomery,
    utils::{
        crypto::{key::X25519PublicKey, rescue_cipher::RescueCipher},
        elliptic_curve::{decompress_montgomery_u, AffineEdwardsPoint},
        field::{BaseField, ScalarField},
    },
};
use quote::ToTokens;

pub fn decrypt_func(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    if args.len() != 2 {
        return Err(Error::new_spanned(
            expr,
            "Wrong number of arguments, should be a method call with 2 argument.",
        ));
    }
    let cipher_arg = args[0];
    let Data::Collection(cipher_struct, CollectionKind::Struct(cipher_id)) =
        interpreter.memory.get_underlying_data(cipher_arg)?
    else {
        return Err(Error::new_spanned(
            expr,
            "This method only works on Enc. It wasn't made with a Enc.",
        ));
    };
    let data_arg = args[1];
    let Data::Collection(data_struct, CollectionKind::Struct(data_id)) =
        interpreter.memory.get_underlying_data(data_arg)?
    else {
        return Err(Error::new_spanned(
            expr,
            "This method only works on Enc. It wasn't used on a Enc.",
        ));
    };

    let SupportedType::Array(output_type, _) = &data_id.types()[1] else {
        return Err(Error::new_spanned(
            expr,
            "This method only works on Enc. It wasn't made with a real Enc.",
        ));
    };
    let output_type = *output_type.clone();
    let Data::Collection(enc_addresses, CollectionKind::Array(_)) =
        interpreter.memory[data_struct[0]].clone()
    else {
        return Err(Error::new_spanned(
            expr,
            "This method only works on Enc. It wasn't made with a Enc. The array is wrong",
        ));
    };
    let encrypted_values: Result<Vec<ArcisValue>, _> = enc_addresses
        .into_iter()
        .map(|x| {
            let Data::Val(ValData(v, _)) = interpreter.memory[x] else {
                return Err(Error::new_spanned(
                    expr,
                    "This method only works on `Enc`. This wasn't one.",
                ));
            };
            Ok(v)
        })
        .collect();
    let encrypted_values = encrypted_values?;

    let decrypted_values: Vec<ArcisValue> = match cipher_id.get_name() {
        "Mxe" => {
            let Data::Val(ValData(nonce, _)) = interpreter.memory[cipher_struct[0]] else {
                return Err(Error::new_spanned(expr, "No Nonce found."));
            };
            let cipher = RescueCipher::new_for_mxe();
            cipher.decrypt(encrypted_values, nonce)
        }
        "Shared" => {
            let Data::Curve(CurveData(pubkey, CurveType::X25519Pubkey)) =
                interpreter.memory[cipher_struct[0]]
            else {
                return Err(Error::new_spanned(expr, "No x25519 pubkey found."));
            };
            let Data::Val(ValData(nonce, _)) = interpreter.memory[cipher_struct[1]] else {
                return Err(Error::new_spanned(expr, "No Nonce found."));
            };
            let cipher = RescueCipher::new_with_client::<
                FieldValue<BaseField>,
                FieldValue<ScalarField>,
                CurveValue,
            >(X25519PublicKey::new(pubkey, true));
            cipher.decrypt(encrypted_values, nonce)
        }
        _ => return Err(Error::new_spanned(expr, "Unsupported cipher.")),
    };
    let output_addr = interpreter.build_from_arcis_values(
        &mut decrypted_values.as_slice(),
        &output_type,
        true,
        expr.span(),
    )?;
    Ok(output_addr.as_pointer())
}

pub fn encrypt_mxe_func(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    if args.len() != 2 {
        return Err(Error::new_spanned(
            expr,
            "Wrong number of arguments, should be a method call with 0 arguments.",
        ));
    }
    if interpreter.memory.in_conditional_execution() {
        return Err(interpreter
            .backtrace
            .make_conditional_execution_error(expr, ".from_arcis(...)"));
    }

    let arg = args[1];
    let data = interpreter.pointer_to_arcis_values(arg)?;
    let cipher_data = interpreter.memory.get_underlying_data(args[0])?;
    let Data::Collection(cipher_addresses, CollectionKind::Struct(_)) = cipher_data.clone() else {
        return Err(Error::new_spanned(expr, "Not a cipher."));
    };
    if cipher_addresses.len() != 1 {
        return Err(Error::new_spanned(expr, "Wrong cipher."));
    }
    let Data::Val(ValData(nonce, _)) = interpreter.memory[cipher_addresses[0]] else {
        return Err(Error::new_spanned(expr, "No Nonce found."));
    };
    let cipher = RescueCipher::new_for_mxe();
    let new_data = cipher.encrypt(data, nonce);
    let mut output_data_addresses = Vec::with_capacity(new_data.len());
    output_data_addresses.extend(new_data.into_iter().map(|val| {
        interpreter
            .memory
            .push(Data::Val(ValData(val, RustType::ArcisField)))
    }));
    let output_data = Data::Collection(
        output_data_addresses,
        CollectionKind::Array(ciphertext_type()),
    );
    Ok(interpreter.memory.push(output_data).as_pointer())
}

fn ciphertext_type() -> SupportedType {
    SupportedType::FieldValue(RustType::ArcisField)
}

pub fn encrypt_client_func(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    if args.len() != 2 {
        return Err(Error::new_spanned(
            expr,
            "Wrong number of arguments, should be a method call with 2 arguments.",
        ));
    }
    if interpreter.memory.in_conditional_execution() {
        return Err(interpreter
            .backtrace
            .make_conditional_execution_error(expr, ".from_arcis(...)"));
    }

    let arg = args[1];
    let data = interpreter.pointer_to_arcis_values(arg)?;
    let cipher_data = interpreter.memory.get_underlying_data(args[0])?;
    let Data::Collection(cipher_addresses, CollectionKind::Struct(_)) = cipher_data.clone() else {
        return Err(Error::new_spanned(expr, "Not a cipher."));
    };
    if cipher_addresses.len() != 2 {
        return Err(Error::new_spanned(expr, "Wrong cipher."));
    }

    let Data::Curve(CurveData(pubkey, CurveType::X25519Pubkey)) =
        interpreter.memory[cipher_addresses[0]]
    else {
        return Err(Error::new_spanned(expr, "No x25519 pubkey found."));
    };
    let cipher = RescueCipher::new_with_client::<
        FieldValue<BaseField>,
        FieldValue<ScalarField>,
        CurveValue,
    >(X25519PublicKey::new(pubkey, true));
    let Data::Val(ValData(nonce, _)) = interpreter.memory[cipher_addresses[1]] else {
        return Err(Error::new_spanned(expr, "No Nonce found."));
    };
    let new_data = cipher.encrypt(data, nonce);
    let mut output_data_addresses = Vec::with_capacity(new_data.len());
    output_data_addresses.extend(new_data.into_iter().map(|val| {
        interpreter
            .memory
            .push(Data::Val(ValData(val, RustType::ArcisField)))
    }));
    let output_data = Data::Collection(
        output_data_addresses,
        CollectionKind::Array(ciphertext_type()),
    );
    let res = interpreter.memory.push(output_data).as_pointer();
    Ok(res)
}

pub fn x_to_pubkey(x: ArcisValue) -> CurveValue {
    // Internally we use (u, v) for Montgomery coordinates, to distinguish from (x, y) which we use
    // for Edwards coordinates (see e.g. https://en.wikipedia.org/wiki/Edwards_curve#The_group_law).
    let (_, (u, v)) = decompress_montgomery_u(x);
    CurveValue::from_affine(AffineEdwardsPoint::from_montgomery((u, v)))
}

/// Construct a new ArcisX25519Pubkey from the x-coordinate of the Montgomery point.
pub fn x25519_pubkey_from_x(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    if args.len() != 1 {
        return Err(Error::new_spanned(expr, "Wrong number of arguments."));
    }
    let Data::Val(ValData(x, _)) = interpreter.memory.get_underlying_data(args[0])? else {
        return Err(Error::new_spanned(expr, "No x-coordinate found."));
    };
    let point = x_to_pubkey(x);
    let point_addr = interpreter
        .memory
        .push(Data::Curve(CurveData(point, CurveType::X25519Pubkey)));
    Ok(point_addr.as_pointer())
}

/// Return the Montgomery x-coordinate of the ArcisX25519Pubkey.
pub fn x25519_pubkey_to_x(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    if args.len() != 1 {
        return Err(Error::new_spanned(expr, "Wrong number of arguments."));
    }
    let Data::Curve(CurveData(point, CurveType::X25519Pubkey)) =
        interpreter.memory.get_underlying_data(args[0])?
    else {
        return Err(Error::new_spanned(expr, "No point found."));
    };
    // Internally we use (u, v) for Montgomery coordinates, to distinguish from (x, y) which we use
    // for Edwards coordinates (see e.g. https://en.wikipedia.org/wiki/Edwards_curve#The_group_law).
    let (u, _) = point.to_montgomery(false);
    let u_addr = interpreter
        .memory
        .push(Data::Val(ValData(u, RustType::BASE_FIELD)));
    Ok(u_addr.as_pointer())
}
