import { BN, Program } from "@coral-xyz/anchor";
import { ComputeBudgetProgram, Connection, PublicKey, Transaction, TransactionInstruction, TransactionSignature } from "@solana/web3.js";
import { Wallet } from "@coral-xyz/anchor/dist/cjs/provider";
import { FundsIDL } from "./fundsIDL";
import { TOKEN_PROGRAM_ID, createAssociatedTokenAccountInstruction, getAssociatedTokenAddressSync } from "@solana/spl-token";
import { applyBlockHashAndPartialSign, sendAndConfirmTransaction, prepareAccounts } from "./txHelpers";
import { RouteData, SWAP_FEE_ACCOUNT, TokenSettings } from "./types";

export async function generateSwapInstruction(
    program: Program<FundsIDL>,
    user: PublicKey,
    tokenListData: TokenSettings[],
    routeData: RouteData,
    fromTokenAccount: PublicKey,
    toTokenAccount: PublicKey,
    slippage: number = 0.5,
): Promise<TransactionInstruction> {
    let { fromTokenId, toTokenId, fromAmount, toAmount } = routeData;
    let minimumReceived = toAmount * (1 - slippage / 100);
    return await program.methods
        .liquidityProvision(
            new BN(fromTokenId),
            new BN(toTokenId),
            new BN(
                Math.floor(
                    fromAmount * 10 ** tokenListData[fromTokenId].decimals
                )
            ),
            new BN(
                Math.floor(
                    minimumReceived * 10 ** tokenListData[toTokenId].decimals
                )
            ),
        )
        .accounts({
            buyer: user,
            fundState: routeData.swapAccounts.fundState,
            pdaAccount: routeData.swapAccounts.authority,
            pdaFromTokenAccount: routeData.swapAccounts.source,
            buyerFromTokenAccount: fromTokenAccount,
            pdaToTokenAccount: routeData.swapAccounts.destination,
            buyerToTokenAccount: toTokenAccount,
            swapFeeAccount: getAssociatedTokenAddressSync(
                routeData.swapAccounts.fees.feeTokenMint,
                routeData.swapAccounts.fees.smfWallet
            ),
            hostFeeAccount: getAssociatedTokenAddressSync(
                routeData.swapAccounts.fees.feeTokenMint,
                routeData.swapAccounts.fees.hostWallet
            ),
            managerFeeAccount: getAssociatedTokenAddressSync(
                routeData.swapAccounts.fees.feeTokenMint,
                routeData.swapAccounts.fees.managerWallet
            ),
            tokenList: routeData.swapAccounts.tokenList,
            curveData: routeData.swapAccounts.curveData,
            tokenProgram: TOKEN_PROGRAM_ID
        })
        .remainingAccounts(routeData.swapAccounts.remainingAccounts)
        .instruction();
}

export async function executeSwap(
    program: Program<FundsIDL>,
    wallet: Wallet,
    connection: Connection,
    tokenListData: TokenSettings[],
    routeData: RouteData,
    userFromTokenAccount: PublicKey|undefined,
    userToTokenAccount: PublicKey|undefined,
    slippage: number = 0.5,
): Promise<any> {
    let {preTransaction, postTransaction, fromTokenAccount, toTokenAccount} = await prepareAccounts(wallet, connection, tokenListData, routeData, userFromTokenAccount, userToTokenAccount);
    let mainTransaction = new Transaction();

    let instructions = [];
    let swapFeeAccount = getAssociatedTokenAddressSync(
        routeData.swapAccounts.fees.feeTokenMint,
        SWAP_FEE_ACCOUNT,
    );
    const info2 =  await program.provider.connection.getAccountInfo(swapFeeAccount, "confirmed");
    if (!info2) {
        instructions.push(
            createAssociatedTokenAccountInstruction(
                wallet.publicKey,
                swapFeeAccount,
                SWAP_FEE_ACCOUNT,
                routeData.swapAccounts.fees.feeTokenMint,
            )
        );
    }
    let hostFeeAccount = getAssociatedTokenAddressSync(
        routeData.swapAccounts.fees.feeTokenMint,
        routeData.swapAccounts.fees.hostWallet,
    );
    const info3 =  await program.provider.connection.getAccountInfo(hostFeeAccount, "confirmed");
    if (!info3 && routeData.swapAccounts.fees.hostWallet.toBase58() !== SWAP_FEE_ACCOUNT.toBase58()
        && routeData.swapAccounts.fees.hostWallet.toBase58() !== wallet.publicKey.toBase58()) {
        instructions.push(
            createAssociatedTokenAccountInstruction(
                wallet.publicKey,
                hostFeeAccount,
                routeData.swapAccounts.fees.hostWallet,
                routeData.swapAccounts.fees.feeTokenMint,
            )
        );
    }
    let managerFeeAccount = getAssociatedTokenAddressSync(
        routeData.swapAccounts.fees.feeTokenMint,
        routeData.swapAccounts.fees.managerWallet,
    );
    const info4 =  await program.provider.connection.getAccountInfo(managerFeeAccount, "confirmed");
    if (!info4 && routeData.swapAccounts.fees.managerWallet.toBase58() !== SWAP_FEE_ACCOUNT.toBase58() 
            && routeData.swapAccounts.fees.managerWallet.toBase58() !== routeData.swapAccounts.fees.hostWallet.toBase58()
                && routeData.swapAccounts.fees.managerWallet.toBase58() !== wallet.publicKey.toBase58()) {
        instructions.push(
            createAssociatedTokenAccountInstruction(
                wallet.publicKey,
                managerFeeAccount,
                routeData.swapAccounts.fees.managerWallet,
                routeData.swapAccounts.fees.feeTokenMint,
            )
        );
    }
    if (instructions.length > 0) {
        preTransaction.add(...instructions);
    }
    
    mainTransaction.add(await generateSwapInstruction(
        program,
        wallet.publicKey,
        tokenListData,
        routeData,
        fromTokenAccount,
        toTokenAccount,
        slippage
    ));
    if(preTransaction.instructions.length > 0) {
        preTransaction.add(ComputeBudgetProgram.setComputeUnitLimit({units: 1000000}));
        preTransaction.add(ComputeBudgetProgram.setComputeUnitPrice({microLamports: 20000}));
    }
    if(postTransaction.instructions.length > 0) {
        postTransaction.add(ComputeBudgetProgram.setComputeUnitLimit({units: 1000000}));
        postTransaction.add(ComputeBudgetProgram.setComputeUnitPrice({microLamports: 20000}));
    }
    mainTransaction.add(ComputeBudgetProgram.setComputeUnitLimit({units: 1000000}));
    mainTransaction.add(ComputeBudgetProgram.setComputeUnitPrice({microLamports: 20000}));
    const { txIndex, transactions } = await applyBlockHashAndPartialSign(
        connection, wallet,
        preTransaction, mainTransaction, postTransaction,
        [], [],
    );
    const signedTxs = await wallet.signAllTransactions(transactions);

    let serialized: any = [];
    let err: any = {};
    let txId;
    for (let i = 0; i < signedTxs.length; i++){
        serialized.push(signedTxs[i].serialize());
    
        for (let i = 0; i < serialized.length; i++) {
            const { signature, response } = await sendAndConfirmTransaction(connection, serialized[i]);

            let failed = false;
            if (!signature || !response || response == "Error" || response.meta.err)
                failed = true;
            if(failed && err.message == undefined){
                if (i != txIndex)
                    err.message = "Couldn't prepare accounts";
                else
                    err.message = "Swap failed";
                err.txId = signature;
            }
            if(i == txIndex) txId = signature;
        }
    }

    if(err.message) throw new FundError(err.message, err.txId);
    return txId;
}

export class FundError extends Error {
    public txId: TransactionSignature;
    constructor (message: string, txId: TransactionSignature = "") {
        super(message);
        this.txId = txId;
    }
}
