import { AccountInfo, SystemProgram, Transaction } from '@solana/web3.js';
import {
    AddressLookupTableAccount, ComputeBudgetProgram, Connection, Keypair, PublicKey, Signer,
    TransactionInstruction, TransactionMessage, VersionedTransaction
} from '@solana/web3.js';
import { getAta } from './instructions/pda';
import { MINTS } from './constants';
import { createAssociatedTokenAccountIdempotentInstruction, createSyncNativeInstruction } from '@solana/spl-token';

type TransactionSignature = string;

export interface Wallet {
    signTransaction<T extends Transaction | VersionedTransaction>(tx: T): Promise<T>;
    signAllTransactions<T extends Transaction | VersionedTransaction>(txs: T[]): Promise<T[]>;
    publicKey: PublicKey;
    /** Keypair of the configured payer (Node only) */
    payer?: Keypair;
}

export async function wrapWsolIxs(
    connection: Connection,
    wallet: PublicKey,
    amount: number,
): Promise<TransactionInstruction[]> {
    if (amount == 0) return [];
    let wsolIxs: TransactionInstruction[] = [];
    let wsolAta = getAta(wallet, MINTS["mainnet"].WSOL);
    let wsolBalance = 0;
    let accountExists = false;
    try {
        let info = await connection.getTokenAccountBalance(wsolAta);
        wsolBalance = parseInt(info.value.amount);
        accountExists = wsolBalance >= 0;
    } catch {}
    if (!accountExists)
        wsolIxs.push(
            createAssociatedTokenAccountIdempotentInstruction(
                wallet,
                wsolAta,
                wallet,
                MINTS["mainnet"].WSOL,
            )
        );
    if (wsolBalance < amount)
        wsolIxs.push(
            SystemProgram.transfer({
                fromPubkey: wallet,
                toPubkey: wsolAta,
                lamports: amount - wsolBalance,
            }),
            createSyncNativeInstruction(wsolAta),
        );
    return wsolIxs;
}
export interface TxData {
    payer: PublicKey,
    instructions: TransactionInstruction[],
    lookupTables: PublicKey[],
}

export interface TxBatchData {
    batches: TxData[][],
}

export interface VersionedTxs {
    blockhash: string;
    lastValidBlockHeight: number;
    batches: VersionedTransaction[][];
}

export interface TxPayloadIxAccountMeta {
    pubkey: string,
    is_signer: boolean,
    is_writable: boolean,
}
export interface TxPayloadIx {
    program_id: string,
    accounts: TxPayloadIxAccountMeta[],
    data: string,
}
export interface TxPayload {
    tx_b64: string,
    message_version: "0" | "legacy",
    recent_blockhash: string,
    payer: string,
    lookup_tables: string[],
    instructions: TxPayloadIx[],
}
export interface TxPayloadBatch {
    transactions: TxPayload[];
}
export interface TxPayloadBatchSequence {
    batches: TxPayloadBatch[];
}

export interface BasketCreationTx extends TxPayloadBatchSequence {
    mint: string;
    basket: string;
}

export function prepareTxPayloadBatchSequence(
    txBatchData: TxBatchData,
    versionedTxs?: VersionedTxs,
): TxPayloadBatchSequence {
    return {
        batches: txBatchData.batches.map((txDatas, i) => {
            return {
                transactions: txDatas.map((txData, j) => {
                    return {
                        tx_b64: versionedTxs ? Buffer.from(versionedTxs.batches[i][j].serialize()).toString('base64') : "",
                        message_version: versionedTxs ? (versionedTxs.batches[i][j].message.version == 0 ? "0" : "legacy") : "0",
                        recent_blockhash: versionedTxs ? versionedTxs.batches[i][j].message.recentBlockhash : "",
                        payer: txData.payer.toBase58(),
                        lookup_tables: txData.lookupTables.map(lookupTable => lookupTable.toBase58()),
                        instructions: txData.instructions.map(
                            instruction => ({
                                program_id: instruction.programId.toBase58(),
                                accounts: instruction.keys.map(account => ({
                                    pubkey: account.pubkey.toBase58(),
                                    is_signer: account.isSigner,
                                    is_writable: account.isWritable,
                                })),
                                data: instruction.data.toString('base64'),
                            })
                        )
                    }
                })
            }
        }),
    };
}

export async function delay(ms: number): Promise<void> {
    return new Promise((resolve) => setTimeout(resolve, ms));
}

export async function getMultipleAccountsInfoBatched(
    connection: Connection,
    pubkeys: PublicKey[],
): Promise<Map<string, AccountInfo<Buffer> | null>> {
    let uniquePubkeys: PublicKey[] = [];
    for (const pubkey of pubkeys) {
        if (uniquePubkeys.find(x => x.equals(pubkey)) === undefined) {
            uniquePubkeys.push(pubkey);
        }
    }
    let chunkSize = 100;
    let promises: Promise<(AccountInfo<Buffer> | null)[]>[] = [];
    for (let i = 0; i < uniquePubkeys.length; i += chunkSize) {
        let chunk = uniquePubkeys.slice(i, i + chunkSize);
        promises.push(connection.getMultipleAccountsInfo(chunk, "confirmed"));
    }
    let resolve = await Promise.all(promises);
    let infos: (AccountInfo<Buffer> | null)[] = [];
    resolve.forEach(result => {
        infos.push(...result);
    });
    let allResults: Map<string, AccountInfo<Buffer> | null> = new Map();
    infos.forEach((info, idx) => {
        allResults.set(uniquePubkeys[idx].toBase58(), info);
    });
    return allResults;
}

export async function getAddressLookupTableAccounts(
    connection: Connection,
    keys: PublicKey[]
): Promise<Map<string, AddressLookupTableAccount>> {
    const addressLookupTableAccountInfos = await getMultipleAccountsInfoBatched(connection, keys);
    let addressLookupTableAccounts: Map<string, AddressLookupTableAccount> = new Map();
    for (const key of keys) {
        const accountInfo = addressLookupTableAccountInfos.get(key.toBase58());
        if (!key.equals(PublicKey.default) && accountInfo) {
            const addressLookupTableAccount = new AddressLookupTableAccount({
                key: new PublicKey(key),
                state: AddressLookupTableAccount.deserialize(accountInfo.data),
            });
            addressLookupTableAccounts.set(key.toBase58(), addressLookupTableAccount);
        }
    }
    return addressLookupTableAccounts;
};

export async function getMultipleAddressLookupTableAccounts(
    connection: Connection,
    batches: PublicKey[][][],
): Promise<AddressLookupTableAccount[][][]> {
    let allLuts: PublicKey[] = [];
    batches.forEach(batch => batch.forEach(luts => luts.forEach(lut => allLuts.push(lut))));
    let addressLookupTableAccounts = await getAddressLookupTableAccounts(connection, allLuts);
    let result: AddressLookupTableAccount[][][] = [];
    for (const batch of batches) {
        let batchLuts: AddressLookupTableAccount[][] = [];
        for (const luts of batch) {
            let batchLut: AddressLookupTableAccount[] = [];
            for (const lut of luts) {
                batchLut.push(addressLookupTableAccounts.get(lut.toBase58())!);
            }
            batchLuts.push(batchLut);
        }
        result.push(batchLuts);
    }
    return result;
}

export function compileVersionedTransaction(
    blockhash: string,
    addressLookupTableAccounts: AddressLookupTableAccount[],
    payerPubkey: PublicKey,
    ixs: TransactionInstruction[],
): VersionedTransaction {
    const txMessage = new TransactionMessage({    
        payerKey: payerPubkey,
        recentBlockhash: blockhash,
        instructions: ixs,
    });
    let versionedTx = new VersionedTransaction(
        txMessage.compileToV0Message(addressLookupTableAccounts)
    );
    try {
        let tt = versionedTx.serialize().length;
        if (tt > 1232) {
            throw new Error("Transaction too large");
        }
    } catch (e: any) {
        throw new Error(e.message);
    }
    return versionedTx;
}

export async function sendVersionedTransaction(
    connection: Connection,
    tx: VersionedTransaction,
    blockhash: string,
    lastValidBlockHeight: number,
    simulateTransactions: boolean,
): Promise<TransactionSignature> {
    const serializedTx = tx.serialize();
    let txId: TransactionSignature;
    if (simulateTransactions) {
        txId = await connection.sendRawTransaction(
            serializedTx,
            { preflightCommitment: "confirmed", maxRetries: 0 }
        ).catch(e => {
            console.log(e.message);
            throw new Error("Simulation failed");
        });
        for (let i = 0; i < 4; i++) {
            await delay(3000).then(() =>
                connection.sendRawTransaction(
                    serializedTx,
                    {preflightCommitment: "confirmed", maxRetries: 0}
                ).catch(() => {})
            );
        }
        console.log("Simulation txId:", txId);
    } else {
        txId = await connection.sendRawTransaction(
            serializedTx,
            { skipPreflight: true, maxRetries: 0 }
        );
        for (let i = 0; i < 4; i++) {
            await delay(3000).then(() =>
                connection.sendRawTransaction(
                    serializedTx,
                    {preflightCommitment: "confirmed", maxRetries: 0}
                ).catch(() => {})
            );
        }
        console.log("Sending tx:", txId);
    }

    let confirmation = null;
    let result = null;

    connection.confirmTransaction({
        blockhash,
        lastValidBlockHeight,
        signature: txId,
    }, "confirmed").catch(() => null).then((res) => confirmation = res);

    let iterations = 10;
    while (confirmation === null && result === null && iterations > 0) {
        await delay(1000);
        result = await connection.getTransaction(txId, {
            commitment: "confirmed",
            maxSupportedTransactionVersion: 0,
        }).catch(() => null);
        if (result && result.meta && result.meta?.err) {
            console.log(result.meta.err);
            throw new Error(txId);
        }
        iterations--;
    }

    if (result) return txId;

    //@ts-ignore
    if (!confirmation || confirmation.value.err) {
        //@ts-ignore
        console.log(confirmation?.value.err);
        throw new Error(txId);
    }

    return txId;
}

export async function prepareVersionedTxs(
    connection: Connection,
    txBatchData: TxBatchData,
): Promise<VersionedTxs> {
    const { batches } = txBatchData;
    const { blockhash, lastValidBlockHeight } = await connection.getLatestBlockhash("confirmed");
    let multipleLookupTableAddresses: PublicKey[][][] = batches.map(batch => batch.map(txData => txData.lookupTables));
    const multipleAddressLookupTableAccounts: AddressLookupTableAccount[][][] = await getMultipleAddressLookupTableAccounts(
        connection,
        multipleLookupTableAddresses
    );
    const batchTxs = batches.map((batch, batchIndex) => {
        const txs = batch.map((txData, txIndex) => {
            let tx = null;
            try {
                tx = compileVersionedTransaction(
                    blockhash,
                    multipleAddressLookupTableAccounts[batchIndex][txIndex],
                    txData.payer,
                    txData.instructions,
                );
            } catch (e: any) {
                console.log("Error signing tx:", e.message);
            }
            return tx;
        }).filter(tx => tx !== null);
        return txs;
    });
    return {
        blockhash,
        lastValidBlockHeight,
        batches: batchTxs,
    };
}

export async function sendVersionedTxs(
    connection: Connection,
    versionedTxs: VersionedTxs,
    simulateTransactions: boolean = false,
): Promise<TransactionSignature[][]> {
    const { batches, blockhash, lastValidBlockHeight } = versionedTxs;
    let batchTxIds: TransactionSignature[][] = [];

    for (const batch of batches) {
        let txIds = await Promise.all(
            batch.map(tx =>
                sendVersionedTransaction(
                    connection,
                    tx,
                    blockhash,
                    lastValidBlockHeight,
                    simulateTransactions
                ).catch(e => {console.log("Error sending tx:", e.message); return "Error"})
            )
        );
        batchTxIds.push(txIds);
    }
    return batchTxIds;
}


export async function signVersionedTxs(
    wallet: Wallet,
    versionedTxs: VersionedTxs,
): Promise<VersionedTxs> {
    let txs: VersionedTransaction[] = [];
    for (const batch of versionedTxs.batches)
        for (const tx of batch)
            txs.push(tx);
    txs = await wallet.signAllTransactions(txs);
    let pointer = 0;
    for (let batchIndex = 0; batchIndex < versionedTxs.batches.length; batchIndex++)
        for (let txIndex = 0; txIndex < versionedTxs.batches[batchIndex].length; txIndex++) {
            versionedTxs.batches[batchIndex][txIndex] = txs[pointer];
            pointer++;
        }
    return versionedTxs;
}

export async function signTxPayloadBatchSequence(
    wallet: Wallet,
    txPayloadBatchSequence: TxPayloadBatchSequence,
): Promise<TxPayloadBatchSequence> {
    let txs: VersionedTransaction[] = [];
    for (const batch of txPayloadBatchSequence.batches)
        for (const tx of batch.transactions)
            txs.push(
                VersionedTransaction.deserialize(Buffer.from(tx.tx_b64, 'base64')) 
            );
    txs = await wallet.signAllTransactions(txs);
    let pointer = 0;
    for (let batchIndex = 0; batchIndex < txPayloadBatchSequence.batches.length; batchIndex++)
        for (let txIndex = 0; txIndex < txPayloadBatchSequence.batches[batchIndex].transactions.length; txIndex++) {
            txPayloadBatchSequence.batches[batchIndex].transactions[txIndex].tx_b64 = Buffer.from(txs[pointer].serialize()).toString('base64');
            pointer++;
        }
    return txPayloadBatchSequence;
}

export async function sendTxPayloadBatchSequence(
    connection: Connection,
    txPayloadBatchSequence: TxPayloadBatchSequence,
    simulateTransactions: boolean = false,
): Promise<TransactionSignature[][]> {
    let txs: VersionedTransaction[][] = [];
    for (const batch of txPayloadBatchSequence.batches) {
        let batchTxs: VersionedTransaction[] = [];
        for (const tx of batch.transactions)
            batchTxs.push(VersionedTransaction.deserialize(Buffer.from(tx.tx_b64, 'base64')));
        txs.push(batchTxs);
    }
    return await sendVersionedTxs(connection, {
        blockhash: txPayloadBatchSequence.batches[0].transactions[0]?.recent_blockhash ?? "",
        lastValidBlockHeight: 0,
        batches: txs,
    }, simulateTransactions);
}