// Core dependencies
import { Program } from "@coral-xyz/anchor";
import { Connection, Keypair, PublicKey, SystemProgram, TransactionInstruction } from "@solana/web3.js";

// Local imports
import { BasketsProgram } from "./idl/types";
import { getAta } from "./utils/programAccounts";
import { prepareV0Transactions, VersionedTxs } from "./utils/txUtils";
import { updateTokenPricesIxs } from "./instructions/update/updateTokenPrices";
import { depositIx } from "./instructions/buy/deposit";
import { withdrawIx } from "./instructions/sell/withdraw";
import { initializeWithdrawStateIx } from "./instructions/sell/initializeWithdrawState";
import { generateSwapInstruction, getQuoteResponseHandler } from "./instructions/jup";
import { fetchWithdrawState, WithdrawState } from "./state/withdrawState";
import { sellDepositAfterRebalanceIx } from "./instructions/sell/depositAfterRebalance";
import { sellWithdrawBeforeRebalanceIx } from "./instructions/sell/withdrawBeforeRebalance";
import { claimTokensIxs } from "./instructions/sell/claimTokens";
import { WSOL_MINT, MAX_CLAIM_TOKENS_PER_TX } from "./utils/constants";
import { createAssociatedTokenAccountInstruction, createSyncNativeInstruction } from "@solana/spl-token";
import { createAtasIxs } from "./utils/createAtas";

export async function buyBasketHandler(
    sdkParams: {
        program: Program<BasketsProgram>,
        payer: PublicKey,
        connection: Connection,
        priorityFee: number,
    },
    params: {
        basket: PublicKey,
        depositAmount: number,
        depositMint: PublicKey,
    }
): Promise<VersionedTxs> {
    const preIxs = [];

    if (params.depositMint.equals(WSOL_MINT)) {
        const ata = getAta(sdkParams.payer, WSOL_MINT);
        const info = await sdkParams.connection.getAccountInfo(ata);
        let toWrap = params.depositAmount;
        if (info)
            toWrap -= parseInt(info.data?.readBigUInt64LE(64).toString() ?? "0");
        else
            preIxs.push(
                createAssociatedTokenAccountInstruction(
                    sdkParams.payer,
                    ata,
                    sdkParams.payer,
                    WSOL_MINT,
                )
            );
        if (toWrap > 0) {
            preIxs.push(
                SystemProgram.transfer({
                    fromPubkey: sdkParams.payer,
                    toPubkey: ata,
                    lamports: toWrap,
                })
            );
            preIxs.push(
                createSyncNativeInstruction(ata)
            );
        }
    }

    const {
        ixs: updatePricesIxs,
        luts: updatePricesLuts,
    } = await updateTokenPricesIxs({
        program: sdkParams.program,
        basket: params.basket,
    })
    const buyIx = await depositIx({
        program: sdkParams.program,
        buyer: sdkParams.payer,
        ...params,
    })
    const txs: TransactionInstruction[][] = updatePricesIxs.map(ix => [ix]);
    const luts: PublicKey[][] = new Array(updatePricesIxs.length).fill(updatePricesLuts);
    const signers: Keypair[][] = new Array(updatePricesIxs.length).fill([]);
    if (preIxs.length > 0) {
        txs.push(preIxs);
        luts.push([]);
        signers.push([]);
    }
    txs.push([buyIx]);
    luts.push([]);
    signers.push([]);
    return await prepareV0Transactions({
        connection: sdkParams.connection,
        payer: sdkParams.payer,
        priorityFee: sdkParams.priorityFee,
        multipleIxs: txs,
        multipleLookupTableAddresses: luts,
        signers: signers,
        batches: [txs.length - 1, 1],
    });
}

export async function sellBasketHandler(
    sdkParams: {
        payer: PublicKey,
        connection: Connection,
        program: Program<BasketsProgram>,
        priorityFee: number,
        jupiterApiKey: string,
        maxAllowedAccounts: number,
    },
    params: {
        basket: PublicKey;
        withdrawStateSeed: number[];
        amountToWithdraw: number,
        destinationMint: PublicKey,
        rebalance: boolean,
    }
): Promise<VersionedTxs> {
    const {
        ixs: updatePricesIxs,
        luts: updatePricesLuts,
    } = await updateTokenPricesIxs({
        program: sdkParams.program,
        basket: params.basket,
    })
    const initIx = await initializeWithdrawStateIx({
        program: sdkParams.program,
        withdrawStateSeed: params.withdrawStateSeed,
    })
    const sellIx = await withdrawIx({
        program: sdkParams.program,
        seller: sdkParams.payer,
        ...params,
    })

    return await prepareV0Transactions({
        connection: sdkParams.connection,
        payer: sdkParams.payer,
        priorityFee: sdkParams.priorityFee,
        multipleIxs: [...updatePricesIxs.map(ix => [ix]), [initIx, sellIx]],
        multipleLookupTableAddresses: [...new Array(updatePricesIxs.length).fill(updatePricesLuts), []],
        signers: [...new Array(updatePricesIxs.length).fill([]), []],
        batches: [updatePricesIxs.length, 1],
    });
}

export async function sellRebalanceTokensIxs(
    sdkParams: {
        payer: PublicKey,
        connection: Connection,
        program: Program<BasketsProgram>,
        priorityFee: number,
        jupiterApiKey: string,
        maxAllowedAccounts: number,
    },
    params: {
        withdrawStateData: WithdrawState;
        fromToken: PublicKey;
    }
): Promise<{
    ixs: TransactionInstruction[];
    luts: PublicKey[];
}> {
    const { withdrawStateData, fromToken } = params;

    const fromTokenIndex = withdrawStateData.compositionMints.findIndex(mint => mint.equals(fromToken));
    const fromAmount = parseInt(withdrawStateData.compositionAmounts[fromTokenIndex].toString());

    const quoteResponse = await getQuoteResponseHandler({
        jupiterApiKey: sdkParams.jupiterApiKey,
        maxAllowedAccounts: sdkParams.maxAllowedAccounts,
        fromToken: params.fromToken,
        toToken: withdrawStateData.destinationMint,
        amount: fromAmount,
        slippageBps: 100,
    });

    const withdrawIx = await sellWithdrawBeforeRebalanceIx({
        program: sdkParams.program,
        withdrawStateData: withdrawStateData,
        payer: sdkParams.payer,
        fromTokenMint: params.fromToken,
    });

    const jupIxAndLuts = await generateSwapInstruction({
        payer: sdkParams.payer,
        jupiterApiKey: sdkParams.jupiterApiKey,
        quoteResponse: quoteResponse,
    }).catch(() => null);
    if (!jupIxAndLuts) {
        return {
            ixs: [],
            luts: [],
        };
    }

    const depositIx = await sellDepositAfterRebalanceIx({
        program: sdkParams.program,
        withdrawStateData: withdrawStateData,
        payer: sdkParams.payer,
        fromTokenMint: params.fromToken,
    });

    return {
        ixs: [withdrawIx, jupIxAndLuts.ix, depositIx],
        luts: jupIxAndLuts.luts,
    };
}

export async function sellRebalanceHandler(
    sdkParams: {
        payer: PublicKey,
        connection: Connection,
        program: Program<BasketsProgram>,
        priorityFee: number,
        jupiterApiKey: string,
        maxAllowedAccounts: number,
    },
    params: {
        withdrawState: PublicKey,
    }
): Promise<VersionedTxs> {
    const withdrawStateData = await fetchWithdrawState(sdkParams.program, params.withdrawState);
    const ixs: TransactionInstruction[][] = [];
    const luts: PublicKey[][] = [];
    const tokenMints: PublicKey[] = [];
    for (let i = 0; i < withdrawStateData.compositionMints.length; i++) {
        const fromToken = withdrawStateData.compositionMints[i];
        const fromAmount = parseInt(withdrawStateData.compositionAmounts[i].toString());
        if (fromAmount > 0) {
            const { ixs: tempIxs, luts: tempLuts } = await sellRebalanceTokensIxs(sdkParams, {
                withdrawStateData,
                fromToken,
            });
            if (tempIxs.length > 0) {
                ixs.push(tempIxs);
                luts.push(tempLuts);
                tokenMints.push(withdrawStateData.destinationMint, fromToken);
            }
        }
    }
    if (ixs.length === 0)
        throw new Error("No rebalance transactions found");
    const preIxs = await createAtasIxs(sdkParams.connection, {
        payer: sdkParams.payer,
        mints: tokenMints,
    });
    return await prepareV0Transactions({
        connection: sdkParams.connection,
        payer: sdkParams.payer,
        priorityFee: sdkParams.priorityFee,
        multipleIxs: [...preIxs, ...ixs],
        multipleLookupTableAddresses: [...new Array(preIxs.length).fill([]), ...luts],
        signers: [...new Array(preIxs.length).fill([]), ...new Array(ixs.length).fill([])],
        batches: [preIxs.length, ixs.length],
    });
}

export async function claimTokensHandler(
    sdkParams: {
        program: Program<BasketsProgram>,
        payer: PublicKey,
        connection: Connection,
        priorityFee: number,
    },
    params: {
        withdrawState: PublicKey;
    }
): Promise<VersionedTxs> {
    const ixs = await claimTokensIxs({
        program: sdkParams.program,
        payer: sdkParams.payer,
        ...params,
    });
    const bundledIxs: TransactionInstruction[][] = [];
    for (let i = 0; i < ixs.length; i += MAX_CLAIM_TOKENS_PER_TX)
        bundledIxs.push(ixs.slice(i, i + MAX_CLAIM_TOKENS_PER_TX));
    return await prepareV0Transactions({
        connection: sdkParams.connection,
        payer: sdkParams.payer,
        priorityFee: sdkParams.priorityFee,
        multipleIxs: bundledIxs,
        multipleLookupTableAddresses: new Array(bundledIxs.length).fill([]),
        signers: new Array(bundledIxs.length).fill([]),
        batches: [bundledIxs.length],
    });
}