import Decimal from 'decimal.js';
import { Address } from '@solana/kit';
import { lamportsToCollDecimal } from '@kamino-finance/farms-sdk';

export interface ReserveAllocationOverview {
  targetWeight: Decimal;
  tokenAllocationCap: Decimal;
  ctokenAllocation: Decimal;
}

export interface VaultAllocationResult {
  targetUnallocatedAmount: Decimal;
  targetReservesAllocation: Map<Address, Decimal>;
}

const ZERO = new Decimal(0);

/**
 * Computes the allocation of vault funds across reserves based on weights and caps.
 * @param vaultAUMInvestableInReserves - Total AUM of the vault, in tokens, that can be invested in reserves.
 * @param vaultUnallocatedWeight - Weight for unallocated funds
 * @param vaultUnallocatedCap - Maximum amount that can remain unallocated
 * @param initialVaultAllocations - Map of reserve addresses to their allocation configurations
 * @param vaultTokenDecimals - The number of decimals of the vault token, needed to compute the min amount
 * @returns Object containing target unallocated amount and target allocations per reserve, in tokens
 */
export function computeReservesAllocation(
  vaultAUMInvestableInReserves: Decimal,
  vaultUnallocatedWeight: Decimal,
  vaultUnallocatedCap: Decimal,
  initialVaultAllocations: Map<Address, ReserveAllocationOverview>,
  vaultTokenDecimals: number
): VaultAllocationResult {
  let totalAllocation = new Decimal(0);
  const reserves = Array.from(initialVaultAllocations.keys());
  const expectedHoldingsDistribution = new Map<Address, Decimal>();

  reserves.forEach((reserve) => {
    expectedHoldingsDistribution.set(reserve, ZERO);
    totalAllocation = totalAllocation.add(initialVaultAllocations.get(reserve)!.targetWeight);
  });

  let totalLeftToInvest = vaultAUMInvestableInReserves;
  const totalAllocationsIncludingUnallocated = totalAllocation.add(vaultUnallocatedWeight);
  if (totalAllocationsIncludingUnallocated.lte(ZERO)) {
    return {
      targetUnallocatedAmount: totalLeftToInvest,
      targetReservesAllocation: expectedHoldingsDistribution,
    };
  }

  // Calculate initial unallocated amount
  let unallocatedAllocation = totalLeftToInvest.mul(vaultUnallocatedWeight).div(totalAllocationsIncludingUnallocated);
  if (unallocatedAllocation.gt(vaultUnallocatedCap)) {
    unallocatedAllocation = vaultUnallocatedCap;
  }
  totalLeftToInvest = totalLeftToInvest.sub(unallocatedAllocation);

  let currentAllocationSum = totalAllocation;

  const reservesCount = reserves.length;
  const maxRemainedUninvestedLamports = lamportsToCollDecimal(new Decimal(reservesCount), vaultTokenDecimals); // invest only if the AUM has more lamports than the number of reserves

  // Allocate in proportional rounds. Reserves that hit a cap are removed from
  // the next round and their unused share is redistributed across the rest.
  while (totalLeftToInvest.gt(maxRemainedUninvestedLamports) && currentAllocationSum.gt(ZERO)) {
    const activeReserves = reserves.filter((reserve) => {
      const allocation = initialVaultAllocations.get(reserve);
      if (!allocation || allocation.targetWeight.lte(ZERO) || allocation.tokenAllocationCap.lte(ZERO)) {
        return false;
      }

      return expectedHoldingsDistribution.get(reserve)!.lt(allocation.tokenAllocationCap);
    });

    currentAllocationSum = activeReserves.reduce(
      (sum, reserve) => sum.add(initialVaultAllocations.get(reserve)!.targetWeight),
      ZERO
    );

    if (currentAllocationSum.lte(ZERO)) {
      break;
    }

    const totalLeftover = totalLeftToInvest;
    let allocatedThisRound = ZERO;

    for (const reserve of activeReserves) {
      const reserveWithWeight = initialVaultAllocations.get(reserve)!;
      const targetAllocation = reserveWithWeight.targetWeight.mul(totalLeftover).div(currentAllocationSum);
      const currentReserveAllocation = expectedHoldingsDistribution.get(reserve)!;
      const remainingCapacity = reserveWithWeight.tokenAllocationCap.sub(currentReserveAllocation);
      const amountToInvest = Decimal.min(targetAllocation, remainingCapacity);

      allocatedThisRound = allocatedThisRound.add(amountToInvest);
      expectedHoldingsDistribution.set(reserve, currentReserveAllocation.add(amountToInvest));
    }

    if (allocatedThisRound.lte(ZERO)) {
      break;
    }

    totalLeftToInvest = totalLeftToInvest.sub(allocatedThisRound);
  }

  return {
    targetUnallocatedAmount: unallocatedAllocation.add(totalLeftToInvest),
    targetReservesAllocation: expectedHoldingsDistribution,
  };
}
