import type {
  BlockhashWithExpiryBlockHeight,
  Connection,
  PublicKey,
  TransactionInstruction,
} from "@solana/web3.js";
import {
  ComputeBudgetProgram,
  TransactionMessage,
  VersionedTransaction,
} from "@solana/web3.js";

import { estimatePriorityFees } from "./estimatePriorityFees.js";
import type { TransactionEnvelope } from "./transactionEnvelope.js";

export const MAX_PRIORITY_FEE = 1_000_000;
export const MIN_PRIORITY_FEE = 20_000;
export const PRIORITY_FEE_MULTIPLIER = 1.5; // 50% above the average

export const getCUsForTx = async (
  connection: Connection,
  latestBlockhash: Awaited<ReturnType<typeof connection.getLatestBlockhash>>,
  txs: TransactionInstruction[],
  payerKey: PublicKey,
) => {
  const messageV0 = new TransactionMessage({
    payerKey,
    recentBlockhash: latestBlockhash.blockhash,
    instructions: txs,
  }).compileToV0Message();
  const transaction = new VersionedTransaction(messageV0);
  const simulation = await connection.simulateTransaction(transaction);
  const CUs =
    !simulation.value.unitsConsumed || simulation.value.unitsConsumed === 0
      ? 1.4e6
      : simulation.value.unitsConsumed;
  console.log("Estimated CUs:", CUs);
  return CUs;
};

/**
 * Arguments for calling `createVersionedTransaction`.
 */
export interface CreateVersionedTransactionArgs {
  connection: Connection;
  tx: TransactionEnvelope;
  payerKey: PublicKey;
  addCUs?: boolean;
  minimumCU?: number;
  doNotAddPriorityFee?: boolean;
}

export interface VersionedTransactionWithBlockhash {
  transaction: VersionedTransaction;
  /**
   * Pass this to `confirmTransaction` to confirm the transaction.
   */
  latestBlockhash: BlockhashWithExpiryBlockHeight;
}

export const createVersionedTransaction = async ({
  connection,
  tx,
  payerKey,
  addCUs,
  minimumCU,
  doNotAddPriorityFee,
}: CreateVersionedTransactionArgs): Promise<VersionedTransactionWithBlockhash> => {
  const latestBlockhash = await connection.getLatestBlockhash("confirmed");
  const ixs = [...tx.instructions];

  if (!doNotAddPriorityFee) {
    const priorityFeeIx = await estimatePriorityFees(
      connection,
      ixs,
      PRIORITY_FEE_MULTIPLIER,
      MAX_PRIORITY_FEE,
      MIN_PRIORITY_FEE,
    );
    ixs.unshift(priorityFeeIx);
  }

  if (addCUs) {
    const estimatedCUs = await getCUsForTx(
      connection,
      latestBlockhash,
      ixs,
      payerKey,
    );
    const CUs = Math.max(minimumCU ?? 28000, estimatedCUs); // Always have at least 28k or minimumCU CU
    ixs.unshift(
      ComputeBudgetProgram.setComputeUnitLimit({
        units: CUs + 3000, // +3k for safety and the CU limit ix itself
      }),
    );
  }

  const messageV0 = new TransactionMessage({
    payerKey,
    recentBlockhash: latestBlockhash.blockhash,
    instructions: ixs,
  }).compileToV0Message();
  const transaction = new VersionedTransaction(messageV0);
  // Add additional signers if provided
  if (tx.signers && tx.signers.length > 0) {
    transaction.sign(tx.signers);
  }
  return { transaction, latestBlockhash };
};