import { RouteParams } from './types';
import { Address, createNoopSigner, Instruction, TransactionSigner } from '@solana/kit';
import { getAssociatedTokenAddress } from '../utils';
import BN from 'bn.js';
import { WRAPPED_SOL_MINT } from '@kamino-finance/limo-sdk';
import { TOKEN_PROGRAM_ADDRESS } from '@solana-program/token';
import {
  CLOSE_ACCOUNT_DISCRIMINATOR,
  getCloseAccountInstruction,
  getSyncNativeInstruction,
  SYNC_NATIVE_DISCRIMINATOR,
} from '@solana-program/token-2022';
import {
  getTransferSolDiscriminatorBytes,
  getTransferSolInstruction,
  SYSTEM_PROGRAM_ADDRESS,
} from '@solana-program/system';
import { ReadonlyUint8Array } from '@solana/codecs-core/dist/types/readonly-uint8array';

// System program discriminators are 4 bytes
const TRANSFER_SOL_DISCRIMINATOR_BYTES = Buffer.of(...getTransferSolDiscriminatorBytes());

export async function wrapAndUnwrapSolIxs(params: RouteParams): Promise<{
  wrapSolIxs: Instruction[];
  unwrapSolIxs: Instruction[];
}> {
  let wrapIxs: Instruction[] = [];
  let unwrapIxs: Instruction[] = [];
  if (params.wrapAndUnwrapSol && params.tokenIn === WRAPPED_SOL_MINT) {
    [wrapIxs, unwrapIxs] = await Promise.all([wrapSolIxs(params, undefined), unwrapSolIxs(params, undefined)]);
  }

  if (params.wrapAndUnwrapSol && params.tokenOut === WRAPPED_SOL_MINT) {
    unwrapIxs = await unwrapSolIxs(params, params.destinationTokenAccount);
  }
  return { wrapSolIxs: wrapIxs, unwrapSolIxs: unwrapIxs };
}

export async function unwrapSolIxs(params: RouteParams, wsolTa: Address | undefined): Promise<Instruction[]> {
  const wsolAta =
    wsolTa !== undefined
      ? wsolTa
      : await getAssociatedTokenAddress(params.executor, WRAPPED_SOL_MINT, TOKEN_PROGRAM_ADDRESS);
  return [
    getCloseAccountInstruction(
      { account: wsolAta, destination: params.executor, owner: params.executor },
      { programAddress: TOKEN_PROGRAM_ADDRESS },
    ),
  ];
}

export async function wrapSolIxs(params: RouteParams, wsolTa: Address | undefined): Promise<Instruction[]> {
  const ata: Address =
    wsolTa !== undefined
      ? wsolTa
      : await getAssociatedTokenAddress(params.executor, WRAPPED_SOL_MINT, TOKEN_PROGRAM_ADDRESS);
  return getDepositWsolIxs(createNoopSigner(params.executor), ata, params.amount);
}

export async function extractOrAddWrapAndUnwrapSolIxs(
  params: RouteParams,
  ixs: Instruction[],
): Promise<{
  wrapSolIxs: Instruction[];
  unwrapSolIxs: Instruction[];
  otherIxs: Instruction[];
}> {
  const {
    wrapSolIxs: maybeWrapSolIxs,
    unwrapSolIxs: maybeUnwrapSolIxs,
    otherIxs,
  } = await splitWrapAndUnwrapSolIxs(params, ixs);

  // Remove unwanted wrap/unwrap SOL ixs if the user disabled it
  if (params.wrapAndUnwrapSol === false) {
    return {
      wrapSolIxs: [],
      unwrapSolIxs: [],
      otherIxs,
    };
  }

  const getWrapSolIxs = async (params: RouteParams, maybeWrapSolIxs: Instruction[]) => {
    if (maybeWrapSolIxs.length > 0) {
      return maybeWrapSolIxs;
    }
    if (params.tokenIn !== WRAPPED_SOL_MINT) {
      return [];
    }
    return wrapSolIxs(params, undefined);
  };

  const getUnwrapSolIxs = async (params: RouteParams, maybeUnwrapSolIxs: Instruction[]) => {
    if (maybeUnwrapSolIxs.length > 0) {
      return maybeUnwrapSolIxs;
    }
    if (params.tokenIn !== WRAPPED_SOL_MINT && params.tokenOut !== WRAPPED_SOL_MINT) {
      return [];
    }
    return unwrapSolIxs(params, params.tokenOut === WRAPPED_SOL_MINT ? params.destinationTokenAccount : undefined);
  };

  const [wrapIxs, unwrapIxs] = await Promise.all([
    getWrapSolIxs(params, maybeWrapSolIxs),
    getUnwrapSolIxs(params, maybeUnwrapSolIxs),
  ]);

  return {
    wrapSolIxs: wrapIxs,
    unwrapSolIxs: unwrapIxs,
    otherIxs,
  };
}

/**
 * Get the deposit WSOL instructions
 * @param owner - the owner address
 * @param ata - the ata address
 * @param amountLamports - the amount in lamports
 * @returns the transaction instructions
 */
function getDepositWsolIxs(owner: TransactionSigner, ata: Address, amountLamports: BN): Instruction[] {
  const ixns: Instruction[] = [];

  // Transfer to WSOL ata
  ixns.push(
    getTransferSolInstruction({
      source: owner,
      destination: ata,
      amount: BigInt(amountLamports.toString()),
    }),
  );

  // Sync wrapped SOL
  ixns.push(
    getSyncNativeInstruction(
      { account: ata },
      {
        programAddress: TOKEN_PROGRAM_ADDRESS,
      },
    ),
  );

  return ixns;
}

async function splitWrapAndUnwrapSolIxs(
  params: RouteParams,
  ixs: Instruction[],
): Promise<{
  wrapSolIxs: Instruction[];
  unwrapSolIxs: Instruction[];
  otherIxs: Instruction[];
}> {
  if (params.tokenIn === WRAPPED_SOL_MINT) {
    return splitWrapAndUnwrapSolTokenInIxs(params, ixs);
  }
  if (params.tokenOut === WRAPPED_SOL_MINT) {
    return splitUnwrapSolTokenOutIxs(params, ixs);
  }
  return { wrapSolIxs: [], unwrapSolIxs: [], otherIxs: ixs };
}

async function splitWrapAndUnwrapSolTokenInIxs(
  params: RouteParams,
  ixs: Instruction[],
): Promise<{
  wrapSolIxs: Instruction[];
  unwrapSolIxs: Instruction[];
  otherIxs: Instruction[];
}> {
  const wrapSolIxs: Instruction[] = [];
  const unwrapSolIxs: Instruction[] = [];
  const otherIxs: Instruction[] = [];

  const solAta = await getAssociatedTokenAddress(params.executor, WRAPPED_SOL_MINT, TOKEN_PROGRAM_ADDRESS);

  for (let i = 0; i < ixs.length; i++) {
    const ix = ixs[i];

    if (ix.programAddress !== TOKEN_PROGRAM_ADDRESS) {
      otherIxs.push(ix);
      continue;
    }
    // wrap SOL
    if (isSyncNativeDiscrim(ix.data)) {
      // Look for a transfer and sync native pair
      if (!ix.accounts || ix.accounts.length === 0) {
        otherIxs.push(ix);
        continue;
      }
      const solTokenAccount = ix.accounts[0].address;
      if (solTokenAccount !== solAta) {
        otherIxs.push(ix);
        continue;
      }

      // Look for a transfer ix before the sync native ix
      const prevIx = ixs[i - 1];
      if (prevIx && prevIx.programAddress === SYSTEM_PROGRAM_ADDRESS && isSolTransferDiscrim(prevIx.data)) {
        if (!prevIx.accounts || prevIx.accounts.length < 2) {
          otherIxs.push(ix);
          continue;
        }
        const source = prevIx.accounts[0].address;
        const destination = prevIx.accounts[1].address;
        if (source === params.executor && destination === solTokenAccount) {
          if (wrapSolIxs.length > 0) {
            // remove duplicate
            otherIxs.pop();
            continue;
          }
          wrapSolIxs.push(prevIx);
          wrapSolIxs.push(ix);
          otherIxs.pop(); // remove prevIx from otherIxs
          continue;
        }
      }
    }

    // unwrap SOL
    if (isCloseTokenTokenAccountIx(ix, params.executor, solAta)) {
      if (unwrapSolIxs.length > 0) {
        // remove duplicate
        continue;
      }
      unwrapSolIxs.push(ix);
      continue;
    }

    otherIxs.push(ix);
  }
  return { wrapSolIxs, unwrapSolIxs, otherIxs };
}

async function splitUnwrapSolTokenOutIxs(
  params: RouteParams,
  ixs: Instruction[],
): Promise<{
  wrapSolIxs: Instruction[];
  unwrapSolIxs: Instruction[];
  otherIxs: Instruction[];
}> {
  const unwrapSolIxs: Instruction[] = [];
  const otherIxs: Instruction[] = [];

  const outSolAta =
    params.destinationTokenAccount ||
    (await getAssociatedTokenAddress(params.executor, WRAPPED_SOL_MINT, TOKEN_PROGRAM_ADDRESS));

  for (let i = 0; i < ixs.length; i++) {
    const ix = ixs[i];

    if (ix.programAddress !== TOKEN_PROGRAM_ADDRESS) {
      otherIxs.push(ix);
      continue;
    }
    if (isCloseTokenTokenAccountIx(ix, params.executor, outSolAta)) {
      if (unwrapSolIxs.length > 0) {
        // remove duplicate
        continue;
      }
      unwrapSolIxs.push(ix);
      continue;
    }

    otherIxs.push(ix);
  }
  return { wrapSolIxs: [], unwrapSolIxs, otherIxs };
}

function isCloseTokenTokenAccountIx(ix: Instruction, executor: Address, account: Address): boolean {
  if (ix.programAddress !== TOKEN_PROGRAM_ADDRESS) {
    return false;
  }
  if (!ix.accounts || ix.accounts.length < 2) {
    return false;
  }
  if (!isCloseAccountDiscrim(ix.data)) {
    return false;
  }

  const tokenAccount = ix.accounts[0].address;
  const destination = ix.accounts[1].address;
  return tokenAccount === account && destination === executor;
}

function isSyncNativeDiscrim(data: ReadonlyUint8Array | undefined): boolean {
  return data !== undefined && data[0] === SYNC_NATIVE_DISCRIMINATOR;
}

function isCloseAccountDiscrim(data: ReadonlyUint8Array | undefined): boolean {
  return data !== undefined && data[0] === CLOSE_ACCOUNT_DISCRIMINATOR;
}

function isSolTransferDiscrim(data: ReadonlyUint8Array | undefined): boolean {
  // System program discriminators are 4 bytes
  return data !== undefined && Buffer.of(...data.subarray(0, 4)).equals(TRANSFER_SOL_DISCRIMINATOR_BYTES);
}