use inco_lightning::cpi::accounts::{Operation, Allow};
use inco_lightning::cpi::{e_add, new_euint128, allow};

/// remaining_accounts:
///   [0] allowance_account (mut) - PDA for granting decrypt access
///   [1] owner_address (readonly) - The owner to grant access to
pub fn add_balances<'info>(
    ctx: Context<'_, '_, '_, 'info, AddBalances<'info>>,
    ciphertext_a: Vec<u8>,
    ciphertext_b: Vec<u8>,
) -> Result<()> {
    let cpi_program = ctx.accounts.inco_lightning_program.to_account_info();
    let signer = ctx.accounts.authority.to_account_info();

    // Convert ciphertexts to encrypted handles
    let cpi_ctx = CpiContext::new(cpi_program.clone(), Operation { signer: signer.clone() });
    let amount_a: Euint128 = new_euint128(cpi_ctx, ciphertext_a, 0)?;

    let cpi_ctx = CpiContext::new(cpi_program.clone(), Operation { signer: signer.clone() });
    let amount_b: Euint128 = new_euint128(cpi_ctx, ciphertext_b, 0)?;

    // Add two encrypted values
    let cpi_ctx = CpiContext::new(cpi_program.clone(), Operation { signer: signer.clone() });
    let result: Euint128 = e_add(cpi_ctx, amount_a, amount_b, 0)?;

    ctx.accounts.account.total = result;

    // Grant decryption access to owner via remaining_accounts
    if ctx.remaining_accounts.len() >= 2 {
        let allowance_account = &ctx.remaining_accounts[0];
        let allowed_address = &ctx.remaining_accounts[1];

        let cpi_ctx = CpiContext::new(
            cpi_program.clone(),
            Allow {
                allowance_account: allowance_account.clone(),
                signer: signer.clone(),
                allowed_address: allowed_address.clone(),
                system_program: ctx.accounts.system_program.to_account_info(),
            }
        );
        allow(cpi_ctx, result.0, true, ctx.accounts.account.owner)?;
    }

    Ok(())
}

#[derive(Accounts)]
pub struct AddBalances<'info> {
    #[account(mut)]
    pub authority: Signer<'info>,
    #[account(mut)]
    pub account: Account<'info, MyAccount>,
    /// CHECK: Inco Lightning program
    #[account(address = INCO_LIGHTNING_ID)]
    pub inco_lightning_program: AccountInfo<'info>,
    pub system_program: Program<'info, System>,
}

e_select(ctx: CpiContext, condition: Ebool, if_true: Euint128, if_false: Euint128, scalar_byte: u8) -> Result<Euint128>

use inco_lightning::cpi::accounts::{Operation, Allow};
use inco_lightning::cpi::{e_add, e_sub, e_ge, e_select, as_euint128, allow};

/// remaining_accounts:
///   [0] source_allowance_account (mut)
///   [1] source_owner_address (readonly)
///   [2] dest_allowance_account (mut)
///   [3] dest_owner_address (readonly)
pub fn confidential_transfer<'info>(
    ctx: Context<'_, '_, '_, 'info, Transfer<'info>>,
    transfer_amount: Euint128,
) -> Result<()> {
    let inco = ctx.accounts.inco_lightning_program.to_account_info();
    let signer = ctx.accounts.authority.to_account_info();

    let source_balance = ctx.accounts.source.balance;
    let dest_balance = ctx.accounts.destination.balance;

    // Check if source has sufficient balance (encrypted comparison)
    let cpi_ctx = CpiContext::new(inco.clone(), Operation { signer: signer.clone() });
    let has_sufficient: Ebool = e_ge(cpi_ctx, source_balance, transfer_amount, 0)?;

    // Create zero for failed transfer case
    let cpi_ctx = CpiContext::new(inco.clone(), Operation { signer: signer.clone() });
    let zero = as_euint128(cpi_ctx, 0)?;

    // Select actual transfer amount: if sufficient balance, use amount; else use 0
    let cpi_ctx = CpiContext::new(inco.clone(), Operation { signer: signer.clone() });
    let actual_amount: Euint128 = e_select(cpi_ctx, has_sufficient, transfer_amount, zero, 0)?;

    // Subtract from source (will subtract 0 if insufficient)
    let cpi_ctx = CpiContext::new(inco.clone(), Operation { signer: signer.clone() });
    let new_source_balance: Euint128 = e_sub(cpi_ctx, source_balance, actual_amount, 0)?;

    // Add to destination (will add 0 if insufficient)
    let cpi_ctx = CpiContext::new(inco.clone(), Operation { signer: signer.clone() });
    let new_dest_balance: Euint128 = e_add(cpi_ctx, dest_balance, actual_amount, 0)?;

    // Update balances
    ctx.accounts.source.balance = new_source_balance;
    ctx.accounts.destination.balance = new_dest_balance;

    // Grant allowance to source owner for their new balance
    if ctx.remaining_accounts.len() >= 2 {
        let cpi_ctx = CpiContext::new(
            inco.clone(),
            Allow {
                allowance_account: ctx.remaining_accounts[0].clone(),
                signer: signer.clone(),
                allowed_address: ctx.remaining_accounts[1].clone(),
                system_program: ctx.accounts.system_program.to_account_info(),
            }
        );
        allow(cpi_ctx, new_source_balance.0, true, ctx.accounts.source.owner)?;
    }

    // Grant allowance to destination owner for their new balance
    if ctx.remaining_accounts.len() >= 4 {
        let cpi_ctx = CpiContext::new(
            inco.clone(),
            Allow {
                allowance_account: ctx.remaining_accounts[2].clone(),
                signer: signer.clone(),
                allowed_address: ctx.remaining_accounts[3].clone(),
                system_program: ctx.accounts.system_program.to_account_info(),
            }
        );
        allow(cpi_ctx, new_dest_balance.0, true, ctx.accounts.destination.owner)?;
    }

    Ok(())
}

#[derive(Accounts)]
pub struct Transfer<'info> {
    #[account(mut)]
    pub authority: Signer<'info>,
    #[account(mut)]
    pub source: Account<'info, TokenAccount>,
    #[account(mut)]
    pub destination: Account<'info, TokenAccount>,
    /// CHECK: Inco Lightning program
    #[account(address = INCO_LIGHTNING_ID)]
    pub inco_lightning_program: AccountInfo<'info>,
    pub system_program: Program<'info, System>,
}

pub fn grant_decrypt_access(
    ctx: Context<GrantAccess>,
    handle: u128,
    allowed_address: Pubkey,
) -> Result<()> {
    let cpi_ctx = CpiContext::new(
        ctx.accounts.inco_lightning_program.to_account_info(),
        Allow {
            allowance_account: ctx.accounts.allowance_account.to_account_info(),
            signer: ctx.accounts.authority.to_account_info(),
            allowed_address: ctx.accounts.allowed_address.to_account_info(),
            system_program: ctx.accounts.system_program.to_account_info(),
        },
    );

    allow(cpi_ctx, handle, true, allowed_address)?;
    Ok(())
}

pub fn check_access(
    ctx: Context<CheckAccess>,
    handle: u128,
) -> Result<bool> {
    let cpi_ctx = CpiContext::new(
        ctx.accounts.inco_lightning_program.to_account_info(),
        IsAllowed {
            allowance_account: ctx.accounts.allowance_account.to_account_info(),
            allowed_address: ctx.accounts.allowed_address.to_account_info(),
        },
    );

    is_allowed(cpi_ctx, handle)
}

is_validsignature(
    ctx: CpiContext<VerifySignature>,
    expected_signature_count: u8,
    handles: Option<Vec<Vec<u8>>>,
    plaintext_values: Option<Vec<Vec<u8>>>,
) -> Result<Vec<SignatureVerificationResult>>

pub fn verify_decryption(
    ctx: Context<VerifyDecryption>,
    handles: Vec<Vec<u8>>,
    plaintext_values: Vec<Vec<u8>>,
) -> Result<()> {
    let cpi_ctx = CpiContext::new(
        ctx.accounts.inco_lightning_program.to_account_info(),
        VerifySignature {
            instructions: ctx.accounts.instructions.to_account_info(),
            signer: ctx.accounts.authority.to_account_info(),
        },
    );

    // Verify Ed25519 signatures from previous instruction
    // The covalidator signs hash(handle + plaintext_value)
    let results = is_validsignature(
        cpi_ctx,
        1,                      // expected signature count
        Some(handles),          // handles being verified
        Some(plaintext_values), // claimed plaintext values
    )?;

    // If we get here, signatures are valid
    // Use the verified plaintext values
    Ok(())
}