use dep::std;
use crate::structs;

fn field_to_u8(_value: Field) -> [u8; 32] {
    let _array = _value.to_be_bytes(32);
    let mut array: [u8; 32] = [0; 32];

    for i in 0 .. 32 {
        array[i] = _array[i];
    }

    array
}

fn hash<N>(data: [Field; N]) -> Field {
    std::hash::pedersen_hash(data)
}

#[export]
fn utxo_to_commitment(utxo: structs::UTXO_New) -> Field {
    utxo.to_commitment()
}

#[export]
fn pedersen_left_right(left: Field, right: Field) -> Field {
    std::hash::pedersen_hash([left, right])
}

#[export]
fn keccak_tx(tx: structs::PublicInputs) -> Field {
    let mut hash_array: [Field; 53] = [0; 53];
    hash_array[0] = tx.current_root;
    hash_array[1] = tx.utxo_root;
    hash_array[2] = tx.deposit_amount;
    hash_array[3] = tx.contract_only_inputs;
    hash_array[4] = tx.withdrawals;
    for i in 0..16 {
        hash_array[5 + i] = tx.commitment_in[i];
        hash_array[21 + i] = tx.commitment_out[i];
        hash_array[37 + i] = tx.nullifier_hashes[i];
    }
    let u8_array = tx_to_u8(hash_array);
    hash_to_field(std::hash::keccak256(u8_array, u8_array.len() as u32))
}

fn bytes_tx_without_deposit(tx: structs::PublicInputs) -> [u8; 1696] {
    let mut hash_array: [Field; 53] = [0; 53];
    hash_array[0] = tx.current_root;
    hash_array[1] = tx.utxo_root;
    hash_array[2] = tx.deposit_amount;
    hash_array[3] = tx.contract_only_inputs;
    hash_array[4] = tx.withdrawals;
    for i in 0..16 {
        hash_array[5 + i] = tx.commitment_in[i];
        hash_array[21 + i] = tx.commitment_out[i];
        hash_array[37 + i] = tx.nullifier_hashes[i];
    }
    tx_to_u8(hash_array)
}

#[export]
fn keccak_contract_only_inputs(contract_only_inputs: structs::ContractOnlyInputs) -> Field {
    let mut hash_array: [Field; 100] = [0; 100];
    hash_array[0] = contract_only_inputs.timestamp;
    hash_array[1] = contract_only_inputs.deadline;
    hash_array[2] = contract_only_inputs.signature_hash;
    hash_array[3] = contract_only_inputs.price_limit;
    for i in 0..16 {
        hash_array[4 + i] = contract_only_inputs.recipients[i];
        hash_array[20 + i] = contract_only_inputs.swap_amounts[i];
        hash_array[36 + i] = contract_only_inputs.uids[i];
        hash_array[52 + (i * 3)] = contract_only_inputs.encrypted_utxo[i].secret;
        hash_array[53 + (i * 3)] = contract_only_inputs.encrypted_utxo[i].amount;
        hash_array[54 + (i * 3)] = contract_only_inputs.encrypted_utxo[i].data;
    }
    let u8_array = contract_only_to_u8(hash_array);
    hash_to_field(std::hash::keccak256(u8_array, u8_array.len() as u32))
}

#[export]
fn keccak_contract_only_inputs_without_deposit(contract_only_inputs: structs::ContractOnlyInputs) -> Field {
    let mut hash_array: [Field; 99] = [0; 99];
    hash_array[0] = contract_only_inputs.timestamp;
    hash_array[1] = contract_only_inputs.deadline;
    hash_array[2] = contract_only_inputs.price_limit;
    for i in 0..16 {
        hash_array[3 + i] = contract_only_inputs.recipients[i];
        hash_array[19 + i] = contract_only_inputs.swap_amounts[i];
        hash_array[35 + i] = contract_only_inputs.uids[i];
        hash_array[51 + (i * 3)] = contract_only_inputs.encrypted_utxo[i].secret;
        hash_array[52 + (i * 3)] = contract_only_inputs.encrypted_utxo[i].amount;
        hash_array[53 + (i * 3)] = contract_only_inputs.encrypted_utxo[i].data;
    }
    let u8_array = contract_only_without_deposit_to_u8(hash_array);
    hash_to_field(std::hash::keccak256(u8_array, u8_array.len() as u32))
}

fn bytes_contract_only_inputs_without_deposit(contract_only_inputs: structs::ContractOnlyInputs) -> [u8; 3168] {
    let mut hash_array: [Field; 99] = [0; 99];
    hash_array[0] = contract_only_inputs.timestamp;
    hash_array[1] = contract_only_inputs.deadline;
    hash_array[2] = contract_only_inputs.price_limit;
    for i in 0..16 {
        hash_array[3 + i] = contract_only_inputs.recipients[i];
        hash_array[19 + i] = contract_only_inputs.swap_amounts[i];
        hash_array[35 + i] = contract_only_inputs.uids[i];
        hash_array[51 + (i * 3)] = contract_only_inputs.encrypted_utxo[i].secret;
        hash_array[52 + (i * 3)] = contract_only_inputs.encrypted_utxo[i].amount;
        hash_array[53 + (i * 3)] = contract_only_inputs.encrypted_utxo[i].data;
    }
    contract_only_without_deposit_to_u8(hash_array)
}

#[export]
fn contract_only_inputs_with_signature_hash(contract_only_inputs: structs::ContractOnlyInputs) -> structs::ContractOnlyInputs {
    let mut hash_array: [Field; 99] = [0; 99];
    hash_array[0] = contract_only_inputs.timestamp;
    hash_array[1] = contract_only_inputs.deadline;
    hash_array[2] = contract_only_inputs.price_limit;
    for i in 0..16 {
        hash_array[3 + i] = contract_only_inputs.recipients[i];
        hash_array[19 + i] = contract_only_inputs.swap_amounts[i];
        hash_array[35 + i] = contract_only_inputs.uids[i];
        hash_array[51 + (i * 3)] = contract_only_inputs.encrypted_utxo[i].secret;
        hash_array[52 + (i * 3)] = contract_only_inputs.encrypted_utxo[i].amount;
        hash_array[53 + (i * 3)] = contract_only_inputs.encrypted_utxo[i].data;
    }
    let u8_array = contract_only_without_deposit_to_u8(hash_array);
    let contract_only_inputs_with_hash = structs::ContractOnlyInputs {
        timestamp: contract_only_inputs.timestamp,
        deadline: contract_only_inputs.deadline,
        signature_hash: hash_to_field(std::hash::keccak256(u8_array, u8_array.len() as u32)),
        price_limit: contract_only_inputs.price_limit,
        recipients: contract_only_inputs.recipients,
        swap_amounts: contract_only_inputs.swap_amounts,
        uids: contract_only_inputs.uids,
        encrypted_utxo: contract_only_inputs.encrypted_utxo
    };
    contract_only_inputs_with_hash
}

fn hash_tx(tx: structs::PublicInputs) -> Field {
    let mut hash_array: [Field; 53] = [0; 53];
    hash_array[0] = tx.current_root;
    hash_array[1] = tx.utxo_root;
    hash_array[2] = tx.deposit_amount;
    hash_array[3] = tx.withdrawals;
    for i in 0..16 {
        hash_array[4 + i] = tx.commitment_in[i];
        hash_array[20 + i] = tx.commitment_out[i];
        hash_array[36 + i] = tx.nullifier_hashes[i];
    }
    hash_array[52] = tx.contract_only_inputs;
    hash(hash_array)
}

fn hash_to_field(hash: [u8; 32]) -> Field {
    let mut keccak_field: Field = 0;    
    for p in 0..32 {
        let bytes_field: Field = hash[31 - p] as Field;
        keccak_field += bytes_field * 256.pow_32(p as Field);
    }

    keccak_field
}


fn tx_to_u8(pi_fields: [Field; 53]) -> [u8; 1696] {
    let mut keccak_array: [u8; 1696] = [0; 1696];
    for i in 0..pi_fields.len() {
        let mut byte_slice = pi_fields[i].to_be_bytes(32);
        for j in 0..32 {
            keccak_array[32*i + j] = byte_slice[j];
        }
    }

    keccak_array
}

fn contract_only_to_u8(pi_fields: [Field; 100]) -> [u8; 3200] {
    let mut keccak_array: [u8; 3200] = [0; 3200];
    for i in 0..pi_fields.len() {
        let mut byte_slice = pi_fields[i].to_be_bytes(32);
        for j in 0..32 {
            keccak_array[32*i + j] = byte_slice[j];
        }
    }

    keccak_array
}

fn contract_only_without_deposit_to_u8(pi_fields: [Field; 99]) -> [u8; 3168] {
    let mut keccak_array: [u8; 3168] = [0; 3168];
    for i in 0..pi_fields.len() {
        let mut byte_slice = pi_fields[i].to_be_bytes(32);
        for j in 0..32 {
            keccak_array[32*i + j] = byte_slice[j];
        }
    }

    keccak_array
}

fn batch_to_u8(pi_fields: [Field; 19]) -> [u8; 608] {
    let mut keccak_array: [u8; 608] = [0; 608];
    for i in 0..pi_fields.len() {
        let mut byte_slice = pi_fields[i].to_be_bytes(32);
        for j in 0..32 {
            keccak_array[32*i + j] = byte_slice[j];
        }
    }

    keccak_array
}

fn hash_tree_four(leaves: [Field; 16]) -> Field {
    let mut tx_tree: [Field; 16] = leaves;
    for l in 0..8 { tx_tree[l] = hash([tx_tree[2*l], tx_tree[2*l + 1]]); }
    for l in 0..4 { tx_tree[l] = hash([tx_tree[2*l], tx_tree[2*l + 1]]); }
    for l in 0..2 { tx_tree[l] = hash([tx_tree[2*l], tx_tree[2*l + 1]]); }
    hash([tx_tree[0], tx_tree[1]])
}

fn compute_merkle_root<N>(leaf: Field, index: Field, hash_path: [Field; N]) -> Field {
    let n = hash_path.len();
    let index_bits = index.to_le_bits(n as u32);
    let mut current = leaf;
    for i in 0..n {
        let path_bit = index_bits[i] as bool;
        let (hash_left, hash_right) = if path_bit {
            (hash_path[i], current)
        } else {
            (current, hash_path[i])
        };
        current = hash([hash_left, hash_right]);
    }
    current
}

fn compute_sibling_path<N>(sibling_path: [Field; N], new_leaf: Field, insertion_index: Field) -> [Field; N] {
    let index_bits: [u1] = insertion_index.to_le_bits(N as u32);
    let mut new_sibling_path: [Field; N] = [0; N];
    let mut current_hash: Field = new_leaf; 
    let mut zero_found: bool = false;

    for i in 0..N {
        let path_bit = index_bits[i] as bool;
        if (!zero_found) {
            if (!path_bit) {
                zero_found = true;
                new_sibling_path[i] = current_hash;
            } else {
                new_sibling_path[i] = structs::zero_hashes[i];
            }
        } else {
            new_sibling_path[i] = sibling_path[i];
        }

        if (path_bit) {
            current_hash = hash([sibling_path[i], current_hash]);
        } else {
            current_hash = hash([current_hash, sibling_path[i]]);
        }
    }
    new_sibling_path

}