use anchor_lang::prelude::*;
use arcium_client::idl::arcium::{
    cpi::{accounts::InitComputationDefinition, init_computation_definition},
    types::{
        Argument,
        CallbackAccount,
        ComputationDefinitionMeta,
        ComputationSignature,
        OffChainCircuitSource,
    },
};
use traits::{FinalizeCompAccs, InitCompDefAccs, QueueCompAccs};

pub mod traits;

// a struct matching this is present in arx/src/utils.rs
// when making changes here, make sure to update the arx version
#[derive(Debug, AnchorSerialize, AnchorDeserialize)]
pub enum ComputationOutputs {
    Bytes(Vec<u8>),
    Abort,
}

pub fn queue_computation<'info, T>(
    accs: &T,
    computation_offset: u64,
    args: Vec<Argument>,
    callback_accs: Vec<CallbackAccount>,
    callback_url: Option<String>,
) -> Result<()>
where
    T: QueueCompAccs<'info>,
{
    let cpi_context = CpiContext::new(accs.arcium_program(), accs.queue_comp_accs());
    arcium_client::idl::arcium::cpi::queue_computation(
        cpi_context,
        computation_offset,
        accs.comp_def_offset(),
        None,
        args,
        callback_accs,
        accs.mxe_program(),
        callback_url,
        0,
        0,
        0,
    )
}

pub fn finalize_computation<'info, T>(accs: &T, computation_offset: u64) -> Result<()>
where
    T: FinalizeCompAccs<'info>,
{
    let cpi_context = CpiContext::new(accs.arcium_program(), accs.finalize_comp_accs());
    arcium_client::idl::arcium::cpi::finalize_computation(
        cpi_context,
        computation_offset,
        accs.comp_def_offset(),
        accs.mxe_program(),
    )
}

pub fn init_comp_def<'info, T>(
    accs: &T,
    finalize_during_callback: bool,
    offchain_circuit_source_override: Option<OffChainCircuitSource>,
    finalize_authority: Option<Pubkey>,
) -> Result<()>
where
    T: InitCompDefAccs<'info>,
{
    let cpi_context = CpiContext::new(
        accs.arcium_program(),
        InitComputationDefinition {
            signer: accs.signer(),
            system_program: accs.system_program(),
            mxe: accs.mxe_acc(),
            comp_def_acc: accs.comp_def_acc(),
        },
    );

    let signature = ComputationSignature {
        parameters: accs.params(),
        outputs: accs.outputs(),
    };
    let computation_definition = ComputationDefinitionMeta {
        circuit_len: accs.compiled_circuit_len(),
        callback_discriminator: accs.circuit_callback_discriminator(),
        signature,
    };
    init_computation_definition(
        cpi_context,
        accs.comp_def_offset(),
        accs.mxe_program(),
        computation_definition,
        offchain_circuit_source_override,
        finalize_authority,
        finalize_during_callback,
    )?;

    Ok(())
}

#[macro_export]
macro_rules! derive_seed {
    ($name:ident) => {
        stringify!($name).as_bytes()
    };
}

pub const fn comp_def_offset(conf_ix_name: &str) -> u32 {
    let hasher = ::sha2_const_stable::Sha256::new();
    let result = hasher.update(conf_ix_name.as_bytes()).finalize();
    u32::from_le_bytes([result[0], result[1], result[2], result[3]])
}

pub const MXE_PDA_SEED: &[u8] = derive_seed!(PersistentMXEAccount);
pub const MEMPOOL_PDA_SEED: &[u8] = b"Mempool";
pub const EXECPOOL_PDA_SEED: &[u8] = b"Execpool";
pub const COMP_PDA_SEED: &[u8] = derive_seed!(ComputationAccount);
pub const COMP_DEF_PDA_SEED: &[u8] = derive_seed!(ComputationDefinitionAccount);
pub const CLUSTER_PDA_SEED: &[u8] = derive_seed!(Cluster);
pub const POOL_PDA_SEED: &[u8] = derive_seed!(StakingPoolAccount);
pub const CLOCK_PDA_SEED: &[u8] = derive_seed!(ClockAccount);

pub const ARCIUM_CLOCK_ACCOUNT_ADDRESS: Pubkey = Pubkey::new_from_array([
    212, 85, 34, 0, 53, 147, 95, 180, 158, 156, 108, 40, 138, 177, 241, 37, 193, 113, 49, 48, 98,
    57, 195, 10, 201, 244, 92, 111, 3, 191, 25, 130,
]);

pub const ARCIUM_STAKING_POOL_ACCOUNT_ADDRESS: Pubkey = Pubkey::new_from_array([
    186, 14, 179, 51, 245, 188, 165, 168, 140, 45, 43, 138, 19, 112, 49, 116, 43, 39, 160, 119,
    185, 51, 233, 185, 82, 73, 153, 192, 213, 226, 190, 150,
]);

#[macro_export]
macro_rules! derive_mxe_pda {
    () => {
        Pubkey::find_program_address(&[MXE_PDA_SEED, ID.to_bytes().as_ref()], &ARCIUM_PROG_ID).0
    };
}

#[macro_export]
macro_rules! derive_mempool_pda {
    () => {
        Pubkey::find_program_address(&[MEMPOOL_PDA_SEED, ID.to_bytes().as_ref()], &ARCIUM_PROG_ID).0
    };
}

#[macro_export]
macro_rules! derive_execpool_pda {
    () => {
        Pubkey::find_program_address(
            &[EXECPOOL_PDA_SEED, ID.to_bytes().as_ref()],
            &ARCIUM_PROG_ID,
        )
        .0
    };
}

#[macro_export]
macro_rules! derive_comp_pda {
    ($computation_offset:expr) => {
        Pubkey::find_program_address(
            &[
                COMP_PDA_SEED,
                ID.to_bytes().as_ref(),
                &$computation_offset.to_le_bytes(),
            ],
            &ARCIUM_PROG_ID,
        )
        .0
    };
}

#[macro_export]
macro_rules! derive_comp_def_pda {
    ($conf_ix_name:expr) => {
        Pubkey::find_program_address(
            &[
                COMP_DEF_PDA_SEED,
                &ID_CONST.to_bytes(),
                &$conf_ix_name.to_le_bytes(),
            ],
            &ARCIUM_PROG_ID,
        )
        .0
    };
}

#[macro_export]
macro_rules! derive_cluster_pda {
    ($mxe_account:expr) => {
        Pubkey::find_program_address(
            &[CLUSTER_PDA_SEED, &$mxe_account.cluster.offset.to_le_bytes()],
            &ARCIUM_PROG_ID,
        )
        .0
    };
}

#[cfg(test)]
mod tests {
    use super::*;
    use arcium_client::idl::arcium::ID_CONST as ARCIUM_PROG_ID;

    fn derive_arcium_pda(seeds: &[&[u8]]) -> Pubkey {
        Pubkey::find_program_address(seeds, &ARCIUM_PROG_ID).0
    }

    #[test]
    fn test_comp_def_offset() {
        let conf_ix_name = "add_together";
        let offset = comp_def_offset(conf_ix_name);
        assert_eq!(offset, 4005749700);
    }

    #[test]
    fn test_clock_account_address() {
        let address = derive_arcium_pda(&[CLOCK_PDA_SEED]);
        assert_eq!(address, ARCIUM_CLOCK_ACCOUNT_ADDRESS);
    }

    #[test]
    fn test_staking_pool_account_address() {
        let address = derive_arcium_pda(&[POOL_PDA_SEED]);
        assert_eq!(address, ARCIUM_STAKING_POOL_ACCOUNT_ADDRESS);
    }
}
