import {
  addAssets,
  Address,
  Assets,
  credentialToAddress,
  getInputIndices,
  LucidEvolution,
  OutRef,
  toHex,
  UTxO,
  validatorToScriptHash,
} from '@lucid-evolution/lucid';
import {
  AccountContent,
  AssetSnapshot,
  AssetState,
  E2S2SIndex,
  E2S2SIndicesPerAsset,
  EpochToScaleKey,
  EpochToScaleToSumEntry,
  fromSPInteger,
  mkSPInteger,
  parseSnapshotEpochToScaleToSumDatumOrThrow,
  ProcessRequestAccountContent,
  SnapshotEpochToScaleToSumContent,
  spAdd,
  spDiv,
  SPInteger,
  spMul,
  spSub,
  StabilityPoolContent,
  StateSnapshot,
  SumSnapshot,
} from './types-new';
import {
  readonlyArray as RA,
  array as A,
  function as F,
  option as O,
} from 'fp-ts';
import {
  AssetClass,
  getInlineDatumOrThrow,
  isSameAssetClass,
  isSameOutRef,
  mkAssetsOf,
} from '@3rd-eye-labs/cardano-offchain-common';
import { repsertWithReadonlyArr } from '../../utils/array-utils';
import { match, P } from 'ts-pattern';
import {
  fromSysParamsCredential,
  SystemParams,
} from '../../types/system-params';
import { mkStabilityPoolValidatorFromSP } from './scripts';

export const BASE_MAX_TX_FEE = 1_250_000n;

export const MAX_E2S2S_ENTRIES_COUNT = 5;
const newScaleMultiplier = 1_000_000_000n;

export const initSumVal: SPInteger = { value: 0n };

export const initSpState: StateSnapshot = {
  productVal: mkSPInteger(1n),
  depositVal: { value: 0n },
  epoch: 0n,
  scale: 0n,
};

export function mkStabilityPoolAddr(
  lucid: LucidEvolution,
  sysParams: SystemParams,
): Address {
  return credentialToAddress(
    lucid.config().network!,
    {
      hash: validatorToScriptHash(
        mkStabilityPoolValidatorFromSP(sysParams.stabilityPoolParams),
      ),
      type: 'Script',
    },
    sysParams.stabilityPoolParams.stakeCredential != null
      ? fromSysParamsCredential(sysParams.stabilityPoolParams.stakeCredential)
      : undefined,
  );
}

export function isSameEpochToScaleKey(
  a: EpochToScaleKey,
  b: EpochToScaleKey,
): boolean {
  return a.epoch === b.epoch && a.scale === b.scale;
}

type StabilityPoolListIdx = {
  spListIdx: number;
  sumSnapshot: readonly [EpochToScaleKey, SumSnapshot];
};
type SnapshotIdx = {
  snapshotUtxo: UTxO;
  snapshotDatum: SnapshotEpochToScaleToSumContent;
  snapshotListIdx: number;
  sumSnapshot: readonly [EpochToScaleKey, SumSnapshot];
};
export type FindE2S2SIdxResult = StabilityPoolListIdx | SnapshotIdx;

function findE2s2sIdx(
  expectedKey: EpochToScaleKey,
  collateralAsset: AssetClass,
  spAssetState: AssetSnapshot,
  e2s2sUtxos: [UTxO, SnapshotEpochToScaleToSumContent][],
): O.Option<FindE2S2SIdxResult> {
  return F.pipe(
    // Try to find the s1 in sp e2s2s.
    F.pipe(
      spAssetState.epoch2scale2sum,
      RA.findIndex(([key, _]) => isSameEpochToScaleKey(key, expectedKey)),
    ),
    O.match<number, O.Option<FindE2S2SIdxResult>>(
      // When such e2sKey non existent in sp e2s2s, find it in e2s2s snapshots
      () =>
        F.pipe(
          e2s2sUtxos,
          A.findFirst(
            ([_, datum]) =>
              isSameAssetClass(datum.collateralAsset, collateralAsset) &&
              F.pipe(
                datum.snapshot,
                RA.exists(([key, _]) =>
                  isSameEpochToScaleKey(expectedKey, key),
                ),
              ),
          ),
          O.map((res) => {
            const listIdx = F.pipe(
              res[1].snapshot,
              RA.findIndex(([key, _]) =>
                isSameEpochToScaleKey(key, expectedKey),
              ),
              O.getOrElse<number>(() => {
                throw new Error(
                  'It was supposed to be there. Some logic error.',
                );
              }),
            );

            return {
              snapshotUtxo: res[0],
              snapshotDatum: res[1],
              snapshotListIdx: listIdx,
              sumSnapshot: res[1].snapshot[listIdx],
            } satisfies SnapshotIdx;
          }),
        ),
      (poolIdx) =>
        O.some({
          spListIdx: poolIdx,
          sumSnapshot: spAssetState.epoch2scale2sum[poolIdx],
        } satisfies StabilityPoolListIdx),
    ),
  );
}

export async function findRelevantE2s2sIdxs(
  lucid: LucidEvolution,
  stabilityPool: StabilityPoolContent,
  accountState: StateSnapshot,
  allSnapshotsOutRefs: OutRef[],
): Promise<[FindE2S2SIdxResult, O.Option<FindE2S2SIdxResult>][]> {
  const e2s2sUtxos = (await lucid.utxosByOutRef(allSnapshotsOutRefs))
    .map((utxo) => {
      return [
        utxo,
        parseSnapshotEpochToScaleToSumDatumOrThrow(getInlineDatumOrThrow(utxo)),
      ] satisfies [UTxO, SnapshotEpochToScaleToSumContent];
    })
    .filter(([_, d]) => toHex(d.iasset) === toHex(stabilityPool.iasset));

  const s1E2sKey: EpochToScaleKey = {
    epoch: accountState.epoch,
    scale: accountState.scale,
  };
  const s2E2sKey: EpochToScaleKey = {
    epoch: accountState.epoch,
    scale: accountState.scale + 1n,
  };

  return stabilityPool.assetStates.map(([collateralAsset, spAssetState]) => {
    const s1Res = F.pipe(
      findE2s2sIdx(s1E2sKey, collateralAsset, spAssetState, e2s2sUtxos),
      // When it's non-existent, we need to find proof it doesn't exist.
      O.getOrElse(() =>
        F.pipe(
          e2s2sUtxos,
          A.findFirst(
            ([_, datum]) =>
              isSameAssetClass(datum.collateralAsset, collateralAsset) &&
              F.pipe(
                datum.snapshot,
                RA.exists(([_, val]) => val.isFirstSnapshot),
              ),
          ),
          O.match<[UTxO, SnapshotEpochToScaleToSumContent], FindE2S2SIdxResult>(
            // Try to find the proof in pool's list.
            () =>
              F.pipe(
                spAssetState.epoch2scale2sum,
                RA.findIndex(([_, dat]) => dat.isFirstSnapshot),
                O.match(
                  () => {
                    throw new Error(
                      "Couldn't find relevant proof for s1 non-existence.",
                    );
                  },
                  (poolIdx) =>
                    ({
                      spListIdx: poolIdx,
                      sumSnapshot: spAssetState.epoch2scale2sum[poolIdx],
                    }) satisfies StabilityPoolListIdx,
                ),
              ),
            (snapshotProof) => {
              const listIdx = F.pipe(
                snapshotProof[1].snapshot,
                RA.findIndex(([_, val]) => val.isFirstSnapshot),
                O.getOrElse<number>(() => {
                  throw new Error(
                    'It was supposed to be there. Some logic error.',
                  );
                }),
              );

              return {
                snapshotUtxo: snapshotProof[0],
                snapshotDatum: snapshotProof[1],
                snapshotListIdx: listIdx,
                sumSnapshot: snapshotProof[1].snapshot[listIdx],
              } satisfies SnapshotIdx;
            },
          ),
        ),
      ),
    );

    const s2Res = findE2s2sIdx(
      s2E2sKey,
      collateralAsset,
      spAssetState,
      e2s2sUtxos,
    );

    // Is actual s1 index?
    if (isSameEpochToScaleKey(s1Res.sumSnapshot[0], s1E2sKey)) {
      return [
        s1Res,
        F.pipe(
          s2Res,
          O.match(
            () => {
              // When s1 is not just a proof and is last in epoch, it proves s2 non-existant.
              if (s1Res.sumSnapshot[1].isLastInEpoch) {
                return O.none;
              } else {
                throw new Error('Expected s2 to be existent.');
              }
            },
            (res) => O.some(res),
          ),
        ),
      ];
    } else {
      // When s1 is just a proof.
      return [
        s1Res,
        F.pipe(
          s2Res,
          O.match(
            () => {
              // When the non-existance proof works for s2 as well
              if (
                s1Res.sumSnapshot[1].isFirstSnapshot &&
                (s1Res.sumSnapshot[0].epoch > s2E2sKey.epoch ||
                  (s1Res.sumSnapshot[0].epoch === s2E2sKey.epoch &&
                    s1Res.sumSnapshot[0].scale > s2E2sKey.scale))
              ) {
                return O.none;
              }

              throw new Error("S1 proof doesn't work for s2.");
            },
            // s2 exists.
            (res) => O.some(res),
          ),
        ),
      ];
    }
  });
}

export function createProcessRequestAccountRedeemer(
  e2s2sIdxs: [FindE2S2SIdxResult, O.Option<FindE2S2SIdxResult>][],
  otherRefInputs: UTxO[],
  currentTime: bigint,
): {
  e2s2sRefInputs: UTxO[];
  mkProcessRequestAccountRedeemerContent: (
    poolInputIdx: bigint,
    accountInputIdx: bigint,
  ) => ProcessRequestAccountContent;
} {
  const allE2s2sSnapshotRefInputs = F.pipe(
    e2s2sIdxs,
    A.flatMap(([s1, s2]) => {
      const s1RefInput = match(s1)
        .with({ snapshotUtxo: P.select() }, (utxo) => O.some(utxo))
        .otherwise(() => O.none);

      const s2RefInput = F.pipe(
        s2,
        O.match(
          () => [],
          (s) =>
            match(s)
              .with({ snapshotUtxo: P.select() }, (utxo) =>
                F.pipe(
                  s1RefInput,
                  O.match(
                    () => [utxo],
                    // when s1 ref input is same as s2 ref input, don't reference it again
                    (s1In) => (isSameOutRef(s1In, utxo) ? [] : [utxo]),
                  ),
                ),
              )
              .otherwise(() => []),
        ),
      );

      return [...F.pipe([s1RefInput], A.compact), ...s2RefInput];
    }),
  );

  const e2s2sInputsIndices = getInputIndices(allE2s2sSnapshotRefInputs, [
    ...otherRefInputs,
    ...allE2s2sSnapshotRefInputs,
  ]);

  const createE2s2sIdx = (s: FindE2S2SIdxResult): E2S2SIndex =>
    match(s)
      .returnType<E2S2SIndex>()
      .with({ spListIdx: P.select() }, (spListIdx) => ({
        StabilityPoolListIdx: BigInt(spListIdx),
      }))
      .with({ snapshotUtxo: P.any }, (obj) => {
        const idxForIdx = F.pipe(
          allE2s2sSnapshotRefInputs,
          A.findIndex((input) => isSameOutRef(input, obj.snapshotUtxo)),
          O.getOrElse<number>(() => {
            throw new Error('Expected to find the index.');
          }),
        );

        return {
          RefInputIdx: {
            refInputIdx: e2s2sInputsIndices[idxForIdx],
            snapshotListIdx: BigInt(obj.snapshotListIdx),
          },
        };
      })
      .exhaustive();

  return {
    e2s2sRefInputs: allE2s2sSnapshotRefInputs,
    mkProcessRequestAccountRedeemerContent: (
      poolInputIdx: bigint,
      accountInputIdx: bigint,
    ) => ({
      poolInputIdx: poolInputIdx,
      accountInputIdx: accountInputIdx,
      e2s2sIdxs: F.pipe(
        e2s2sIdxs,
        A.reduce<
          [FindE2S2SIdxResult, O.Option<FindE2S2SIdxResult>],
          E2S2SIndicesPerAsset
        >([], (acc, [s1, s2]) => {
          return [
            ...acc,
            [
              createE2s2sIdx(s1),
              F.pipe(
                s2,
                O.match(
                  () => null,
                  (s) => createE2s2sIdx(s),
                ),
              ),
            ],
          ];
        }),
      ),
      currentTime: currentTime,
    }),
  };
}

function calculateReward(
  s1: SPInteger,
  s2: SPInteger,
  accountSumVal: SPInteger,
  accountState: StateSnapshot,
): bigint {
  const a1 = spSub(s1, accountSumVal);
  const a2 = spDiv(spSub(s2, s1), mkSPInteger(newScaleMultiplier));

  return F.pipe(
    spDiv(
      spMul(spAdd(a1, a2), accountState.depositVal),
      accountState.productVal,
    ),
    fromSPInteger,
  );
}

function getE2s2sEntry(
  idx: FindE2S2SIdxResult,
  assetState: AssetSnapshot,
): EpochToScaleToSumEntry {
  return match(idx)
    .returnType<EpochToScaleToSumEntry>()
    .with(
      { spListIdx: P.select() },
      (spIdx) => assetState.epoch2scale2sum[spIdx],
    )
    .with(
      { snapshotUtxo: P.any },
      (obj) => obj.snapshotDatum.snapshot[obj.snapshotListIdx],
    )
    .exhaustive();
}

function rewardsPerAsset(
  poolAssetStates: readonly AssetState[],
  e2s2sIdxs: [FindE2S2SIdxResult, O.Option<FindE2S2SIdxResult>][],
  accountAssetSums: readonly (readonly [AssetClass, SPInteger])[],
  accountState: StateSnapshot,
): [AssetClass, bigint][] {
  const expectedS1Key: EpochToScaleKey = {
    epoch: accountState.epoch,
    scale: accountState.scale,
  };

  return F.pipe(
    RA.zip(e2s2sIdxs)(poolAssetStates),
    RA.reduce<
      readonly [AssetState, [FindE2S2SIdxResult, O.Option<FindE2S2SIdxResult>]],
      [AssetClass, bigint][]
    >([], (acc, [[poolAsset, poolAssetState], [s1Idx, s2Idx]]) => {
      const s1Res = getE2s2sEntry(s1Idx, poolAssetState);

      const s1 = isSameEpochToScaleKey(s1Res[0], expectedS1Key)
        ? s1Res[1].sumVal
        : initSumVal;

      const s2 = F.pipe(
        s2Idx,
        O.match(
          () => s1,
          (s2Res) => s2Res.sumSnapshot[1].sumVal,
        ),
      );

      const reward = calculateReward(
        s1,
        s2,
        F.pipe(
          accountAssetSums,
          RA.findFirst(([accountAsset, _]) =>
            isSameAssetClass(accountAsset, poolAsset),
          ),
          O.match(
            // When account doesn't have this asset in snapshots
            () => initSumVal,
            (accountSumVal) => accountSumVal[1],
          ),
        ),
        accountState,
      );

      return [[poolAsset, reward], ...acc];
    }),
  );
}

export function getUpdatedAccountDeposit(
  poolState: StateSnapshot,
  accountState: StateSnapshot,
): SPInteger {
  if (poolState.epoch > accountState.epoch) {
    return mkSPInteger(0n);
  } else if (poolState.scale - accountState.scale > 1) {
    return mkSPInteger(0n);
  } else if (poolState.scale > accountState.scale) {
    return spDiv(spMul(accountState.depositVal, poolState.productVal), {
      value: accountState.productVal.value * newScaleMultiplier,
    });
  } else {
    return spDiv(
      spMul(accountState.depositVal, poolState.productVal),
      accountState.productVal,
    );
  }
}

export function updateAccount(
  pool: StabilityPoolContent,
  account: AccountContent,
  e2s2sIdxs: [FindE2S2SIdxResult, O.Option<FindE2S2SIdxResult>][],
): {
  updatedAccountContent: AccountContent;
  reward: Assets;
} {
  const accountState = account.state;
  const fund = getUpdatedAccountDeposit(pool.state, accountState);

  const rewards = rewardsPerAsset(
    pool.assetStates,
    e2s2sIdxs,
    account.assetSums,
    accountState,
  );

  const newDepositVal =
    fund.value <
    spDiv(accountState.depositVal, mkSPInteger(1_000_000_000n)).value
      ? mkSPInteger(0n)
      : fund;

  return {
    updatedAccountContent: {
      ...account,
      state: { ...pool.state, depositVal: newDepositVal },
      assetSums: F.pipe(
        pool.assetStates,
        RA.map(([key, assetState]) => [key, assetState.currentSumVal]),
      ),
    },
    reward: F.pipe(
      rewards,
      A.reduce<[AssetClass, bigint], Assets>({}, (acc, [asset, amt]) =>
        addAssets(acc, mkAssetsOf(asset, amt)),
      ),
    ),
  };
}
export function liquidationHelper(
  poolContent: StabilityPoolContent,
  collateralAsset: AssetClass,
  iassetBurnAmt: bigint,
  /**
   * The collateral absorbed
   */
  reward: bigint,
): StabilityPoolContent {
  const lossPerUnitStaked = spDiv(
    mkSPInteger(iassetBurnAmt),
    poolContent.state.depositVal,
  );

  const productFactor = spSub(mkSPInteger(1n), lossPerUnitStaked);

  const isScaleIncrease =
    spMul(poolContent.state.productVal, productFactor).value <
    newScaleMultiplier;

  const newSnapshotP = spMul(
    {
      value:
        poolContent.state.productVal.value *
        (isScaleIncrease ? newScaleMultiplier : 1n),
    },
    productFactor,
  );

  const currentS = F.pipe(
    poolContent.assetStates,
    RA.findFirstMap(([ac, assetSnap]) =>
      isSameAssetClass(ac, collateralAsset)
        ? O.some(assetSnap.currentSumVal)
        : O.none,
    ),
    O.getOrElse(() => initSumVal),
  );

  const newSnapshotS = spAdd(
    currentS,
    spDiv(
      spMul(mkSPInteger(reward), poolContent.state.productVal),
      poolContent.state.depositVal,
    ),
  );

  const isEpochIncrease = newSnapshotP.value <= 0;

  const newState: StateSnapshot = isEpochIncrease
    ? { ...initSpState, epoch: poolContent.state.epoch + 1n }
    : {
        productVal: newSnapshotP,
        depositVal: spSub(
          poolContent.state.depositVal,
          mkSPInteger(iassetBurnAmt),
        ),
        epoch: poolContent.state.epoch,
        scale: poolContent.state.scale + (isScaleIncrease ? 1n : 0n),
      };

  const currentE2S2SKey: EpochToScaleKey = {
    epoch: poolContent.state.epoch,
    scale: poolContent.state.scale,
  };

  const newCollateralAssetSumVal = isEpochIncrease ? initSumVal : newSnapshotS;

  const newAssetStates = (() => {
    const updatedAssetStates = F.pipe(
      poolContent.assetStates,
      repsertWithReadonlyArr<AssetClass, AssetSnapshot>(
        (key) => isSameAssetClass(key, collateralAsset),
        (assetSnap) => ({
          currentSumVal: newCollateralAssetSumVal,
          epoch2scale2sum: F.pipe(
            assetSnap.epoch2scale2sum,
            RA.modifyAt(0, ([key, val]) => [
              key,
              { ...val, sumVal: newSnapshotS } satisfies SumSnapshot,
            ]),
            O.getOrElse<readonly EpochToScaleToSumEntry[]>(() => {
              throw new Error('There has to be first entry');
            }),
          ),
        }),
        () => ({
          currentSumVal: newCollateralAssetSumVal,
          epoch2scale2sum: [
            [
              currentE2S2SKey,
              {
                sumVal: newSnapshotS,
                isFirstSnapshot: true,
                isLastInEpoch: true,
              },
            ],
          ],
        }),
        () => collateralAsset,
      ),
    );

    if (isEpochIncrease) {
      return F.pipe(
        updatedAssetStates,
        RA.map<AssetState, AssetState>(([key, assetSnap]) => [
          key,
          {
            currentSumVal: initSumVal,
            epoch2scale2sum: [
              [
                {
                  epoch: poolContent.state.epoch + 1n,
                  scale: 0n,
                },
                {
                  sumVal: initSumVal,
                  isLastInEpoch: true,
                  isFirstSnapshot: false,
                },
              ],
              ...assetSnap.epoch2scale2sum,
            ],
          },
        ]),
      );
    } else if (isScaleIncrease) {
      return F.pipe(
        updatedAssetStates,
        RA.map<AssetState, AssetState>(([key, assetSnap]) => [
          key,
          {
            ...assetSnap,
            epoch2scale2sum: match(assetSnap.epoch2scale2sum)
              .returnType<readonly EpochToScaleToSumEntry[]>()
              .with([[P.any, P.any], ...P.array()], ([[key, val], ...rest]) => {
                const newScaleEntry: EpochToScaleToSumEntry = [
                  {
                    epoch: poolContent.state.epoch,
                    scale: poolContent.state.scale + 1n,
                  },
                  {
                    sumVal: val.sumVal,
                    isLastInEpoch: true,
                    isFirstSnapshot: false,
                  },
                ];

                return [
                  newScaleEntry,
                  [key, { ...val, isLastInEpoch: false } satisfies SumSnapshot],
                  ...rest,
                ];
              })
              .otherwise(() => {
                throw new Error('There has to be at least 1 entry');
              }),
          },
        ]),
      );
    } else {
      return updatedAssetStates;
    }
  })();

  return {
    ...poolContent,
    assetStates: newAssetStates,
    state: newState,
  };
}

export function updatePoolStateWhenWithdrawalFee(
  withdrawalFeeAmt: bigint,
  updatedPoolState: StateSnapshot,
): StateSnapshot {
  if (withdrawalFeeAmt === 0n) {
    return updatedPoolState;
  } else {
    const withdrawalFeeSpInt = mkSPInteger(withdrawalFeeAmt);
    const newDepositVal = spAdd(
      updatedPoolState.depositVal,
      withdrawalFeeSpInt,
    );

    const productFactor = spAdd(
      mkSPInteger(1n),
      spDiv(withdrawalFeeSpInt, updatedPoolState.depositVal),
    );

    const newProductVal = spMul(updatedPoolState.productVal, productFactor);

    return {
      ...updatedPoolState,
      productVal: newProductVal,
      depositVal: newDepositVal,
    };
  }
}

export function partitionEpochToScaleToSums(
  spContent: StabilityPoolContent,
): readonly [
  readonly SnapshotEpochToScaleToSumContent[],
  readonly AssetState[],
] {
  const res = F.pipe(
    spContent.assetStates,
    RA.map<AssetState, [SnapshotEpochToScaleToSumContent[], AssetState]>(
      ([collateralAsset, assetState]) => {
        if (assetState.epoch2scale2sum.length >= MAX_E2S2S_ENTRIES_COUNT) {
          const { right: remaining, left: snapshotMapItems } = F.pipe(
            assetState.epoch2scale2sum,
            RA.partition(
              ([e2sKey, _]) =>
                e2sKey.epoch === spContent.state.epoch &&
                e2sKey.scale === spContent.state.scale,
            ),
          );

          return [
            [
              {
                iasset: spContent.iasset,
                collateralAsset: collateralAsset,
                snapshot: snapshotMapItems,
              },
            ],
            [collateralAsset, { ...assetState, epoch2scale2sum: remaining }],
          ];
        } else {
          return [[], [collateralAsset, assetState]];
        }
      },
    ),
  );

  const [newSnapshots, newAssetStates] = RA.unzip(res);

  return [RA.flatten(newSnapshots), newAssetStates];
}