airdrop 0.1.0

use airdrop_api::{
    consts::{CAMPAIGN_STATUS_ACTIVE, MAX_MERKLE_PROOF_LENGTH, TREASURY},
    instruction::Claim,
    loaders::AirdropAccountInfoValidation,
    merkle::{verify_merkle_proof, MerkleProof},
    pda::{allocation_pda, campaign_pda, campaign_treasury_pda},
    prelude::*,
};
use airdrop_api::{Allocation, Campaign};
use solana_program::clock::Clock;
use steel::*;

pub fn process_claim(accounts: &[AccountInfo<'_>], data: &[u8]) -> ProgramResult {
    // Parse instruction data
    let args = Claim::try_from_bytes(data)?;
    let campaign_id = args.campaign_id;
    let method = args.method;
    let amount = u64::from_le_bytes(args.amount);

    // Load accounts
    // For on-chain mode: [recipient, campaign, allocation, treasury, treasury_tokens, recipient_tokens, mint, token_program, associated_token_program, system_program]
    // For off-chain mode: [recipient, campaign, allocation_or_empty, treasury, treasury_tokens, recipient_tokens, mint, token_program, associated_token_program, system_program]
    // For merkle mode: [recipient, campaign, allocation_or_empty, treasury, treasury_tokens, recipient_tokens, mint, token_program, associated_token_program, system_program]
    let [recipient_info, campaign_info, allocation_or_empty_info, treasury_info, treasury_tokens_info, recipient_tokens_info, mint_info, token_program, associated_token_program, system_program] =
        accounts
    else {
        return Err(ProgramError::NotEnoughAccountKeys);
    };

    // Validate recipient
    recipient_info.is_signer()?;

    // Verify campaign PDA matches campaign_id
    let (expected_campaign, _campaign_bump) = campaign_pda(&campaign_id);

    // Load and validate campaign
    let campaign = campaign_info
        .is_campaign()?
        .has_address(&expected_campaign)?
        .as_account_mut::<Campaign>(&airdrop_api::ID)?
        .assert_mut(|c| c.status == CAMPAIGN_STATUS_ACTIVE)?;

    // Verify treasury PDA
    let (expected_treasury, _treasury_bump) = campaign_treasury_pda(&campaign_id);
    treasury_info.has_address(&expected_treasury)?;

    // Verify treasury token account
    treasury_tokens_info
        .is_writable()?
        .has_address(&campaign.treasury)?;

    // Verify recipient token account (will be created if needed)
    recipient_tokens_info.is_writable()?;

    // Verify mint
    mint_info.has_address(&campaign.mint)?;

    token_program.is_program(&spl_token::ID)?;
    associated_token_program.is_program(&spl_associated_token_account::ID)?;
    system_program.is_program(&system_program::ID)?;

    // Process claim based on method
    // Method 0: fulfill_method='allocation' - Claim against on-chain Allocation account
    // Method 1: fulfill_method='record' - Claim against off-chain DB record (with signature)
    // Method 2: fulfill_method='merkle' - Claim against merkle proof
    let claim_amount = match method {
        0 => {
            // Method 0: fulfill_method='allocation' - Verify using Allocation account
            let allocation = allocation_or_empty_info
                .as_account_mut::<Allocation>(&airdrop_api::ID)?
                .assert_mut(|a| a.campaign == *campaign_info.key)?
                .assert_mut(|a| a.recipient == *recipient_info.key)?
                .assert_mut(|a| a.claimed == 0)?;

            // Verify allocation account PDA
            let (expected_allocation, _) = allocation_pda(*campaign_info.key, *recipient_info.key);
            allocation_or_empty_info.has_address(&expected_allocation)?;

            let claim_amount = allocation.amount;

            // Mark as claimed
            let clock = Clock::get()?;
            allocation.claimed = 1;
            allocation.claimed_at = clock.unix_timestamp;

            claim_amount
        }
        1 => {
            // Method 1: fulfill_method='record' - Amount is provided in instruction data
            // Off-chain validation (e.g., signature verification) should be done before calling this instruction
            if amount == 0 {
                return Err(AirdropError::InvalidAmount.into());
            }
            amount
        }
        2 => {
            // Method 2: fulfill_method='merkle' - Verify merkle proof against campaign's merkle root
            if campaign.merkle_root == [0u8; 32] {
                return Err(AirdropError::MerkleRootNotSet.into());
            }
            if amount == 0 {
                return Err(AirdropError::InvalidAmount.into());
            }

            // Parse merkle proof from instruction data
            let proof = parse_merkle_proof(data, args.merkle_proof_length)?;

            // Verify merkle proof
            if !verify_merkle_proof(&proof, recipient_info.key, amount, campaign.merkle_root) {
                return Err(AirdropError::InvalidMerkleProof.into());
            }

            amount
        }
        _ => return Err(AirdropError::InvalidAmount.into()),
    };

    // Create recipient token account if needed
    if recipient_tokens_info.data_is_empty() {
        create_associated_token_account(
            recipient_info,
            recipient_info,
            recipient_tokens_info,
            mint_info,
            system_program,
            token_program,
            associated_token_program,
        )?;
    }

    // Transfer tokens from treasury to recipient using PDA-signed transfer
    transfer_signed(
        treasury_info,
        treasury_tokens_info,
        recipient_tokens_info,
        token_program,
        claim_amount,
        &[TREASURY, &campaign_id],
    )?;

    // Update campaign totals
    campaign.total_claimed = campaign
        .total_claimed
        .checked_add(claim_amount)
        .ok_or(ProgramError::ArithmeticOverflow)?;

    // Emit event
    TokensClaimedEvent {
        campaign: *campaign_info.key,
        recipient: *recipient_info.key,
        amount: claim_amount,
        method,
        _padding: [0u8; 7],
    }
    .log();

    Ok(())
}

/// Parse merkle proof from instruction data
///
/// ## Data Layout (for merkle mode):
/// Fixed: [Claim struct]
/// Variable: [proof_path_hashes] + [proof_indices]
///   - proof_path_hashes: proof_length * 32 bytes
///   - proof_indices: proof_length bytes (0=false, 1=true)
fn parse_merkle_proof(data: &[u8], proof_length: u8) -> Result<MerkleProof, ProgramError> {
    // Validate proof length to prevent DoS attacks
    if proof_length > MAX_MERKLE_PROOF_LENGTH {
        return Err(AirdropError::InvalidMerkleProof.into());
    }

    let proof_length = proof_length as usize;
    let fixed_size = std::mem::size_of::<Claim>();
    let variable_data = &data[fixed_size..];

    let mut offset = 0;

    // Parse merkle proof PATH (32-byte hashes)
    let proof_path_size = proof_length * 32;
    if offset + proof_path_size > variable_data.len() {
        return Err(ProgramError::InvalidInstructionData);
    }

    let mut proof_path = Vec::new();
    for i in 0..proof_length {
        let start = offset + i * 32;
        let end = start + 32;
        let hash: [u8; 32] = variable_data[start..end]
            .try_into()
            .map_err(|_| ProgramError::InvalidInstructionData)?;
        proof_path.push(hash);
    }
    offset += proof_path_size;

    // Parse merkle proof INDICES (1-byte booleans: 0=false, 1=true)
    if offset + proof_length > variable_data.len() {
        return Err(ProgramError::InvalidInstructionData);
    }

    let mut proof_indices = Vec::new();
    for i in 0..proof_length {
        let index_byte = variable_data[offset + i];
        proof_indices.push(index_byte != 0);
    }

    Ok(MerkleProof::new(proof_path, proof_indices))
}