import { PublicKey } from "@solana/web3.js";
import BN from "bn.js";
import {
  ActionVertexInfo,
  isNonSpecificActionGroup,
  NonSpecificConstruction,
  NonSpecificAction,
  NonSpecificActionGroup,
  castToActionGroup,
  NextNode,
  NextNodeNoNextMint,
  NonSpecificConstructionFilledUnserializable,
  _NonSpecificActionFilled,
} from "./interfaces";
import { deepCloneObject, flattenArrObject } from "./utils/object";

export const _checkThatAllHitInActionGroup = (
  group: NonSpecificActionGroup
): boolean => {
  const ids = Object.keys(group.actions);
  const hit = ids.map((_) => false);
  const frontier = [group.first];
  while (frontier.length !== 0) {
    const elem = frontier.pop();
    const elemIdx = ids.indexOf(elem);
    hit[elemIdx] = true;
    for (const neighborId in flattenArrObject(
      group.actions[elem].nextActions
    )) {
      const neighborIdx = ids.indexOf(neighborId);
      // If the neighborIdx is negative 1, then it is referencing an action outside
      // the current group
      if (neighborIdx !== -1 && !hit[neighborIdx]) {
        frontier.push(neighborId);
      }
    }
  }
  return hit.every((v) => v === true);
};

export const getInDegrees = (graph: number[][][]): number[] => {
  const inDegrees = Array(graph.length).fill(0);
  graph.forEach((vertex) =>
    vertex.forEach((vertexNexts) =>
      vertexNexts.forEach((v) => (inDegrees[v] += 1))
    )
  );
  return inDegrees;
};

/**
 * Get a sort on the graph. This should be DFS based as we want
 * to go down "legs" of the tree 1 at a time
 *
 * The algorithm works by using DFS and only following the dfs if
 * the amount of times a node is hit equals its in degree
 */
export const topologicalSort = (
  graph: number[][][],
  inDegrees: number[],
  init: number[]
): number[] => {
  const hitCount = Array(graph.length).fill(0);
  // Ensure that the initial in the frontier do not have other in degrees
  const frontier = [...init.filter((idx) => inDegrees[idx] === hitCount[idx])];
  frontier.map((i) => (hitCount[i] += 1));
  const inFrontier = Array(graph.length).fill(false);
  return topologicalSortRecurse(
    graph,
    hitCount,
    frontier,
    inDegrees,
    inFrontier
  );
};

const topologicalSortRecurse = (
  graph: number[][][],
  hitCount: number[],
  frontier: number[],
  inDegrees: number[],
  inFrontier: boolean[]
): number[] => {
  const ret = [];
  while (frontier.length !== 0) {
    const idx = frontier.pop();
    // Conditionally add to the frontier
    graph[idx].forEach((nextNodes) =>
      nextNodes.forEach((next) => {
        hitCount[next] += 1;
        if (hitCount[next] === inDegrees[next] && !inFrontier[next]) {
          frontier.push(next);
          inFrontier[next] = true;
        }
      })
    );
    ret.push(
      idx,
      ...topologicalSortRecurse(
        graph,
        hitCount,
        frontier,
        inDegrees,
        inFrontier
      )
    );
  }
  return ret;
};

/**
 * Convert to an adjacency list where every next action edge pointing
 * to outside of the action group is thrown out
 */
const _convertActionGroupToInternalDoubleAdjacency = (
  group: NonSpecificActionGroup,
  actionGroupIds: string[]
): number[][][] => {
  const nextActions = actionGroupIds.map((id) =>
    group.actions[id].nextActions.map((nObj) => Object.keys(nObj))
  );
  return nextActions.map((outerNexts) =>
    outerNexts.map((nexts) =>
      nexts.map((id) => actionGroupIds.indexOf(id)).filter((idx) => idx !== -1)
    )
  );
};

/**
 * Convert a set of action datas to an adjacency list
 */
export const _convertToDoubleAdjacencyList = (
  actionDatasFormatted: NonSpecificConstruction["actionDatas"],
  actionIdxToId: string[],
  initActions: string[]
): { graph: number[][][]; initActions: number[] } => {
  const newActionDatas: any = deepCloneObject(actionDatasFormatted);
  const actionIds = actionIdxToId;
  actionIds.forEach((id) => {
    if (isNonSpecificActionGroup(actionDatasFormatted[id])) {
      const group = newActionDatas[id]
        .actionGroup as unknown as NonSpecificActionGroup;
      if (!_checkThatAllHitInActionGroup(group)) {
        throw `Expected no actions to be left uncalled`;
      }
      const groupActionIds = Object.keys(group.actions);
      const nextActionsMap = groupActionIds
        .map((id) => group.actions[id].nextActions)
        .flat();
      newActionDatas[id].nextActions = nextActionsMap.map((nextActionsObj) =>
        Object.keys(nextActionsObj)
          // Only look at the actions which are found in the external actions
          .map((k) => actionIds.indexOf(k))
          .filter((idx) => idx !== -1)
      ) as number[][];
    } else {
      const action = newActionDatas[id].action as unknown as NonSpecificAction;
      newActionDatas[id].nextActions = action.nextActions.map(
        (nextActionsObj) =>
          Object.keys(nextActionsObj).map((k) => {
            const idx = actionIds.indexOf(k);
            if (idx === -1)
              throw `There is an unexpected next action of id ${k} for action group ${id}`;
            return idx;
          })
      ) as number[][];
    }
  });
  const graph = actionIds.map(
    (actionId) => newActionDatas[actionId].nextActions
  );
  const _initActions = initActions.map((id) => actionIds.indexOf(id));
  return {
    graph,
    initActions: _initActions,
  };
};

/**
 * The first number corresponds to the outer index and the inner
 * corresponds to the set of inner indices of an action group.
 * If it is just an action, then it is left empty
 */
type AccessIdx = [number, number[]];

/**
 * Take in a non specific construction and format all action id's within a group to match the @code{createInnerActionId}
 */
const formatActionIds = (
  actions: NonSpecificConstruction["actionDatas"]
): NonSpecificConstruction["actionDatas"] => {
  // TODO: this may be too slow...
  const actionsNew = deepCloneObject(actions);
  const oldOuterIds = Object.keys(actionsNew);
  for (let i = 0; i < oldOuterIds.length; i++) {
    const outerAction = oldOuterIds[i];
    // get the group new action ids
    if (isNonSpecificActionGroup(actionsNew[outerAction])) {
      const group = castToActionGroup(actionsNew[outerAction]);
      const oldInnerActionIds = Object.keys(group.actions);
      for (let i = 0; i < oldInnerActionIds.length; i++) {
        const innerAction = oldInnerActionIds[i];
        const newNexts = group.actions[innerAction].nextActions.map(
          (nextActions) => {
            const newNext: NonSpecificAction["nextActions"][number] = {
              ...nextActions,
            };
            for (const next in nextActions) {
              if (oldInnerActionIds.includes(next)) {
                newNext[createInnerActionId(outerAction, next)] = newNext[next];
                delete newNext[next];
              }
            }
            return newNext;
          }
        );
        group.actions[createInnerActionId(outerAction, innerAction)] =
          group.actions[innerAction];
        group.actions[
          createInnerActionId(outerAction, innerAction)
        ].nextActions = newNexts;
        delete group.actions[innerAction];
      }
      group.first = createInnerActionId(outerAction, group.first);
      actionsNew[outerAction] = { actionGroup: group };
    }
  }
  return actionsNew;
};

const buildSeqListOfActionCallsNonFlat = (
  actionDatas: NonSpecificConstruction["actionDatas"],
  initActions: string[]
): {
  order: AccessIdx[];
  actionIdxToInnerId: string[][];
  actionIdxToOuterId: string[];
  initialActionIdxs: number[];
  inDegrees: number[][];
} => {
  const actionIds = Object.keys(actionDatas);
  const { graph, initActions: initIndices } = _convertToDoubleAdjacencyList(
    actionDatas,
    actionIds,
    initActions
  );

  if (graph.length !== actionIds.length) throw `Expected no isolated actions`;

  const inDegreesOuter = getInDegrees(graph);
  const outerOrder = topologicalSort(graph, inDegreesOuter, initIndices);
  const inDegrees2D = Array(inDegreesOuter.length).fill([]);

  // Now use the outer order to get inner order for action groups
  const orderMapped = outerOrder.map((actionIdx) => {
    if (isNonSpecificActionGroup(actionDatas[actionIds[actionIdx]])) {
      const group = castToActionGroup(actionDatas[actionIds[actionIdx]]);
      const actionGroupIds = Object.keys(group.actions);
      const actionGroupGraph = _convertActionGroupToInternalDoubleAdjacency(
        group,
        actionGroupIds
      );
      const actionInDegrees = getInDegrees(actionGroupGraph);
      // Add the in degrees coming into the group for the first element
      const firstIdx = actionGroupIds.indexOf(group.first);
      if (firstIdx === -1)
        throw `Expected the first id in group ${actionIds[actionIdx]} to be an action in the group`;
      actionInDegrees[firstIdx] += inDegreesOuter[actionIdx];

      inDegrees2D[actionIdx] = actionInDegrees;
      return {
        order: [
          actionIdx,
          topologicalSort(actionGroupGraph, actionInDegrees, [
            actionGroupIds.indexOf(group.first),
          ]),
        ] as AccessIdx,
      };
    }
    inDegrees2D[actionIdx] = [inDegreesOuter[actionIdx]];
    return {
      order: [actionIdx, []] as AccessIdx,
    };
  });

  const order = orderMapped.map((o) => o.order);
  // const inDegrees2D = orderMapped.map((o) => o.inDegrees);

  const actionIdxToInnerId = actionIds.map((id) =>
    isNonSpecificActionGroup(actionDatas[id])
      ? Object.keys(castToActionGroup(actionDatas[id]).actions).map(
          (innerActionId) => innerActionId
        )
      : [id]
  );
  return {
    order,
    actionIdxToInnerId,
    actionIdxToOuterId: actionIds,
    initialActionIdxs: initIndices,
    inDegrees: inDegrees2D,
  };
};

/**
 * Return a list of the hit order where the outer list corresponds to
 * the order of transactions and the inner list corresponds to the actions
 * grouped into an instruction
 *
 * Every action idx is in reference to 1 FLAT list of actions. This algorithm works
 * on the assumption that every vertex is only every visited once and that the input graph is a DAG
 *
 *
 * @return actionIdxToId - a list of strings where the string at id index i corresponds to action of the same id
 * @return order - The order hit with the outer group being for action/ action group and the inner one being the order for internal actions. Order maps from 
 * call order index -> action index
 * @return initActionIdxs - the initial action index's which get hit
 *
 * NOTE: this returns an unknown result for cyclic graphs
 */
export const buildSeqListOfActionCalls = (
  actionDatas: NonSpecificConstruction["actionDatas"],
  initActions: string[],
  numberOfInputTokens: number
): {
  order: number[][];
  initActionIdxs: number[];
  actionIdxToId: string[];
  inDegrees: number[];
  nextNodes: NextNodeNoNextMint[][][];
} => {
  const actionDatasFormattedIds = formatActionIds(actionDatas);
  const {
    order,
    actionIdxToInnerId,
    actionIdxToOuterId,
    initialActionIdxs,
    inDegrees,
  } = buildSeqListOfActionCallsNonFlat(actionDatasFormattedIds, initActions);
  const { orderFlat, actionIdxToIdFlat, initActionIdxsFlat, inDegreesFlat } =
    flattenOrder(
      order,
      actionIdxToInnerId,
      actionIdxToOuterId,
      initialActionIdxs,
      inDegrees,
      actionDatasFormattedIds
    );
  const groupingLengths = order.map((accessIdx) =>
    accessIdxIsGroup(accessIdx) ? accessIdx[1].length : 1
  );

  const orderByGroups = groupingLengths.map((len) => orderFlat.splice(0, len));
  if (orderByGroups.length !== order.length)
    throw `Unexpected error in creating an ordering`;

  const nextNodes = actionIdxToIdFlat.map((id) =>
    getNextNodes(
      actionIdxToIdFlat,
      getNonSpecificActionDataFromId(id, actionDatasFormattedIds, true),
      actionDatasFormattedIds
    )
  );
  // Check to ensure that there are the appropriate # of actions hit in the ordering
  if (
    orderByGroups.reduce((prev, group) => prev + group.length, 0) !==
    actionIdxToIdFlat.length
  )
    throw "Expected the action grouping length to equal the number of ids";

  // Ensure that all action idxs are "found" elements
  orderByGroups.forEach((group) => {
    group.forEach((idx) => {
      if (idx === -1)
        throw `An error occurred with an unexpected next action name`;
    });
  });

  // add number of input mints to all the initial in degrees to account for the initial inputs
  // this means that they get "hit" once for each init token
  initActionIdxsFlat.forEach((idx) => {
    inDegreesFlat[idx] += numberOfInputTokens;
  });

  return {
    order: orderByGroups,
    initActionIdxs: initActionIdxsFlat,
    actionIdxToId: actionIdxToIdFlat,
    inDegrees: inDegreesFlat,
    nextNodes,
  };
};

const getNextNodes = (
  idxToActionId: string[],
  action: NonSpecificAction,
  actionDatas: NonSpecificConstruction["actionDatas"]
): NextNodeNoNextMint[][] => {
  const getInnerNextNode = (
    innerNextNode: NonSpecificAction["nextActions"][number]
  ): NextNodeNoNextMint[] => {
    const nexts = Object.keys(innerNextNode);
    return nexts.map((id) => {
      return {
        actionIdx: getIndexOfWithGroups(id, idxToActionId, actionDatas),
        fraction: innerNextNode[id],
      };
    });
  };
  return action.nextActions.map(getInnerNextNode);
};

const flattenOrder = (
  order: AccessIdx[],
  actionIdxToInnerId: string[][],
  actionIdxToOuterId: string[],
  initActionIdxs: number[],
  inDegrees: number[][],
  actionDatasFormatted: NonSpecificConstruction["actionDatas"]
): {
  orderFlat: number[];
  actionIdxToIdFlat: string[];
  initActionIdxsFlat: number[];
  inDegreesFlat: number[];
} => {
  const actionIdxToIdFlat = actionIdxToInnerId.flat();
  const newOrder = order
    .map((accessIdx) => {
      if (accessIdxIsGroup(accessIdx)) {
        const ids = accessIdx[1].map(
          (innerIdx) => actionIdxToInnerId[accessIdx[0]][innerIdx]
        );
        const outerId = actionIdxToOuterId[accessIdx[0]];
        return ids.map((innerId) => actionIdxToIdFlat.indexOf(innerId));
      } else {
        const outerId = actionIdxToOuterId[accessIdx[0]];
        const idx = getIndexOfWithGroups(
          outerId,
          actionIdxToIdFlat,
          actionDatasFormatted
        );
        return [idx];
      }
    })
    .flat();
  const initActionIdxsFlat = initActionIdxs.map((i) => {
    // find the access info corresponding to the access info which corresponds
    // to the index in the init action list
    const accessIdx = order.find((o) => o[0] === i);
    if (accessIdxIsGroup(accessIdx)) {
      const outerId = actionIdxToOuterId[accessIdx[0]];
      // Get the first element called in the group
      const innerActionIdx = accessIdx[1][0];
      const innerId = actionIdxToInnerId[accessIdx[0]][innerActionIdx];
      return actionIdxToIdFlat.indexOf(innerId);
    } else {
      const outerId = actionIdxToOuterId[accessIdx[0]];
      return actionIdxToIdFlat.indexOf(outerId);
    }
  });
  return {
    orderFlat: newOrder,
    actionIdxToIdFlat,
    initActionIdxsFlat,
    inDegreesFlat: inDegrees.flat(),
  };
};

const accessIdxIsGroup = (a: AccessIdx) => a[1].length !== 0;

const ACTION_DELIMITER = "!@#@!";

export const createInnerActionId = (actionGroupId: string, actionId: string) =>
  `${actionGroupId}${ACTION_DELIMITER}${actionId}`;

export const prettyPrintActionId = (id: string) => {
  if (id.includes(ACTION_DELIMITER)) {
    return `group: ${id.replace(ACTION_DELIMITER, ", inner-action: ")}`;
  }
  return `${id}`;
};

export const getNonSpecificActionDataFromId = (
  id: string,
  actionDatas: NonSpecificConstruction["actionDatas"],
  actionDatasIdsFormatted?: boolean
): NonSpecificAction => {
  const split = id.split(ACTION_DELIMITER);
  if (split.length === 1) {
    return (actionDatas[split[0]] as any).action as NonSpecificAction;
  } else {
    return castToActionGroup(actionDatas[split[0]]).actions[
      actionDatasIdsFormatted ? id : split[1]
    ];
  }
};

export const updateNonSpecificActionDataFromId = (
  id: string,
  construction:
    | NonSpecificConstruction
    | NonSpecificConstructionFilledUnserializable,
  update: NonSpecificAction | _NonSpecificActionFilled<PublicKey, BN>,
  actionDatasIdsFormatted?: boolean
) => {
  const split = id.split(ACTION_DELIMITER);
  if (split.length === 1) {
    (construction.actionDatas[split[0]] as any).action = update;
  } else {
    castToActionGroup(construction.actionDatas[split[0]]).actions[
      actionDatasIdsFormatted ? id : split[1]
    ] = update;
  }
};

const getIndexOfWithGroups = (
  id: string,
  idxToActionId: string[],
  actionDatasFormatted: NonSpecificConstruction["actionDatas"]
) => {
  let idx = idxToActionId.indexOf(id);
  if (idx === -1) {
    // If the next id is a non specific action, get the id of the first action in the group
    if (isNonSpecificActionGroup(actionDatasFormatted[id])) {
      const group = castToActionGroup(actionDatasFormatted[id]);
      idx = idxToActionId.indexOf(group.first);
    }
    if (idx === -1) throw `Could not find next action of id ${id}`;
  }
  return idx;
};