// Core dependencies
import { Program } from "@coral-xyz/anchor";
import { PublicKey, SystemProgram, TransactionInstruction } from "@solana/web3.js";
import { ASSOCIATED_TOKEN_PROGRAM_ID, TOKEN_PROGRAM_ID } from "@solana/spl-token";

// Local imports
import { BasketsProgram } from "../../idl/types";
import { getAta, getBasketPda, getStateCreatorAccount } from "../../utils/programAccounts";
import { fetchWithdrawState } from "../../state/withdrawState";

export async function claimTokensIxs(params: {
    program: Program<BasketsProgram>,
    payer: PublicKey, 
    withdrawState: PublicKey,
}): Promise<TransactionInstruction[]> {
    // Destructure params and fetch state data
    const { program, payer, withdrawState } = params;
    const withdrawStateData = await fetchWithdrawState(program, withdrawState);

    // Get key data from states
    const withdrawStateSeed = withdrawStateData.withdrawStateSeed;
    const basket = withdrawStateData.basket;
    const basketPda = getBasketPda(basket);
    const user = withdrawStateData.owner;
    const stateCreator = getStateCreatorAccount();

    // Build list of token mints to claim
    const tokenMints: PublicKey[] = [];
    
    // Add mints with non-zero amounts
    for (let i = 0; i < parseInt(withdrawStateData.numTokens.toString()); i++) {
        const amount = parseInt(withdrawStateData.compositionAmounts[i].toString());
        if (amount > 0) {
            tokenMints.push(withdrawStateData.compositionMints[i]);
        }
    }

    // Add destination mint if not already included and valid
    const destinationAmount = parseInt(withdrawStateData.destinationAmount.toString());
    const isDestinationMintIncluded = tokenMints.find(mint => 
        mint.equals(withdrawStateData.destinationMint)
    );
    
    if (!isDestinationMintIncluded && destinationAmount !== 0) {
        tokenMints.push(withdrawStateData.destinationMint);
    }

    // Get token accounts for basket and user
    const basketTokenAccounts = tokenMints.map(mint => getAta(basketPda, mint));
    const userTokenAccounts = tokenMints.map(mint => getAta(user, mint));

    // Build claim instructions for each token
    const ixs = await Promise.all(tokenMints.map((mint, index) => 
        program.methods
            .withdrawClaim(
                withdrawStateSeed
            )
            .accountsStrict({
                worker: payer,
                basket,
                basketPda,
                withdrawState,
                tokenMint: mint,
                basketTokenAccount: basketTokenAccounts[index],
                user,
                userTokenAccount: userTokenAccounts[index],
                stateCreator,
                systemProgram: SystemProgram.programId,
                tokenProgram: TOKEN_PROGRAM_ID,
                associatedTokenProgram: ASSOCIATED_TOKEN_PROGRAM_ID,
            })
            .instruction()
    ));

    return ixs;
}