import { BN, Instruction } from "@project-serum/anchor";
import { Provider, Wallet } from "@project-serum/anchor";
import {
  transfer,
  WRAPPED_SOL_MINT,
} from "@project-serum/serum/lib/token-instructions";
import {
  AccountInfo,
  ASSOCIATED_TOKEN_PROGRAM_ID,
  MintInfo,
  Token,
  TOKEN_PROGRAM_ID,
  u64,
} from "@solana/spl-token";
import {
  Connection,
  PublicKey,
  Signer,
  SystemProgram,
  Transaction,
  TransactionInstruction,
} from "@solana/web3.js";
import type { Action, BNIsh } from "../interfaces";
import { parseMintAccount, parseTokenAccount } from "@project-serum/common";
import { conditionallyIncludeSigner } from "./solana";
import type { SendTxRequest } from "@project-serum/anchor/dist/provider";

export const tryCreateAssociatedAccountInst = async (
  mint: PublicKey,
  owner: PublicKey,
  connection: Connection,
  payer: PublicKey
): Promise<{
  associateTokAccount: PublicKey;
  instr?: TransactionInstruction;
}> => {
  const associated = (
    await PublicKey.findProgramAddress(
      [
        owner.toBuffer(),
        TOKEN_PROGRAM_ID.toBuffer(),
        new PublicKey(mint).toBuffer(),
      ],
      ASSOCIATED_TOKEN_PROGRAM_ID
    )
  )[0];
  const data = await connection.getAccountInfo(associated);
  if (!data) {
    const inst = Token.createAssociatedTokenAccountInstruction(
      ASSOCIATED_TOKEN_PROGRAM_ID,
      TOKEN_PROGRAM_ID,
      new PublicKey(mint),
      associated,
      owner,
      payer
    );
    return {
      instr: inst,
      associateTokAccount: associated,
    };
  } else {
    return {
      associateTokAccount: associated,
    };
  }
};

export const tryCreateAssociatedAccount = async (
  mint: PublicKey,
  owner: PublicKey,
  provider: Provider
) => {
  const { associateTokAccount, instr } = await tryCreateAssociatedAccountInst(
    mint,
    owner,
    provider.connection,
    provider.wallet.publicKey
  );
  if (!instr) return associateTokAccount;
  const tx = new Transaction();
  tx.add(instr);
  await provider.send(tx);
  return associateTokAccount;
};

export interface AccountByMint {
  [mint: string]: { publicKey: PublicKey };
}

/**
 * Transfer to the token account being used by a construction input.
 * If the source account does not have enough balance, then throw an error.
 * This will also return null if the source equals the destination.
 * Transfer checked will also check if the mint is Wrapped Sol. If it is
 * and the source account is not a token account, the returned tx will transfer Native Solana
 * from the source
 */
export const transferToTxChecked = async (
  source: PublicKey,
  destination: PublicKey,
  amount: BN,
  authority: Signer | Wallet,
  mint: PublicKey,
  connection: Connection
): Promise<SendTxRequest | null> => {
  // Skip the transfer if the source or mint is a "dummy" account
  if (source.equals(PublicKey.default) || mint.equals(PublicKey.default))
    return null;

  const signer = conditionallyIncludeSigner([], authority);
  const isTokAccount = await isTokenAccount(source, mint, connection);

  // Check if the mint is a wrapped sol mint and if the source is not
  // a token account
  if (mint.equals(WRAPPED_SOL_MINT) && !isTokAccount) {
    const nativeBal = await connection.getBalance(source);
    if (new BN(nativeBal.toString()).lt(amount))
      throw `Expected account ${source.toBase58()} to have at least ${amount.toString()} lamports`;
    const insts = getWrapSolInsts(source, destination, amount);
    const tx = new Transaction();
    tx.add(...insts);
    return { tx, signers: signer };
  }

  if (!isTokAccount)
    throw `Expected input token account, ${source.toBase58()}, to be a token account`;

  const bal = await connection.getTokenAccountBalance(source);
  if (new BN(bal.value.amount).lt(amount)) {
    throw `Expected token account for mint ${mint.toBase58()} to have at least ${toReadableNumber(
      amount,
      bal.value.decimals
    )}`;
  }

  if (source.equals(destination)) return null;

  const inst = transfer({
    source,
    destination,
    amount,
    owner: authority.publicKey,
  });
  const tx = new Transaction();
  tx.add(inst);
  return { tx, signers: signer };
};

export type PreferredTokenAccounts = { [mint: string]: PublicKey };

/**
 * Get the token account which the malloc sdk will use for a mint
 *
 * @param preferredTokenAccounts - a map of mints to token accounts which the caller would prefer to use per mint. If the account does not exist,
 * and is needed it will be created in the the returned txs
 */
export const getTokenAccountUsedByMalloc = async (
  tokenAuthority: PublicKey | string,
  mintAccount: PublicKey | string,
  preferredTokenAccounts?: PreferredTokenAccounts
): Promise<PublicKey> => {
  const mintAccountPK = new PublicKey(mintAccount);
  const tokenAuthPk = new PublicKey(tokenAuthority);
  if (
    preferredTokenAccounts &&
    preferredTokenAccounts[mintAccountPK.toBase58()]
  ) {
    return preferredTokenAccounts[mintAccountPK.toBase58()];
  }
  return await findAssociatedTokenAddress(tokenAuthPk, mintAccountPK);
};

/**
 * Takes in a mapping of mints to token accounts. If a token account has yet to exist, one will be made
 * which is associated to the authority's address
 *
 * @param preferredTokenAccounts - a map of mints to token accounts which the caller would prefer to use per mint. If the account does not exist,
 * and is needed it will be created in the the returned txs
 */
export const createTokenAccountsForActions = async (
  provider: Provider,
  tokenAuthority: Signer | Wallet,
  actions: Action[],
  fundingAccount: Signer | Wallet,
  preferredTokenAccounts?: PreferredTokenAccounts
): Promise<{ accounts: AccountByMint; txs: SendTxRequest[] }> => {
  const allMints = actions
    .map((a) => {
      const mintsWithDummies: PublicKey[] = [];
      if (a.opts?.tokenCreationOptions?.skipOutTokenCreations || false) {
        const skipIdxs = a.opts?.tokenCreationOptions.skipOutTokenCreations;
        // Skip all the mints which are included by the skipIdxs
        a.tokenMintOuts.forEach((mint, i) => {
          if (!skipIdxs.includes(i)) mintsWithDummies.push(mint);
        });
      } else {
        mintsWithDummies.push(...a.tokenMintOuts);
      }

      if (a.opts?.tokenCreationOptions?.skipInTokenCreation === true) {
      } else {
        mintsWithDummies.push(...a.tokenMintsIn);
      }
      if (a.opts?.tokenCreationOptions?.extraTokenMintsUsed) {
        mintsWithDummies.push(
          ...a.opts.tokenCreationOptions.extraTokenMintsUsed
        );
      }

      // Remove tall the dummy mints, signified by the default pk
      const mints = mintsWithDummies.filter(
        (m) => !m.equals(PublicKey.default)
      );
      return mints;
    })
    .flat();

  const mints = removeDuplicateMints(allMints);

  let accounts: AccountByMint = {};
  for (let i = 0; i < mints.length; i++) {
    const mint = mints[i];
    if ((preferredTokenAccounts || {})[mint])
      accounts[mint] = { publicKey: preferredTokenAccounts[mint] };
    // create the token account if not provided
    else {
      const pk = await findAssociatedTokenAddress(
        tokenAuthority.publicKey,
        mint
      );
      accounts[mint] = {
        publicKey: pk,
      };
    }
  }

  const createTokTx = await Promise.all(
    mints.map(async (mint) => {
      const associatedToSignerAccount = accounts[mint].publicKey;
      const accountExists = await provider.connection.getAccountInfo(
        associatedToSignerAccount
      );

      // Return if the account already exists
      if (accountExists) {
        return null;
      }
      const insts = Token.createAssociatedTokenAccountInstruction(
        ASSOCIATED_TOKEN_PROGRAM_ID,
        TOKEN_PROGRAM_ID,
        new PublicKey(mint),
        accounts[mint].publicKey,
        tokenAuthority.publicKey,
        fundingAccount.publicKey
      );
      const tx = new Transaction();
      tx.add(insts);
      const signers = conditionallyIncludeSigner([], fundingAccount);
      return {
        tx,
        signers: signers,
      };
    })
  );

  return { accounts: accounts, txs: createTokTx.filter((i) => !!i) };
};

// TODO: if there are 2 token accounts of the same type, one of them not being made, then there will be an error out

// See https://spl.solana.com/associated-token-account
const findAssociatedTokenAddress = async (
  walletAddress: PublicKey,
  tokenMintAddress: PublicKey | string
): Promise<PublicKey> => {
  return (
    await PublicKey.findProgramAddress(
      [
        walletAddress.toBuffer(),
        TOKEN_PROGRAM_ID.toBuffer(),
        new PublicKey(tokenMintAddress).toBuffer(),
      ],
      ASSOCIATED_TOKEN_PROGRAM_ID
    )
  )[0];
};

export const removeDuplicateMints = (mints: PublicKey[]) => [
  ...new Set(mints.map(pkToStr)),
];

export const removeMints = (
  mints: PublicKey[],
  toRemove: PublicKey[]
): PublicKey[] =>
  mints.filter((m) => !toRemove.find((toRem) => toRem.equals(m)));

const pkToStr = (pk: PublicKey) => pk.toBase58();
const pkFromStr = (pk: string) => new PublicKey(pk);

export const toReadableNumber = (number: u64, decimals: number): string => {
  function trimTrailingZeroes(value: string): string {
    return value.replace(/\.?0*$/, "");
  }
  function formatWithCommas(value: string): string {
    const pattern = /(-?\d+)(\d{3})/;
    while (pattern.test(value)) {
      value = value.replace(pattern, "$1,$2");
    }
    return value;
  }

  const balance = number.toString();
  const wholeStr = balance.substring(0, balance.length - decimals) || "0";
  const fractionStr = balance
    .substring(balance.length - decimals)
    .padStart(decimals, "0")
    .substring(0, decimals);

  return trimTrailingZeroes(`${formatWithCommas(wholeStr)}.${fractionStr}`);
};

export const fromReadableNumber = (
  number: number,
  decimals: number
): string => {
  function trimLeadingZeroes(value: string): string {
    value = value.replace(/^0+/, "");
    if (value === "") {
      return "0";
    }
    return value;
  }

  const split = number.toString().split(".");
  const wholePart = split[0];
  const fracPart = split[1] || "";
  if (split.length > 2 || fracPart.length > decimals) {
    throw new Error(`Cannot parse '${number}' as token amount`);
  }
  return trimLeadingZeroes(wholePart + fracPart.padEnd(decimals, "0"));
};

/**
 * Split the amounts and add the remainder to the last amount
 */
export const splitTokensByFractions = (amount: BN, splits: number[]): BN[] => {
  if (splits.length === 0) return [];
  const sum = splits.reduce((a, b) => a + b, 0);
  const amounts = splits.map((split) => amount.muln(split).divn(sum));
  const amountSum = amounts.reduce((a, b) => a.add(b), new BN(0));
  const unused = amount.sub(amountSum);
  amounts[amounts.length - 1] = amounts[amounts.length - 1].add(unused);
  return amounts;
};

export const getMintInfo = async (
  connection: Connection,
  mintAddr: PublicKey
): Promise<MintInfo | null> => {
  const account = await connection.getAccountInfo(mintAddr);
  if (!account) return null;
  return parseMintAccount(account.data);
};

export const getTokenAccountInfo = async (
  connection: Connection,
  tokenAccount: PublicKey
): Promise<AccountInfo | null> => {
  const account = await connection.getAccountInfo(tokenAccount);
  if (!account) return null;
  return parseTokenAccount(account.data);
};

export const checkInMintsAreTheSame = (actions: Action[]) => {
  if (actions.length === 0) return true;
  const initMints = actions[0].tokenMintsIn.map((m) => m.toBase58());
  return initMints.every((mint, i) =>
    actions.every((a) => {
      if (a.tokenMintsIn.length <= i) {
        return false;
      }
      return a.tokenMintsIn[i].toBase58() === mint;
    })
  );
};

/**
 * Get the transaction to wrap sol
 *
 * @param tokenAccount - the token account which contains the wrapped sol.
 * Assume that the account exists
 */
const getWrapSolInsts = (
  payer: PublicKey,
  tokenAccount: PublicKey,
  lamports: BNIsh
): TransactionInstruction[] => {
  const transferNonWrappedInst = SystemProgram.transfer({
    programId: TOKEN_PROGRAM_ID,
    fromPubkey: payer,
    toPubkey: tokenAccount,
    lamports: new BN(lamports).toNumber(),
  });
  const syncInst = (Token as any).createSyncNativeInstruction(
    TOKEN_PROGRAM_ID,
    tokenAccount
  );
  return [transferNonWrappedInst, syncInst];
};

const isTokenAccount = async (
  account: PublicKey,
  expectedMint: PublicKey,
  connection: Connection
): Promise<boolean> => {
  try {
    const info = await getTokenAccountInfo(connection, account);
    return info.mint.equals(expectedMint);
  } catch (e) {
    console.error(
      "Failed to account as token account, assuming that it is not"
    );
    return false;
  }
};