import {
  AccountPoolReward,
  AccumulatedReward,
  AccumulatedPoolReward,
  HistoryElement,
  Reward,
  RewardType,
} from "../types";
import { SubstrateEvent } from "@subql/types";
import {
  eventId,
  eventIdFromBlockAndIdxAndAddress,
  timestamp,
  eventIdWithAddress,
  blockNumber,
} from "./common";
import {
  handleGenericForTxHistory,
  updateAccumulatedGenericReward,
} from "./Rewards";
import { getPoolMembers } from "./Cache";
import { Option } from "@pezkuwi/types";
import {
  PEZKUWI_RELAY_GENESIS,
  PEZKUWI_ASSET_HUB_GENESIS,
  STAKING_TYPE_NOMINATION_POOL,
} from "./constants";

export async function handlePoolReward(
  rewardEvent: SubstrateEvent,
): Promise<void> {
  await handlePoolRewardForTxHistory(rewardEvent);
  let accumulatedReward = await updateAccumulatedPoolReward(rewardEvent, true);
  let {
    event: {
      data: [accountId, poolId, amount],
    },
  } = rewardEvent;
  await updateAccountPoolRewards(
    rewardEvent,
    accountId.toString(),
    (amount as any).toBigInt(),
    (poolId as any).toNumber(),
    RewardType.reward,
    accumulatedReward.amount,
  );
  // Save to multi-staking Reward entity for both relay and Asset Hub networkIds
  await savePoolMultiStakingReward(
    rewardEvent,
    accountId.toString(),
    (amount as any).toBigInt(),
    RewardType.reward,
  );
}

async function handlePoolRewardForTxHistory(
  rewardEvent: SubstrateEvent,
): Promise<void> {
  const {
    event: {
      data: [account, poolId, amount],
    },
  } = rewardEvent;
  handleGenericForTxHistory(
    rewardEvent,
    account.toString(),
    async (element: HistoryElement) => {
      element.poolReward = {
        eventIdx: rewardEvent.idx,
        amount: amount.toString(),
        isReward: true,
        poolId: (poolId as any).toNumber(),
      };
      return element;
    },
  );
}

async function updateAccumulatedPoolReward(
  event: SubstrateEvent,
  isReward: boolean,
): Promise<AccumulatedReward> {
  let {
    event: {
      data: [accountId, _, amount],
    },
  } = event;
  return await updateAccumulatedGenericReward(
    AccumulatedPoolReward,
    accountId.toString(),
    (amount as any).toBigInt(),
    isReward,
  );
}

async function updateAccountPoolRewards(
  event: SubstrateEvent,
  accountAddress: string,
  amount: bigint,
  poolId: number,
  rewardType: RewardType,
  accumulatedAmount: bigint,
): Promise<void> {
  let id = eventIdWithAddress(event, accountAddress);
  let accountPoolReward = new AccountPoolReward(
    id,
    accountAddress,
    blockNumber(event),
    timestamp(event.block),
    amount,
    accumulatedAmount,
    rewardType,
    poolId,
  );
  await accountPoolReward.save();
}

export async function handlePoolBondedSlash(
  bondedSlashEvent: SubstrateEvent,
): Promise<void> {
  const {
    event: {
      data: [poolIdEncoded, slash],
    },
  } = bondedSlashEvent;
  const poolId = (poolIdEncoded as any).toNumber();

  const poolOption = (await api.query.nominationPools.bondedPools(
    poolId,
  )) as Option<any>;
  const pool = poolOption.unwrap();

  await handleRelaychainPooledStakingSlash(
    bondedSlashEvent,
    poolId,
    pool.points.toBigInt(),
    (slash as any).toBigInt(),
    (member: any): bigint => {
      return member.points.toBigInt();
    },
  );
}

export async function handlePoolUnbondingSlash(
  unbondingSlashEvent: SubstrateEvent,
): Promise<void> {
  const {
    event: {
      data: [poolId, era, slash],
    },
  } = unbondingSlashEvent;
  const poolIdNumber = (poolId as any).toNumber();
  const eraIdNumber = (era as any).toNumber();

  const unbondingPools = (
    (await api.query.nominationPools.subPoolsStorage(
      poolIdNumber,
    )) as Option<any>
  ).unwrap();

  const pool =
    (unbondingPools.withEra as any).get(eraIdNumber) ?? unbondingPools.noEra;

  await handleRelaychainPooledStakingSlash(
    unbondingSlashEvent,
    poolIdNumber,
    pool.points.toBigInt(),
    (slash as any).toBigInt(),
    (member: any): bigint => {
      return (
        (member.unbondingEras as any).get(eraIdNumber)?.toBigInt() ?? BigInt(0)
      );
    },
  );
}

async function handleRelaychainPooledStakingSlash(
  event: SubstrateEvent,
  poolId: number,
  poolPoints: bigint,
  slash: bigint,
  memberPointsCounter: (member: any) => bigint,
): Promise<void> {
  if (poolPoints == BigInt(0)) {
    return;
  }

  const members = await getPoolMembers(blockNumber(event));

  for (const [accountId, member] of members) {
    let memberPoints: bigint;
    if (member.poolId.toNumber() === poolId) {
      memberPoints = memberPointsCounter(member);
      if (memberPoints != BigInt(0)) {
        const personalSlash = (slash * memberPoints) / poolPoints;

        await handlePoolSlashForTxHistory(
          event,
          poolId,
          accountId,
          personalSlash,
        );
        let accumulatedReward = await updateAccumulatedGenericReward(
          AccumulatedPoolReward,
          accountId,
          personalSlash,
          false,
        );
        await updateAccountPoolRewards(
          event,
          accountId,
          personalSlash,
          poolId,
          RewardType.slash,
          accumulatedReward.amount,
        );
        // Save to multi-staking Reward entity
        await savePoolMultiStakingReward(
          event,
          accountId,
          personalSlash,
          RewardType.slash,
        );
      }
    }
  }
}

async function handlePoolSlashForTxHistory(
  slashEvent: SubstrateEvent,
  poolId: number,
  accountId: string,
  personalSlash: bigint,
): Promise<void> {
  const extrinsic = slashEvent.extrinsic;
  const block = slashEvent.block;
  const blockNum = block.block.header.number.toString();
  const blockTimestamp = timestamp(block);
  const evtId = eventIdFromBlockAndIdxAndAddress(
    blockNum,
    slashEvent.idx.toString(),
    accountId,
  );

  const element = HistoryElement.create({
    id: evtId,
    timestamp: blockTimestamp,
    blockNumber: block.block.header.number.toNumber(),
    extrinsicHash:
      extrinsic !== undefined ? extrinsic.extrinsic.hash.toString() : null,
    extrinsicIdx: extrinsic !== undefined ? extrinsic.idx : null,
    address: accountId,
    poolReward: {
      eventIdx: slashEvent.idx,
      amount: personalSlash.toString(),
      isReward: false,
      poolId: poolId,
    },
  });

  await element.save();
}

/**
 * Save pool reward/slash to the multi-staking Reward entity
 * Saves for both relay chain and Asset Hub networkIds since nom pools
 * are accessible from both contexts
 */
async function savePoolMultiStakingReward(
  event: SubstrateEvent,
  accountAddress: string,
  amount: bigint,
  rewardType: RewardType,
): Promise<void> {
  const ts = timestamp(event.block);
  const bn = blockNumber(event);
  const baseId = `${eventId(event)}-${accountAddress}-pool`;

  // Save for Asset Hub networkId (nomination pools are accessed via Asset Hub)
  const ahReward = Reward.create({
    id: `${baseId}-ah`,
    networkId: PEZKUWI_ASSET_HUB_GENESIS,
    stakingType: STAKING_TYPE_NOMINATION_POOL,
    address: accountAddress,
    type: rewardType,
    amount,
    timestamp: ts,
    blockNumber: bn,
  });
  await ahReward.save();
}