import { Address, createNoopSigner, Instruction } from '@solana/kit';
import { RouteParams } from './types';
import { createAtaIdempotent } from '../utils';
import { ASSOCIATED_TOKEN_PROGRAM_ADDRESS } from '@solana-program/token-2022';
import { ReadonlyUint8Array } from '@solana/codecs-core/dist/types/readonly-uint8array';
import { WRAPPED_SOL_MINT } from '@kamino-finance/limo-sdk';

export async function extractOrAddAtaIxs(
  params: RouteParams,
  ixs: Instruction[],
  tokenInProgramId: Address,
  tokenOutProgramId: Address,
): Promise<{
  createInAtaIxs: Instruction[];
  createOutAtaIxs: Instruction[];
  otherIxs: Instruction[];
}> {
  const {
    createInAtaIxs: maybeCreateInAtas,
    createOutAtaIxs: maybeCreateOutAtas,
    otherIxs,
  } = splitCreateAtaIxs(params, ixs);
  if (params.includeSetupIxs === false) {
    return {
      createInAtaIxs: [],
      createOutAtaIxs: [],
      otherIxs,
    };
  }

  const getCreateInIxs = async (params: RouteParams, maybeCreateInAtas: Instruction[]) => {
    // We only need to create ATA for SOL input when wrapping
    // Otherwise, we assume the user already has the ATA for the input token
    if (params.tokenIn !== WRAPPED_SOL_MINT || params.wrapAndUnwrapSol === false) {
      return [];
    }
    if (maybeCreateInAtas.length > 0) {
      return maybeCreateInAtas;
    }
    const { createAtaIx } = await createAtaIdempotent(
      params.executor,
      createNoopSigner(params.executor),
      params.tokenIn,
      tokenInProgramId,
    );
    return [createAtaIx];
  };

  const getCreateOutIxs = async (params: RouteParams, maybeCreateOutAtas: Instruction[]) => {
    if (params.destinationTokenAccount !== undefined) {
      // If the user specifies a destination token account, we assume they have created it
      return [];
    } else if (maybeCreateOutAtas.length > 0) {
      return maybeCreateOutAtas;
    }
    const { createAtaIx } = await createAtaIdempotent(
      params.executor,
      createNoopSigner(params.executor),
      params.tokenOut,
      tokenOutProgramId,
    );
    return [createAtaIx];
  };

  const [createInAtaIxs, createOutAtaIxs] = await Promise.all([
    getCreateInIxs(params, maybeCreateInAtas),
    getCreateOutIxs(params, maybeCreateOutAtas),
  ]);

  return {
    createInAtaIxs,
    createOutAtaIxs,
    otherIxs,
  };
}

function splitCreateAtaIxs(
  params: RouteParams,
  ixs: Instruction[],
): {
  createInAtaIxs: Instruction[];
  createOutAtaIxs: Instruction[];
  otherIxs: Instruction[];
} {
  const createInAtaIxs: Instruction[] = [];
  const createOutAtaIxs: Instruction[] = [];
  const otherIxs: Instruction[] = [];
  for (const ix of ixs) {
    if (ix.programAddress !== ASSOCIATED_TOKEN_PROGRAM_ADDRESS) {
      otherIxs.push(ix);
      continue;
    }
    if (isCreateOrCreateIdempotentAtaIxDiscrim(ix.data)) {
      if (!ix.accounts || ix.accounts.length < 4) {
        otherIxs.push(ix);
        continue;
      }
      const owner = ix.accounts[2].address;
      const mint = ix.accounts[3].address;
      if (owner === params.executor && mint === params.tokenIn) {
        if (createInAtaIxs.length > 0) {
          // remove duplicate
          continue;
        }
        if (isCreateAtaIxDiscrim(ix.data)) {
          createInAtaIxs.push({
            ...ix,
            data: Uint8Array.of(1), // overwrite to idempotent ix
          });
        } else {
          createInAtaIxs.push(ix);
        }
        continue;
      } else if (owner === params.executor && mint === params.tokenOut) {
        if (createOutAtaIxs.length > 0) {
          // remove duplicate
          continue;
        }
        if (isCreateAtaIxDiscrim(ix.data)) {
          createOutAtaIxs.push({
            ...ix,
            data: Uint8Array.of(1), // overwrite to idempotent ix
          });
        } else {
          createOutAtaIxs.push(ix);
        }
        continue;
      }
    }
    otherIxs.push(ix);
  }
  return { createInAtaIxs, createOutAtaIxs, otherIxs };
}

function isCreateOrCreateIdempotentAtaIxDiscrim(data: ReadonlyUint8Array | undefined): boolean {
  return (data !== undefined && data[0] === 1) || isCreateAtaIxDiscrim(data);
}

/**
 * Create ATA ix (non-idempotent) can either have no discriminator or 0
 * Reference:
 * https://github.com/solana-program/associated-token-account/blob/main/program/src/processor.rs#L45
 * @param data
 */
function isCreateAtaIxDiscrim(data: ReadonlyUint8Array | undefined): boolean {
  return !data || data.length === 0 || data[0] === 0;
}