import { BN, Program, Provider, Wallet } from "@project-serum/anchor";
import { PublicKey, Signer } from "@solana/web3.js";
import {
  AccountsArray,
  Action,
  ActionPIDToActionUID,
  BuildActionMap,
  Construction,
  errors,
  PartialNextActionCtx,
  TxRequestsTyped,
} from "../index";
import { ConstructionCompiled, ConstructionInstr } from "../interfaces";
import {
  checkInMintsAreTheSame,
  createTokenAccountsForActions,
  transferToTxChecked,
} from "../utils/tokens";
import { addActionsToConstruction, getProcessActionsTx } from "./actions";
import { createAndInitConstruction, prepareConstruction } from "./construciton";
import { TOKEN_PROGRAM_ID } from "@solana/spl-token";
import { cleanup } from "./cleanup";
import { SendTxRequest } from "@project-serum/anchor/dist/provider";

export const compileConstruction = async (
  provider: Provider,
  mallocProgram: Program,
  { construction, amountInAccounts, actionMap, opts }: ConstructionInstr
): Promise<ConstructionCompiled> => {
  const { actions, initialSplits, initialActionIndices, amountIns } =
    construction;

  const _inAuth = opts?.amountInAuthority;
  let amountInAuthority = _inAuth ?? (provider.wallet as Wallet);
  const _authority = opts?.authority;
  let authority = _authority || (provider.wallet as Wallet);

  const { accounts: tokenAccounts, txs: tokenTxs } =
    await createTokenAccountsForActions(
      provider,
      authority,
      actions,
      provider.wallet as Wallet,
      opts?.preferredTokenAccounts
    );

  let constructionTx: SendTxRequest | undefined;
  let constructionSigner: Signer;

  if (!opts?.constructionAccount) {
    const ret = await createAndInitConstruction(
      mallocProgram,
      authority.publicKey,
      actions.length
    );
    constructionTx = ret.tx;
    constructionSigner = ret.construction;
  } else {
    constructionSigner = opts.constructionAccount;
  }

  const addConstructionActionsTx = addActionsToConstruction(
    mallocProgram,
    constructionSigner.publicKey,
    actions,
    authority
  );

  if (!checkInMintsAreTheSame(initialActionIndices.map((i) => actions[i]))) {
    throw errors.newCompilerError(
      "Expected all initial token types to be the same"
    );
  }

  if (initialActionIndices.length === 0)
    throw errors.newCompilerError(
      "Expected initial actions to have length more than 0"
    );

  // TODO: check all token mints for init actions are the same
  const initMints = actions[initialActionIndices[0]].tokenMintsIn;
  if (initMints.length !== amountIns.length) {
    throw errors.newCompilerError(
      "Expected number of initial amounts to equal number of initial mints"
    );
  }

  if (amountInAccounts.length !== amountIns.length) {
    throw errors.newCompilerError(
      "Expected the number of initial accounts to equal the number of amounts in"
    );
  }

  // TODO: have this transfer transaction be conditional,
  // See https://github.com/Lev-Stambler/malloc-solana2/issues/4
  const transferTxs = (
    await Promise.all(
      amountIns.map((amount, i) =>
        initMints[i].equals(PublicKey.default)
          ? null
          : transferToTxChecked(
              amountInAccounts[i],
              tokenAccounts[initMints[i].toBase58()].publicKey,
              new BN(amount),
              amountInAuthority,
              initMints[i],
              provider.connection
            )
      )
    )
  ).filter((tx) => tx !== null);

  const prepareConstructionTx = prepareConstruction(
    mallocProgram,
    constructionSigner.publicKey,
    initialActionIndices,
    initialSplits,
    authority,
    amountIns
  );

  // The additional accounts added onto action action accounts follow conventional malloc standards for token
  const actionCtxFilled: PartialNextActionCtx[] = actions.map((action, i) => {
    const toAdd = action.opts?.addDefaultTokAccounts
      ? [
          {
            address: tokenAccounts[action.tokenMintsIn[0].toBase58()].publicKey,
            isWriteable: true,
            isSigner: false,
          },
          {
            address: TOKEN_PROGRAM_ID,
            isWriteable: false,
            isSigner: false,
          },
          {
            address: authority.publicKey,
            isSigner: true,
            isWriteable: false,
          },
        ]
      : [];

    return {
      actionProgram: action.actionProgram,
      actionAccounts: [...toAdd, ...action.actionAccounts] as AccountsArray,
    };
  });

  const processActionsTxs = construction.actionOrders.map((actionsInTx) =>
    getProcessActionsTx(
      mallocProgram,
      actionsInTx,
      constructionSigner.publicKey,
      actionsInTx.map((idx) => actionCtxFilled[idx]),
      authority,
      actionsInTx.map((idx) => {
        return {
          data: actions[idx].actionData,
          instructions: actions[idx].opts?.instructions,
          additionalSigners: actions[idx].opts?.additionalSigners,
          preActionTxs: actions[idx].opts?.preActionTxs,
        };
      })
    )
  );

  const actionPIDToUID = getActionPIDToUID(actions, actionMap);

  const atStartOfPrepareTxs = getAllAtStartOfPrepareTxs(actions);

  const cleanUpTxs = opts?.cleanupAtEnd
    ? await cleanup(
        provider,
        mallocProgram,
        construction,
        constructionSigner,
        authority,
        opts?.cleanupOpts
      )
    : [];

  const txs: TxRequestsTyped[] = [
    { type: "prepended", txs: atStartOfPrepareTxs },
    { type: "token", txs: tokenTxs },
    { type: "transfer", txs: transferTxs },
    { type: "construction", txs: constructionTx ? [constructionTx] : [] },
    { type: "add construction actions", txs: addConstructionActionsTx },
    { type: "prepare construction", txs: prepareConstructionTx },
    {
      type: "action calls",
      txs: processActionsTxs,
    },
    {
      type: "cleanup",
      txs: cleanUpTxs,
    },
  ];

  return {
    constructionSigner,
    construction,
    txs,
    tokenAccounts,
    authority,
    actionMap,
    actionPIDToUID,
  };
};

const getActionPIDToUID = (
  actions: Action[],
  actionMap: BuildActionMap
): ActionPIDToActionUID => {
  const ret: ActionPIDToActionUID = {};
  actions.forEach((action) => {
    ret[action.actionProgram.toBase58()] = action.actionTypeUID;
  });
  return ret;
};

const getAllAtStartOfPrepareTxs = (actions: Action[]) =>
  actions.map((a) => a.opts?.atStartOfPrepareTxs || []).flat();