import { AddressLookupTableAccount, ComputeBudgetProgram, Connection, Keypair, PublicKey, TransactionInstruction, TransactionMessage, TransactionSignature, VersionedTransaction } from "@solana/web3.js";
import { COMPUTE_UNITS } from "../constants";

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

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

export async function getAddressLookupTableAccounts(
    connection: Connection,
    keys: PublicKey[]
): Promise<AddressLookupTableAccount[]> {
    const addressLookupTableAccountInfos =
      await connection.getMultipleAccountsInfo(
        keys,
        "confirmed"
      );
  
    return addressLookupTableAccountInfos.reduce((acc, accountInfo, index) => {
      const addressLookupTableAddress = keys[index];
      if (accountInfo) {
        const addressLookupTableAccount = new AddressLookupTableAccount({
          key: new PublicKey(addressLookupTableAddress),
          state: AddressLookupTableAccount.deserialize(accountInfo.data),
        });
        acc.push(addressLookupTableAccount);
      }
  
      return acc;
    }, new Array<AddressLookupTableAccount>());
};

export async function getMultipleAddressLookupTableAccounts(
    connection: Connection,
    keys: PublicKey[][],
): Promise<AddressLookupTableAccount[][]> {
    let allLuts: string[] = [];
    keys.forEach(luts => 
        luts.forEach(lut => allLuts.push(lut.toBase58()))
    )
    allLuts = [...new Set(allLuts)];
    const addressLookupTableAccounts: AddressLookupTableAccount[] = await getAddressLookupTableAccounts(
        connection,
        allLuts.map(lut => new PublicKey(lut))
    );
    const map: {[key: string]: AddressLookupTableAccount} = {};
    allLuts.forEach((pubkey, id) => map[pubkey] = addressLookupTableAccounts[id]);
    return keys.map(luts => luts.map(lut => map[lut.toBase58()]));
}

export function wrapV0Transaction(
    blockhash: string,
    addressLookupTableAccounts: AddressLookupTableAccount[],
    payerPubkey: PublicKey,
    priorityFee: number,
    ixs: TransactionInstruction[],
): VersionedTransaction {
    ixs = [
        ...ixs,
        ComputeBudgetProgram.setComputeUnitLimit({units: COMPUTE_UNITS}),
        ComputeBudgetProgram.setComputeUnitPrice({microLamports: priorityFee}),
    ];
    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 sendV0Transaction(
    connection: Connection,
    tx: VersionedTransaction,
    blockhash: string,
    lastValidBlockHeight: number,
    simulateTransactions: boolean,
): Promise<TransactionSignature> {
    const serializedTx = tx.serialize();
    let txId: TransactionSignature = "Error";
    if (simulateTransactions) {
        txId = await connection.sendRawTransaction(
            serializedTx,
            { preflightCommitment: "confirmed"}
        ).catch(e => {console.log(e.message); return "Error"});
        console.log("Simulation txId:", txId);
        for (let i = 0; i < 4; i++) {
            await delay(1000).then(() =>
                connection.sendRawTransaction(
                    serializedTx,
                    {preflightCommitment: "confirmed"}
                ).catch(() => {})
            );
        }
        if (txId === "Error")
            throw new Error("Simulation failed");
    } else {
        connection.sendRawTransaction(serializedTx, {preflightCommitment: "confirmed"})
            .catch(e => {console.log(e.message)});

        txId = await connection.sendRawTransaction(
            serializedTx,
            { skipPreflight: true}
        );
        console.log("Sending tx:", txId);

        for (let i = 0; i < 4; i++) {
            await delay(1000).then(() =>
                connection.sendRawTransaction(serializedTx, {skipPreflight: true}).catch(() => {})
            );
        }
    }

    let confirmation = null;
    let result = null;

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

    let iterations = 20; 
    while (confirmation === null && result === null && iterations > 0) {
        await delay(2500);
        result = await connection.getTransaction(txId, {
            commitment: "confirmed",
            maxSupportedTransactionVersion: 0,
        });
        if (result && result.meta && 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 prepareV0Transactions(params: {
    connection: Connection,
    payer: PublicKey,
    priorityFee: number,
    multipleIxs: TransactionInstruction[][],
    multipleLookupTableAddresses: PublicKey[][],
    signers: Keypair[][],
    batches: number[],
},): Promise<VersionedTxs> {
    const { connection, payer, priorityFee, multipleIxs, multipleLookupTableAddresses, signers, batches } = params;
    const { blockhash, lastValidBlockHeight } = await connection.getLatestBlockhash();
    const multipleAddressLookupTableAccounts: AddressLookupTableAccount[][] = await getMultipleAddressLookupTableAccounts(
        connection,
        multipleLookupTableAddresses
    );
    const txs = multipleIxs.map((ixs, index) => {
        let tx = null;
        try {
            tx = wrapV0Transaction(
                blockhash,
                multipleAddressLookupTableAccounts[index],
                payer,
                priorityFee,
                ixs,
            );
            if (signers[index].length > 0)
                tx.sign(signers[index]);
        } catch (e: any) {
            console.log("Error signing tx:", e.message);
        }
        return tx;
    }).filter(tx => tx !== null);
    return {
        blockhash,
        lastValidBlockHeight,
        versionedTxs: txs,
        batches,
    };
}

export async function sendV0Transactions(
    connection: Connection,
    txParams: VersionedTxs,
    simulateTransactions: boolean,
): Promise<TransactionSignature[]> {
    const { versionedTxs, blockhash, lastValidBlockHeight, batches } = txParams;
    const signedTxs = versionedTxs;
    let lastIndex = 0;
    let txIds: TransactionSignature[] = [];
    for (let i = 0; i < batches.length; i++) {
        const ids = await Promise.all(
            signedTxs.slice(lastIndex, lastIndex + batches[i]).map(
                signedTx => sendV0Transaction(
                    connection,
                    signedTx,
                    blockhash,
                    lastValidBlockHeight,
                    simulateTransactions
                ).catch(e =>  {
                    console.log("Transaction failed:", e.message);
                    return "Error"
                })
            )
        );
        txIds = [...txIds, ...ids];
        lastIndex += batches[i];
    }
    return txIds;
}