import {
    ASSOCIATED_TOKEN_PROGRAM_ID, TOKEN_2022_PROGRAM_ID, TOKEN_PROGRAM_ID
} from '@solana/spl-token';
import {  PublicKey, SystemProgram, TransactionInstruction } from '@solana/web3.js';

import { BASKETS_V3_PROGRAM_ID } from '../../constants';
import {
    getAta, getBasketFeesPda, getBasketState, getGlobalConfigPda, getWithdrawBasketFeesPda
} from '../pda';

const WITHDRAW_FEES_DISCRIMINATOR = Buffer.from([159, 222, 146, 66, 254, 108, 55, 51]);

const CLAIM_FEE_TOKENS_FROM_BASKET_DISCRIMINATOR = Buffer.from([53, 109, 220, 141, 193, 193, 13, 32]);

export function withdrawFeesIx(params: {
    claimer: PublicKey,
    basketTokenMint: PublicKey,
    feeType: number,
}): TransactionInstruction {
    let { claimer, basketTokenMint, feeType } = params;
    let basket = getBasketState(basketTokenMint);
    let withdrawBasketFees = getWithdrawBasketFeesPda(basket, feeType);
    let basketFeesWallet = getBasketFeesPda(basket);
    let basketFeesTokenAccount = getAta(basketFeesWallet, basketTokenMint);
    let globalConfig = getGlobalConfigPda();
    
    const discriminator = WITHDRAW_FEES_DISCRIMINATOR;
    const args = Buffer.from([feeType]);
    const data = Buffer.concat([discriminator, args]);

    const keys = [
        { pubkey: claimer, isWritable: true, isSigner: true },
        { pubkey: basket, isWritable: true, isSigner: false },
        { pubkey: withdrawBasketFees, isWritable: true, isSigner: false },
        { pubkey: basketTokenMint, isWritable: true, isSigner: false },
        { pubkey: basketFeesWallet, isWritable: true, isSigner: false },
        { pubkey: basketFeesTokenAccount, isWritable: true, isSigner: false },
        { pubkey: globalConfig, isWritable: false, isSigner: false },
        { pubkey: SystemProgram.programId, isWritable: false, isSigner: false },
        { pubkey: TOKEN_PROGRAM_ID, isWritable: false, isSigner: false },
        { pubkey: ASSOCIATED_TOKEN_PROGRAM_ID, isWritable: false, isSigner: false },
    ];

    return new TransactionInstruction({
        keys,
        programId: BASKETS_V3_PROGRAM_ID,
        data,
    });
}

export function claimFeeTokensFromBasketIx(params: {
    signer: PublicKey,
    basket: PublicKey,
    withdrawBasketFees: PublicKey,
    rentPayer: PublicKey,
    owners: PublicKey[],
    tokenMints: PublicKey[],
}): TransactionInstruction {
    let { signer, basket, withdrawBasketFees, rentPayer, owners, tokenMints } = params;
    let numOwners = owners.length;
    
    const discriminator = CLAIM_FEE_TOKENS_FROM_BASKET_DISCRIMINATOR;
    const args = Buffer.from([numOwners]);
    const data = Buffer.concat([discriminator, args]);

    const keys = [
        { pubkey: signer, isWritable: true, isSigner: true },
        { pubkey: basket, isWritable: true, isSigner: false },
        { pubkey: withdrawBasketFees, isWritable: true, isSigner: false },
        { pubkey: rentPayer, isWritable: true, isSigner: false },
        { pubkey: SystemProgram.programId, isWritable: false, isSigner: false },
        { pubkey: TOKEN_PROGRAM_ID, isWritable: false, isSigner: false },
        { pubkey: TOKEN_2022_PROGRAM_ID, isWritable: false, isSigner: false },
        { pubkey: ASSOCIATED_TOKEN_PROGRAM_ID, isWritable: false, isSigner: false },
    ];

    // Remaining accounts: for each token: (mint, basketAta, claimer1Ata, claimer2Ata, ...)
    for (let i = 0; i < tokenMints.length; i++) {
        const mint = tokenMints[i];
        const basketAta = getAta(basket, mint);
        const claimerATAs = owners.map(owner => getAta(owner, mint));
        keys.push({ pubkey: mint, isWritable: false, isSigner: false });
        keys.push({ pubkey: basketAta, isWritable: true, isSigner: false });
        
        for (let j = 0; j < numOwners; j++) {
            keys.push({ pubkey: claimerATAs[j], isWritable: true, isSigner: false });
        }
    }

    return new TransactionInstruction({
        keys,
        programId: BASKETS_V3_PROGRAM_ID,
        data,
    });
}

export function claimFeeTokensFromBasketIxs(
    signer: PublicKey,
    basket: PublicKey,
    withdrawBasketFees: PublicKey,
    rentPayer: PublicKey,
    owners: PublicKey[],
    tokenMints: PublicKey[],
): TransactionInstruction[] {
    if (owners.length === 0 || tokenMints.length === 0) return [];
    
    // Account limit: 8 fixed accounts + (owners + 2) * number_of_mints <= 30
    // So: (owners + 2) * number_of_mints <= 24
    // Max tokens per tx: Math.floor(24 / (owners.length + 2))
    const accountsPerToken = owners.length + 2; // mint, basketAta, and one account per owner
    const maxRemainingAccounts = 24; // 30 - 8 fixed accounts
    const maxTokensPerTx = Math.floor(maxRemainingAccounts / accountsPerToken);
    
    let ixs: TransactionInstruction[] = [];

    for (let i = 0; i < tokenMints.length; i += maxTokensPerTx) {
        const batchTokenMints = tokenMints.slice(i, i + maxTokensPerTx);
        ixs.push(claimFeeTokensFromBasketIx({
            signer: signer,
            basket: basket,
            withdrawBasketFees: withdrawBasketFees,
            rentPayer: rentPayer,
            owners: owners,
            tokenMints: batchTokenMints,
        }));
    }
    
    return ixs;
}