import { AccountRole, Address, Instruction } from '@solana/kit';
import { SYSTEM_PROGRAM_ADDRESS } from '@solana-program/system';
import { RouterInstructions } from './types';
import { compileRouteIxs } from './compileRouteIxs';

/**
 * With @solana/kit invoking a program and having the account marked as writable results in an error being thrown:
 *
 * > SolanaError: This transaction includes an address (`11111111111111111111111111111111`) which is both invoked and marked writable. Program addresses may not be writable
 *
 * This function always marks system program accounts as read-only as well since it's a surefire way to increase tx costs unnecessarily.
 *
 * There is also a bug in older anchor libs where optional accounts (marked as `none` by sending the programAddress) are marked as writable unnecessarily. This function also works around that issue.
 *
 * @param ixs
 * @returns Instructions with invoked program accounts marked as read-only
 */
export function markInvokedProgramAccountsReadonly(ixs: RouterInstructions): RouterInstructions {
  const invokedProgramIds = new Set([...compileRouteIxs(ixs).map((ix) => ix.programAddress), SYSTEM_PROGRAM_ADDRESS]);

  return {
    createInAtaIxs: marketInvokedProgramAccountsReadonlyIxArray(invokedProgramIds, ixs.createInAtaIxs),
    createOutAtaIxs: marketInvokedProgramAccountsReadonlyIxArray(invokedProgramIds, ixs.createOutAtaIxs),
    wrapSolIxs: marketInvokedProgramAccountsReadonlyIxArray(invokedProgramIds, ixs.wrapSolIxs),
    swapIxs: marketInvokedProgramAccountsReadonlyIxArray(invokedProgramIds, ixs.swapIxs),
    unwrapSolIxs: marketInvokedProgramAccountsReadonlyIxArray(invokedProgramIds, ixs.unwrapSolIxs),
  };
}

function marketInvokedProgramAccountsReadonlyIxArray(
  invokedProgramIds: Set<Address>,
  ixs: Instruction[],
): Instruction[] {
  return ixs.map((ix) => markInvokedProgramAccountsReadonlyIx(invokedProgramIds, ix));
}

function markInvokedProgramAccountsReadonlyIx(invokedProgramIds: Set<Address>, ix: Instruction): Instruction {
  if (!ix.accounts) {
    return ix;
  }
  return {
    ...ix,
    accounts: ix.accounts.map((acc) => ({
      ...acc,
      role: invokedProgramIds.has(acc.address) ? AccountRole.READONLY : acc.role,
    })),
  };
}