import { Provider } from "@project-serum/anchor";
import { PublicKey } from "@solana/web3.js";
import BN from "bn.js";
import { checkServerIdentity } from "tls";
import {
  Action,
  ActionVertexInfo,
  BNIsh,
  BuildActionMap,
  Construction,
  errors,
  IsolatedAction,
  NextNode,
  NonSpecificAtionOverride,
  NonSpecificConstruction,
  NonSpecificAction,
  MintWithAmount,
  MintAndTokenAccount,
  NextNodeNoNextMint,
  EstimateOutFN,
  ActionVertexInfoIsolated,
  BuildIsolateActionConfig,
  NonSpecificConstructionFilledUnserializable,
  _NonSpecificActionFilled,
  NonSpecificConstructionFilled,
  isNonSpecificActionGroup,
  castToActionGroup,
  castToStandaloneAction,
  IsolatedActionWithNonSpecific,
  MallocError,
  MallocActionBuilderError,
} from "../index";
import { newBuilderError } from "../error";
import {
  buildSeqListOfActionCalls,
  getNonSpecificActionDataFromId,
  prettyPrintActionId,
  updateNonSpecificActionDataFromId,
} from "../graph-utils";
import { deepCloneObject, removeDuplicates } from "../utils/object";
import {
  getTokenAccountUsedByMalloc,
  PreferredTokenAccounts,
  removeDuplicateMints,
  removeMints,
} from "../utils/tokens";
import { EstimateRet, getConstructionEstimates } from "./estimates";

/**
 * Get the set of previous actions for an action,
 */
const getPreviousActions = (
  actionIdx: number,
  actions: (IsolatedAction | undefined)[],
  nextNodes: NextNodeNoNextMint[][][]
): { idx: number; idxInNextNodes: [number, number][] }[] => {
  // First add indices to each action
  // Then remove the undefined actions
  // Then search if the action has the actionIdx
  return (
    actions
      .map((a, i) => {
        return a ? { ...a, idx: i } : a;
      })
      .filter((a) => !!a)
      .map((a: IsolatedAction & { idx: number }) => {
        const allIdxInNextNodes: [number, number][] = [];
        // note this assumes that each action is only called once from a prior action
        // see https://github.com/Lev-Stambler/malloc-solana2/issues/5
        for (let i = 0; i < nextNodes[a.idx].length; i++) {
          const idxInNextNode = [-1, -1];
          idxInNextNode[1] = nextNodes[a.idx][i]
            .map((item) => item.actionIdx)
            .indexOf(actionIdx);
          if (idxInNextNode[1] !== -1) {
            idxInNextNode[0] = i;
            allIdxInNextNodes.push(idxInNextNode as [number, number]);
          }
        }
        if (allIdxInNextNodes.length === 0) return undefined;
        return {
          idx: a.idx,
          idxInNextNodes: allIdxInNextNodes,
        };
      })
      // Ensure that the previous action has the output mint
      .filter(
        (i) =>
          i &&
          actions[i.idx] &&
          i.idxInNextNodes.every(
            (idxInNext) => actions[i.idx].tokenMintOuts[idxInNext[0]]
          )
      )
  );
};

/**
 * Get the mint from prior actions
 *
 * This assumes that the action given is not an initial action
 */
const getMintAccountsFromPriorActions = (
  actionIdx: number,
  actions: (IsolatedAction | undefined)[],
  nextNodes: NextNodeNoNextMint[][][]
): PublicKey[] => {
  const previousActions = getPreviousActions(actionIdx, actions, nextNodes);
  const mints = previousActions
    .map((prev) => {
      const action = actions[prev.idx] as Action; // assume its not undefined because getPreviousAction ensures its not undefined
      const mintsForPrevAction = prev.idxInNextNodes.map(
        (idxInNextNodeInstance) =>
          action.tokenMintOuts[idxInNextNodeInstance[0]]
      );
      return mintsForPrevAction;
    })
    .flat();

  return removeDuplicateMints(mints).map((pk) => new PublicKey(pk));
};

const fillMintForNextNode = (
  nextNodesNoMint: NextNodeNoNextMint[][],
  outMints: PublicKey[],
  allInputMints: MintAndTokenAccount[][]
): NextNode[][] => {
  if (nextNodesNoMint.length === 0) return [];
  if (nextNodesNoMint.length !== outMints.length) {
    throw `Expected the number of out mints to equal ${outMints.length}`;
  }
  return nextNodesNoMint.map((next, mintIdx) => {
    const mint = outMints[mintIdx];
    return next.map((n) => {
      const nextActionInputMints = allInputMints[n.actionIdx].map((p) =>
        p.mint.toBase58()
      );
      const nextInputMintIdx = nextActionInputMints.indexOf(mint.toBase58());
      if (nextInputMintIdx === -1) {
        console.log(outMints.map((o) => o.toBase58()));
        throw `Unexpected error: Expected to find mint ${mint.toBase58()} in the set of input mints for the next action with action index ${
          n.actionIdx
        }`;
      }
      return {
        ...n,
        nextInputMintIdx,
      };
    });
  });
};

const getActionVertexInfoNoMinOuts = (
  action: IsolatedAction,
  nextNodes: NextNodeNoNextMint[][],
  tokenMintOuts: PublicKey[],
  allInputMintAndTokenAccounts: MintAndTokenAccount[][],
  inDegree: number,
  actionIdx: number
): ActionVertexInfoIsolated => {
  const nextNodesFilled = fillMintForNextNode(
    nextNodes,
    tokenMintOuts,
    allInputMintAndTokenAccounts
  );
  const vertexInfo: ActionVertexInfoIsolated = {
    inDegree,
    nextNodes: nextNodesFilled,
    numberInputMints: allInputMintAndTokenAccounts[actionIdx].length,
    actionType: action.opts?.actionType ?? { multiFungibleTokens: {} }, // assume fungible token
  };
  return vertexInfo;
};

/**
 * Fill in an isolated action to contain info on the next nodes to hit,
 * the minimum outputs, and output token accounts
 */
const isolatedActionToActionFilled = async (
  actionIdx: number,
  isolatedAction: IsolatedActionWithNonSpecific,
  outputEstimates: MintWithAmount<BNIsh, PublicKey>[],
  inputEstimates: MintWithAmount<BNIsh, PublicKey>[],
  nonSpecificAction: NonSpecificAction,
  allInputMintAndTokenAccounts: MintAndTokenAccount[][],
  authority: PublicKey,
  vertexInfo: ActionVertexInfoIsolated,
  preferredTokenAccounts?: PreferredTokenAccounts
): Promise<Action> => {
  if (isolatedAction.tokenMintOuts.length !== outputEstimates.length)
    throw `Expected the number of token mints out to equal the number of estimate outs`;
  const tokenAccountOuts = await Promise.all(
    isolatedAction.tokenMintOuts.map((mint) =>
      getTokenAccountUsedByMalloc(authority, mint, preferredTokenAccounts)
    )
  );

  const outputEstAmounts = outputEstimates.map((est) => est.amount);
  const minOuts = nonSpecificAction.opts?.minimumOutToEstimateRatio
    ? getMinOutsFromEstimates(
        outputEstAmounts,
        nonSpecificAction.opts.minimumOutToEstimateRatio
      )
    : null;

  const vertexInfoFilled = vertexInfo as ActionVertexInfo;

  // Get the minimum outs
  vertexInfoFilled.minOuts = nonSpecificAction.overrides?.minOuts ?? minOuts;

  return {
    ...isolatedAction,
    actionVertexInfo: vertexInfoFilled,
    inputTokenAccounts: allInputMintAndTokenAccounts[actionIdx].map(
      (inp) => inp.tokenAccount
    ),
    tokenMintsIn: allInputMintAndTokenAccounts[actionIdx].map(
      (inp) => inp.mint
    ),
    actionTypeUID: nonSpecificAction.actionTypeUID,
    tokenAccountOuts,
    estimateOuts: outputEstimates,
    estimateIns: inputEstimates,
    nonSpecificAction,
  };
};

const getMinOutsFromEstimates = (
  estimates: BNIsh[],
  estimateMultiplier: number[]
): BN[] => {
  if (estimates.length !== estimateMultiplier.length)
    throw `Expected the minimum output multipliers to be the same length as the number of output estimates`;
  return estimates.map((est, i) =>
    new BN(est).mul(new BN(estimateMultiplier[i] * 1000)).divn(1000)
  );
};

/**
 * Get the mints used for an action based off of past inputs and the init mints
 *
 * If an action has both a prior action and an init mint, it will add mints from both
 * while keeping the init mints in the beginning of the list to ensure that initing
 * still works properly (it relies on the init mints being the first amounts)
 */
const getMintAccounts = (
  actionIdx: number,
  isolatedActions: (IsolatedAction | undefined)[],
  nextNodes: NextNodeNoNextMint[][][],
  initActionIdxs: number[],
  inMintAccounts: PublicKey[]
): PublicKey[] => {
  let mintAccounts = [];
  if (initActionIdxs.includes(actionIdx)) {
    mintAccounts.push(...inMintAccounts);
  }
  const mintsFromPrior = getMintAccountsFromPriorActions(
    actionIdx,
    isolatedActions,
    nextNodes
  );
  const mintsFromPriorFiltered = removeMints(mintsFromPrior, mintAccounts);
  mintAccounts.push(...mintsFromPriorFiltered);

  return mintAccounts;
};

const getIsolatedActionInfo = async (
  actionIdxToId: string[],
  order: number[][],
  initActionIdxs: number[],
  nextNodes: NextNodeNoNextMint[][][],
  inMints: PublicKey[],
  construction: NonSpecificConstruction,
  buildActionMap: BuildActionMap,
  provider: Provider,
  authority: PublicKey,
  config: BuildIsolateActionConfig,
  preferredTokenAccounts?: PreferredTokenAccounts
): Promise<{
  isolatedActions: IsolatedActionWithNonSpecific[];
  estimateOutFns: EstimateOutFN[];
  mintInsWithTokens: MintAndTokenAccount[][];
  constructionFilled: NonSpecificConstructionFilledUnserializable;
}> => {
  const constructionFilled = deepCloneObject(
    construction
  ) as NonSpecificConstructionFilledUnserializable;
  const isolatedActions: (IsolatedAction | undefined)[] = [
    ...Array(actionIdxToId.length).fill(undefined),
  ];

  const estimateOutFns: (EstimateOutFN | undefined)[] = [
    ...Array(actionIdxToId.length).fill(undefined),
  ];

  const mintInsWithTokens: (MintAndTokenAccount[] | undefined)[] = [
    ...Array(actionIdxToId.length).fill(undefined),
  ];

  let i = 0;
  for (let o = 0; o < order.length; o++) {
    const orderGroup = order[o];
    for (let q = 0; q < orderGroup.length; q++) {
      const actionIdx = orderGroup[q];

      const mintAccountsIn = getMintAccounts(
        actionIdx,
        isolatedActions,
        nextNodes,
        initActionIdxs,
        inMints
      );

      const actionData = getNonSpecificActionDataFromId(
        actionIdxToId[actionIdx],
        construction.actionDatas
      ) as _NonSpecificActionFilled<PublicKey, BNIsh>;

      // Update the filled construction
      actionData.inMints = [...mintAccountsIn];
      updateNonSpecificActionDataFromId(
        actionIdxToId[actionIdx],
        constructionFilled,
        actionData
      );

      const buildFN =
        buildActionMap[actionData.actionTypeUID].buildIsolatedAction;
      if (!buildFN)
        throw errors.newBuilderError(
          `Could not find a builder function for action with UID ${actionData.actionTypeUID}`
        );

      // Get the token account which malloc will use
      const tokenWithMintAccounts: MintAndTokenAccount[] = await Promise.all(
        mintAccountsIn.map(async (mint) => {
          return {
            tokenAccount: await getTokenAccountUsedByMalloc(
              authority,
              new PublicKey(mint),
              preferredTokenAccounts
            ),
            mint,
          };
        })
      );

      mintInsWithTokens[actionIdx] = tokenWithMintAccounts;

      try {
        const { isolatedAction: _isolatedAction, estimateOuts } = await buildFN(
          {
            inp: actionData.inputs,
            inputTokens: tokenWithMintAccounts,
            actionPID: new PublicKey(actionData.actionPID),
            provider,
            authority,
            config,
            opts: {
              preferredTokenAccounts: preferredTokenAccounts,
            },
          }
        );
        const isolatedAction: IsolatedActionWithNonSpecific = {
          ..._isolatedAction,
          nonSpecificAction: actionData,
          actionName: prettyPrintActionId(actionIdxToId[actionIdx]),
        };
        isolatedActions[actionIdx] = isolatedAction;
        estimateOutFns[actionIdx] = estimateOuts;

        // Update the filled action
        actionData.outMints = [...isolatedAction.tokenMintOuts];
        updateNonSpecificActionDataFromId(
          actionIdxToId[actionIdx],
          constructionFilled,
          actionData
        );
      } catch (e) {
        console.error(e);
        throw errors.newActionBuilderError(
          actionData.actionTypeUID,
          prettyPrintActionId(actionIdxToId[actionIdx]),
          actionIdx,
          e,
          nspFilledUnserialToSerial(constructionFilled)
        );
      }

      i++;
    }
  }
  return {
    estimateOutFns: estimateOutFns as EstimateOutFN[],
    isolatedActions: isolatedActions as IsolatedActionWithNonSpecific[],
    mintInsWithTokens: mintInsWithTokens,
    constructionFilled,
  };
};

interface BuildConstructionOpts {
  authority?: PublicKey;
  preferredTokenAccounts?: PreferredTokenAccounts;
}

const _buildConstructionFromNonSpecific = async (
  construction: NonSpecificConstruction,
  buildActionMap: BuildActionMap,
  provider: Provider,
  config: BuildIsolateActionConfig,
  opts?: BuildConstructionOpts
): Promise<{
  construction: Construction;
  filledNonSpecific: NonSpecificConstructionFilledUnserializable;
}> => {
  const initialMints = Object.keys(construction.source.mints);
  const amountsIn = Object.keys(construction.source.mints).map(
    (mint) => construction.source.mints[mint]
  );
  const initialActions = Object.keys(construction.source.nextActions);
  const initialSplits = Object.values(construction.source.nextActions);

  const initMintAccounts = initialMints.map((p) => {
    try {
      return new PublicKey(p);
    } catch (error) {
      throw newBuilderError(
        `The input mint, ${p}, is not a valid base-58 public key`
      );
    }
  });

  const authorityChecked = opts?.authority ?? provider.wallet.publicKey;

  const { order, actionIdxToId, initActionIdxs, inDegrees, nextNodes } =
    buildSeqListOfActionCalls(
      construction.actionDatas,
      initialActions,
      initMintAccounts.length
    );

  const {
    isolatedActions,
    estimateOutFns,
    mintInsWithTokens,
    constructionFilled,
  } = await getIsolatedActionInfo(
    actionIdxToId,
    order,
    initActionIdxs,
    nextNodes,
    initMintAccounts,
    construction,
    buildActionMap,
    provider,
    authorityChecked,
    config,
    opts?.preferredTokenAccounts
  );

  const vertexInfos = isolatedActions.map((action, actionIdx) => {
    try {
      return getActionVertexInfoNoMinOuts(
        action,
        nextNodes[actionIdx],
        action.tokenMintOuts,
        mintInsWithTokens,
        inDegrees[actionIdx],
        actionIdx
      );
    } catch (e) {
      console.error(e);
      throw errors.newActionBuilderError(
        action.nonSpecificAction.actionTypeUID,
        prettyPrintActionId(actionIdxToId[actionIdx]),
        actionIdx,
        e,
        nspFilledUnserialToSerial(constructionFilled)
      );
    }
  });

  let estimateRet: EstimateRet;
  try {
    estimateRet = await getConstructionEstimates(
      order,
      isolatedActions,
      construction,
      initActionIdxs,
      initMintAccounts,
      vertexInfos,
      estimateOutFns,
      mintInsWithTokens
    );
  } catch (e) {
    if ((e as MallocError).mallocErrorMarker === "MALLOC ERROR MARKER") {
      throw {
        ...e,
        filledNonSpecificConstruction:
          nspFilledUnserialToSerial(constructionFilled),
      } as MallocActionBuilderError;
    } else throw e;
  }

  const { estimateIns, estimateOuts } = estimateRet;

  // Update the estimates for the actions in the filled construction
  actionIdxToId.forEach((id, actionIdx) => {
    console.log("ID", id);
    const action = getNonSpecificActionDataFromId(
      id,
      constructionFilled.actionDatas
    ) as _NonSpecificActionFilled<PublicKey, BNIsh>;
    action.inEstimates = estimateIns[actionIdx];
    action.outEstimates = estimateOuts[actionIdx];
    console.log(action);
    updateNonSpecificActionDataFromId(id, constructionFilled, action);
  });

  const actions = await Promise.all(
    isolatedActions.map(async (isolated, i) => {
      const actionData = getNonSpecificActionDataFromId(
        actionIdxToId[i],
        construction.actionDatas
      );
      return await isolatedActionToActionFilled(
        i,
        isolated,
        estimateOuts[i],
        estimateIns[i],
        actionData,
        mintInsWithTokens,
        authorityChecked,
        vertexInfos[i],
        opts?.preferredTokenAccounts
      );
    })
  );

  return {
    construction: {
      actions,
      actionOrders: order,
      initialActionIndices: initActionIdxs,
      initialSplits: initialSplits,
      amountIns: amountsIn.map((s) => new BN(s)),
    },
    filledNonSpecific: constructionFilled,
  };
};

// TODO: CLEAN UP AND TEST this function (estimate outs better)
// TODO: test the split in tokens
export const buildConstructionFromNonSpecific = async (
  construction: NonSpecificConstruction,
  buildActionMap: BuildActionMap,
  provider: Provider,
  config: BuildIsolateActionConfig,
  opts?: BuildConstructionOpts
): Promise<Construction> => {
  const { construction: constructionBuilt } =
    await _buildConstructionFromNonSpecific(
      construction,
      buildActionMap,
      provider,
      config,
      opts
    );
  return constructionBuilt;
};

export const buildConstructionFromNonSpecificWithFilled = async (
  construction: NonSpecificConstruction,
  buildActionMap: BuildActionMap,
  provider: Provider,
  config: BuildIsolateActionConfig,
  opts?: BuildConstructionOpts
) => {
  const { construction: constructionBuilt, filledNonSpecific } =
    await _buildConstructionFromNonSpecific(
      construction,
      buildActionMap,
      provider,
      config,
      opts
    );
  return {
    constructionBuilt,
    filledNonSpecific,
  };
};

export const nspFilledUnserialToSerial = (
  filledUnserial: NonSpecificConstructionFilledUnserializable
): NonSpecificConstructionFilled => {
  const actionDatasSerial: NonSpecificConstructionFilled["actionDatas"] = {};
  Object.keys(filledUnserial.actionDatas).forEach((k) => {
    const actionWrapped = filledUnserial.actionDatas[k];
    if (isNonSpecificActionGroup(actionWrapped)) {
      const group = castToActionGroup(actionWrapped);
      actionDatasSerial[k] = {
        actionGroup: {
          ...group,
          actions: Object.keys(group.actions).reduce((map, actionId) => {
            map[actionId] = serializeActionFilled(group.actions[actionId]);
            return map;
          }, {}),
        },
      };
    } else {
      actionDatasSerial[k] = {
        action: serializeActionFilled(
          castToStandaloneAction(filledUnserial.actionDatas[k])
        ),
      };
    }
  });
  filledUnserial.actionDatas = actionDatasSerial;
  return {
    ...filledUnserial,
  };
};

const serializeActionFilled = (
  action: _NonSpecificActionFilled<PublicKey, BNIsh>
): _NonSpecificActionFilled<string, string> => {
  return {
    ...action,
    inMints: action.inMints?.map((m) => m.toString()),
    inEstimates: action.inEstimates?.map(serializeMintWithAmount),
    outEstimates: action.outEstimates?.map(serializeMintWithAmount),
    outMints: action.outMints?.map((m) => m.toString()),
  };
};

function serializeMintWithAmount(
  m: MintWithAmount<BNIsh, PublicKey>
): MintWithAmount<string, string> {
  return {
    mint: m.mint.toString(),
    amount: m.amount.toString(),
  };
}