import { UtxoMetaData } from '@saturnbtcio/arch-sdk';
import {
  IdentifiableLiquidityPool,
  IdentifiableLiquidityPoolShard,
  LiquidityPoolShard,
} from '@saturnbtcio/pool-serde-sdk';
import { PoolErrorException, PoolErrorType } from '../error/pool.error';
import { getMempoolEntryFromShard } from './mempool';
import { getTxIdFromShard } from './mempool';
import { MempoolInfoMap } from '../providers/bitcoin.provider';
import { DUST_LIMIT } from './constants';

export enum UpdateLiquidityBy {
  Liquidity,
  BtcAmount,
  RuneAmount,
}

export function getShardsKey(
  shard: LiquidityPoolShard,
  modifyBy: UpdateLiquidityBy,
): bigint {
  switch (modifyBy) {
    case UpdateLiquidityBy.Liquidity:
      return shard.liquidity;
    case UpdateLiquidityBy.BtcAmount:
      return BigInt(shard.btcUtxos.reduce((sum, utxo) => sum + utxo.value, 0n));
    case UpdateLiquidityBy.RuneAmount:
      return shard.runeUtxo?.runes[0]?.amount ?? BigInt(0);
  }
}

/**
 * Walk the ancestor chain of the shard's controlling UTXO (using the depends[]
 * links returned by Bitcoin Core) and compute the *smallest* remaining
 * descendant capacity among all ancestors.  That is:
 *     min( 25 - descendantsCount )
 * for every unconfirmed ancestor in the chain.  The smaller this number, the
 * closer the shard is to the 25-descendant cliff.
 */
function getRemainingDescendantCapacity(
  shard: IdentifiableLiquidityPoolShard,
  mempoolInfo: MempoolInfoMap,
): number {
  const queue: string[] = [getTxIdFromShard(shard)];
  const visited: Set<string> = new Set();

  let minRemaining = 25; // maximum possible head-room

  while (queue.length) {
    const txid = queue.pop() as string;
    if (visited.has(txid)) continue;
    visited.add(txid);

    const entry = mempoolInfo.get(txid);
    if (!entry) {
      // confirmed ancestor → no mempool limits apply further up this branch.
      continue;
    }

    const remaining = 25 - entry.descendantsCount;
    if (remaining < minRemaining) minRemaining = remaining;

    // traverse further ancestors
    for (const parent of entry.depends ?? []) {
      if (!visited.has(parent)) queue.push(parent);
    }
  }

  return minRemaining;
}

export function sortShardsByPriority(
  shards: IdentifiableLiquidityPoolShard[],
  updateBy: UpdateLiquidityBy,
  mempoolInfo: MempoolInfoMap,
): void {
  shards.sort((a, b) => {
    // Then, prioritize shards with fewer descendants
    // Default limit of descendants is 25.
    const aMempoolInfo = getMempoolEntryFromShard(a, mempoolInfo);
    const bMempoolInfo = getMempoolEntryFromShard(b, mempoolInfo);
    if (aMempoolInfo && bMempoolInfo) {
      if (aMempoolInfo.descendantsCount !== bMempoolInfo.descendantsCount) {
        return Number(
          aMempoolInfo.descendantsCount - bMempoolInfo.descendantsCount,
        );
      }
    }

    // Then, prioritize shards with fewer ancestors size
    // Default limit of ancestors is 100kb - now we are using count
    if (aMempoolInfo && bMempoolInfo) {
      if (aMempoolInfo.ancestorsCount !== bMempoolInfo.ancestorsCount) {
        return Number(
          aMempoolInfo.ancestorsCount - bMempoolInfo.ancestorsCount,
        );
      }
    }

    // Prefer shards whose ancestor chain still has more head-room before
    // hitting the 25 descendant limit.  We compute the minimum remaining
    // capacity across all ancestors; the higher the value, the safer.
    const aRemaining = getRemainingDescendantCapacity(a, mempoolInfo);
    const bRemaining = getRemainingDescendantCapacity(b, mempoolInfo);

    if (aRemaining !== bRemaining) {
      // larger remaining capacity first
      return bRemaining - aRemaining;
    }

    // Finally, sort by size (biggest first)
    const keyA = getShardsKey(a, updateBy);
    const keyB = getShardsKey(b, updateBy);
    return Number(keyB - keyA);
  });
}

export function selectShardsToRemoveMultipleFrom(
  pool: IdentifiableLiquidityPool,
  amountToRemove: Map<UpdateLiquidityBy, bigint>,
  shardsMempoolInfo: MempoolInfoMap,
): IdentifiableLiquidityPoolShard[] {
  const selectedShards = new Map<string, Map<UpdateLiquidityBy, bigint>>();
  const remainingAmounts = new Map(amountToRemove);

  console.log('amountToRemove', amountToRemove);

  for (const [removeBy, totalAmount] of amountToRemove.entries()) {
    let amountLeft = totalAmount;

    console.log('Selecting best shards for: ', removeBy, totalAmount);

    const shardsToRemoveFrom = selectBestShardsToRemoveFrom(
      pool,
      totalAmount,
      removeBy,
      shardsMempoolInfo,
    );

    console.log('Selected shards to remove from: ', shardsToRemoveFrom);

    for (const shard of shardsToRemoveFrom) {
      const shardKey = shard.pubkey;

      // Skip already selected shards.
      if (selectedShards.has(shardKey)) {
        continue;
      }

      const availableAmount = getShardsKey(shard, removeBy);
      const remainingAmount = remainingAmounts.get(removeBy)!;
      const amountToRemoveFromShard =
        availableAmount < remainingAmount ? availableAmount : remainingAmount;

      if (!selectedShards.has(shardKey)) {
        selectedShards.set(shardKey, new Map<UpdateLiquidityBy, bigint>());
      }

      selectedShards.get(shardKey)!.set(removeBy, amountToRemoveFromShard);

      remainingAmounts.set(removeBy, remainingAmount - amountToRemoveFromShard);

      // Update other types for this shard
      for (const [otherType, otherAmount] of remainingAmounts.entries()) {
        if (otherType !== removeBy) {
          const otherAvailable = getShardsKey(shard, otherType);
          const otherToRemove =
            otherAvailable < otherAmount ? otherAvailable : otherAmount;
          remainingAmounts.set(otherType, otherAmount - otherToRemove);

          selectedShards.get(shardKey)!.set(otherType, otherToRemove);
        }
      }
    }
  }

  // Double check if we need all of the shards or can remove some
  let finalShards: IdentifiableLiquidityPoolShard[] = [];
  const totalRemoved = new Map<UpdateLiquidityBy, bigint>();

  for (const [shardKey, amounts] of selectedShards.entries()) {
    let keepShard = false;

    for (const [removeBy, amount] of amounts.entries()) {
      const needed = amountToRemove.get(removeBy)!;
      const removed = totalRemoved.get(removeBy) || BigInt(0);

      if (removed < needed) {
        const toRemove = amount < needed - removed ? amount : needed - removed;
        totalRemoved.set(removeBy, removed + toRemove);

        if (toRemove > BigInt(0)) {
          keepShard = true;
        }
      }
    }

    if (keepShard) {
      const shard = pool.shards.find((s) => s.pubkey === shardKey)!;
      finalShards.push(shard);
    }
  }

  // Check if we've removed enough for all types
  for (const [removeBy, needed] of amountToRemove.entries()) {
    if ((totalRemoved.get(removeBy) || BigInt(0)) < needed) {
      throw new PoolErrorException({
        message: `Not enough liquidity to remove for ${removeBy}`,
        type: PoolErrorType.InsufficientLiquidity,
        token: removeBy.toString(),
        maxAmount: (totalRemoved.get(removeBy) || BigInt(0)).toString(),
      });
    }
  }

  // Enforce BTC dust remainder on the final selection when BTC is being removed
  const btcNeeded = amountToRemove.get(UpdateLiquidityBy.BtcAmount) || 0n;
  if (btcNeeded > 0n) {
    finalShards = enforceBtcDustRemainder(
      pool,
      finalShards,
      btcNeeded,
      shardsMempoolInfo,
    );
  }

  return finalShards;
}

export function validateAmountsAreWithinThreshold(
  amounts: bigint[],
  thresholdPercent: bigint, // 100n = 100%
): void {
  if (amounts.length === 0) {
    return;
  }

  const totalAmount = amounts.reduce((sum, amount) => sum + amount, 0n);
  const numAmounts = BigInt(amounts.length);

  const idealAmount = totalAmount / numAmounts;

  // Calculate the acceptable deviation as a percentage of the ideal amount
  const acceptableDeviation = (thresholdPercent * idealAmount + 99n) / 100n;

  for (const amount of amounts) {
    // Compute the absolute difference between the amount and the ideal amount
    const diff =
      amount > idealAmount ? amount - idealAmount : idealAmount - amount;

    if (diff > acceptableDeviation) {
      throw new PoolErrorException({
        message: `Amounts are not within threshold`,
        type: PoolErrorType.InvalidPsbt,
      });
    }
  }
}

export function selectBestShardsToRemoveFrom(
  pool: IdentifiableLiquidityPool,
  amountToRemove: bigint,
  removeBy: UpdateLiquidityBy,
  mempoolInfo: MempoolInfoMap,
): IdentifiableLiquidityPoolShard[] {
  const shardsToRemoveFrom: number[] = [];
  let amountRemoved = BigInt(0);

  // First, gather shards with positive capacity for the requested dimension
  const positiveShards = pool.shards
    .filter((shard) => getShardsKey(shard, removeBy) > 0n)
    .slice();

  if (positiveShards.length === 0) {
    throw new PoolErrorException({
      message: `No eligible shards available for ${
        removeBy === UpdateLiquidityBy.RuneAmount
          ? pool.config.token0
          : pool.config.token1
      }`,
      type: PoolErrorType.ShardsUnavailable,
      reason: 'NoBalance',
      token:
        removeBy === UpdateLiquidityBy.RuneAmount
          ? `${pool.config.token0}`
          : `${pool.config.token1}`,
    });
  }

  // Then apply mempool safety filters
  // Sort shards by priority (least updated first, then by size)
  const filteredShards = positiveShards
    // Skip shards whose controlling UTXO is already at or above the mempool limits so we
    // don't create an "too-long-mempool-chain" rejection.
    .filter((shard, i) => {
      const entry = getMempoolEntryFromShard(shard, mempoolInfo);

      // If there is no mempool entry the UTXO is confirmed (safe to use).
      if (!entry) return true;

      const descendantCapacity = getRemainingDescendantCapacity(
        shard,
        mempoolInfo,
      );

      if (descendantCapacity === 0) {
        return false;
      }

      return entry.ancestorsCount < 23 && entry.descendantsCount < 23;
    })
    .slice();

  // If all positive-balance shards were filtered out by mempool constraints, signal that.
  if (filteredShards.length === 0) {
    throw new PoolErrorException({
      message: `No eligible shards available due to mempool constraints for ${
        removeBy === UpdateLiquidityBy.RuneAmount
          ? pool.config.token0
          : pool.config.token1
      }`,
      type: PoolErrorType.ShardsUnavailable,
      reason: 'MempoolConstraints',
      token:
        removeBy === UpdateLiquidityBy.RuneAmount
          ? `${pool.config.token0}`
          : `${pool.config.token1}`,
    });
  }

  sortShardsByPriority(filteredShards, removeBy, mempoolInfo);

  for (let i = 0; i < filteredShards.length; i++) {
    const shard = filteredShards[i];

    const key = getShardsKey(shard, removeBy);
    const amountToRemoveFromShard =
      key < amountToRemove - amountRemoved
        ? key
        : amountToRemove - amountRemoved;

    if (amountToRemoveFromShard > BigInt(0)) {
      shardsToRemoveFrom.push(i);
      amountRemoved += amountToRemoveFromShard;
    }

    if (amountRemoved >= amountToRemove) {
      break;
    }
  }

  if (amountRemoved < amountToRemove) {
    throw new PoolErrorException({
      message: `Pool ${pool.config.token0} ${pool.config.token1} does not have liquidity`,
      type: PoolErrorType.InsufficientLiquidity,
      token:
        removeBy === UpdateLiquidityBy.RuneAmount
          ? `${pool.config.token0}`
          : `${pool.config.token1}`,
      maxAmount: amountRemoved.toString(),
    });
  }

  // Now, from the selected shards, we sort by amount (biggest first)
  shardsToRemoveFrom.sort((a, b) => {
    const keyA = getShardsKey(filteredShards[a], removeBy);
    const keyB = getShardsKey(filteredShards[b], removeBy);
    return Number(keyB - keyA);
  });

  // Check if we need the last shards at the end.
  const lastShardsToRemoveFrom: IdentifiableLiquidityPoolShard[] = [];
  amountRemoved = BigInt(0);

  for (const index of shardsToRemoveFrom) {
    const shard = filteredShards[index];
    const key = getShardsKey(shard, removeBy);
    const amountToRemoveFromShard =
      key < amountToRemove - amountRemoved
        ? key
        : amountToRemove - amountRemoved;
    amountRemoved += amountToRemoveFromShard;

    if (amountToRemoveFromShard > BigInt(0)) {
      lastShardsToRemoveFrom.push(shard);
    }

    if (amountRemoved >= amountToRemove) {
      break;
    }
  }

  // For BTC-only selection, ensure the post-removal remainder is either 0 or >= dust.
  if (removeBy === UpdateLiquidityBy.BtcAmount) {
    return enforceBtcDustRemainder(
      pool,
      lastShardsToRemoveFrom,
      amountToRemove,
      mempoolInfo,
    );
  }

  return lastShardsToRemoveFrom;
}

export const splitRemainingAmountAmongShards = (
  shards: IdentifiableLiquidityPoolShard[],
  usedUtxos: UtxoMetaData[],
  remainingAmount: bigint,
  modifiedBy: UpdateLiquidityBy,
): bigint[] => {
  if (shards.length === 0) {
    return [remainingAmount];
  }

  const numShards = shards.length;
  const remainingAmountPerShard: bigint[] = new Array(numShards).fill(
    BigInt(0),
  );

  // Calculate the current amount in each shard, excluding used UTXOs
  for (let index = 0; index < numShards; index++) {
    const shard = shards[index];
    let remainingAmountInShard: bigint = BigInt(0);

    switch (modifiedBy) {
      case UpdateLiquidityBy.BtcAmount:
        // Sum up the values of btc_utxos that are not in usedUtxos
        remainingAmountInShard = shard.btcUtxos
          .filter(
            (utxoInfo) =>
              !usedUtxos.some(
                (usedUtxo) =>
                  usedUtxo.txid === utxoInfo.txid &&
                  usedUtxo.vout === utxoInfo.vout,
              ),
          )
          .reduce((sum, utxoInfo) => sum + utxoInfo.value, BigInt(0));
        break;

      case UpdateLiquidityBy.RuneAmount:
        if (
          shard.runeUtxo &&
          !usedUtxos.some(
            (usedUtxo) =>
              usedUtxo.txid === shard.runeUtxo!.txid &&
              usedUtxo.vout === shard.runeUtxo!.vout,
          )
        ) {
          remainingAmountInShard = shard.runeUtxo!.runes[0].amount;
        } else {
          remainingAmountInShard = BigInt(0);
        }
        break;

      case UpdateLiquidityBy.Liquidity:
        remainingAmountInShard = shard.liquidity;
        break;

      default:
        remainingAmountInShard = BigInt(0);
        break;
    }

    remainingAmountPerShard[index] = remainingAmountInShard;
  }

  // Calculate total current amount and desired amount per shard
  const totalCurrentAmount = remainingAmountPerShard.reduce(
    (sum, amount) => sum + amount,
    BigInt(0),
  );
  const totalAmountAfterDistribution = totalCurrentAmount + remainingAmount;
  const desiredAmountPerShard =
    totalAmountAfterDistribution / BigInt(numShards);

  // Calculate the needed amount for each shard to reach the desired amount
  const neededAmounts: bigint[] = remainingAmountPerShard.map(
    (currentAmount) =>
      desiredAmountPerShard > currentAmount
        ? desiredAmountPerShard - currentAmount
        : BigInt(0),
  );

  // Calculate total needed amount
  const totalNeededAmount = neededAmounts.reduce(
    (sum, amount) => sum + amount,
    BigInt(0),
  );
  let assignedAmounts: bigint[] = [];

  if (totalNeededAmount <= remainingAmount) {
    // If we have enough remainingAmount to fulfill needed amounts
    assignedAmounts = neededAmounts.slice();
    const remainingAmountLeftover = remainingAmount - totalNeededAmount;
    const perShardExtra = remainingAmountLeftover / BigInt(numShards);
    let extraLeftover = remainingAmountLeftover % BigInt(numShards);

    for (let i = 0; i < numShards; i++) {
      assignedAmounts[i] += perShardExtra;
      if (extraLeftover > BigInt(0)) {
        assignedAmounts[i] += BigInt(1);
        extraLeftover -= BigInt(1);
      }
    }
  } else {
    // Distribute the remainingAmount proportionally to the needed amounts
    assignedAmounts = new Array(numShards).fill(BigInt(0));
    let cumulative = BigInt(0);
    let cumulativeNeeded = BigInt(0);

    for (let i = 0; i < numShards; i++) {
      cumulativeNeeded += neededAmounts[i];
      const proportionalAmount =
        (remainingAmount * cumulativeNeeded) / totalNeededAmount;
      assignedAmounts[i] = proportionalAmount - cumulative;
      cumulative = proportionalAmount;
    }
  }

  console.log('assignedAmounts', assignedAmounts);

  return assignedAmounts;
};

export function selectBestShardToAddTo(
  poolShards: LiquidityPoolShard[],
  shardIndexes: number[] | undefined,
  addBy: UpdateLiquidityBy,
): number {
  shardIndexes =
    shardIndexes ?? Array.from({ length: poolShards.length }, (_, i) => i);

  if (shardIndexes.length === 0) {
    return -1;
  }

  // Find the index of the shard with the least liquidity
  return shardIndexes.reduce((minIndex, currentIndex) => {
    const minKey = getShardsKey(poolShards[minIndex], addBy);
    const currentKey = getShardsKey(poolShards[currentIndex], addBy);
    return currentKey < minKey ? currentIndex : minIndex;
  });
}

// Ensures that for a BTC-only removal, the BTC that remains across the selected
// shards is either zero or at least the dust limit. If not, it expands the
// selection (by BTC shard priority) until the remainder is safe or throws.
export function enforceBtcDustRemainder(
  pool: IdentifiableLiquidityPool,
  selection: IdentifiableLiquidityPoolShard[],
  amountToRemoveBtc: bigint,
  mempoolInfo: MempoolInfoMap,
): IdentifiableLiquidityPoolShard[] {
  // Compute total BTC capacity of current selection
  let totalSelectedBtc = selection.reduce(
    (sum, shard) => sum + getShardsKey(shard, UpdateLiquidityBy.BtcAmount),
    0n,
  );

  if (totalSelectedBtc <= amountToRemoveBtc) {
    // no positive remainder to worry about
    return selection;
  }

  let remainder = totalSelectedBtc - amountToRemoveBtc;
  if (remainder === 0n || remainder >= DUST_LIMIT) {
    return selection;
  }

  // Try to add more shards (by BTC priority) until remainder reaches dust limit
  const selectedSet = new Set(selection.map((s) => s.pubkey));

  const candidates = pool.shards
    // Only shards with some BTC available and not already selected
    .filter((shard) => {
      if (selectedSet.has(shard.pubkey)) return false;
      const key = getShardsKey(shard, UpdateLiquidityBy.BtcAmount);
      if (key <= 0n) return false;

      // Apply mempool safety filters similar to selection
      const entry = getMempoolEntryFromShard(shard, mempoolInfo);
      if (!entry) return true; // confirmed → safe

      const descendantCapacity = getRemainingDescendantCapacity(
        shard,
        mempoolInfo,
      );
      if (descendantCapacity === 0) return false;

      return entry.ancestorsCount < 23 && entry.descendantsCount < 23;
    })
    .slice();

  // Sort by the same shard priority rules (for BTC)
  sortShardsByPriority(candidates, UpdateLiquidityBy.BtcAmount, mempoolInfo);

  const updated: IdentifiableLiquidityPoolShard[] = selection.slice();

  for (const shard of candidates) {
    if (remainder >= DUST_LIMIT) break;
    const key = getShardsKey(shard, UpdateLiquidityBy.BtcAmount);
    if (key <= 0n) continue;

    updated.push(shard);
    selectedSet.add(shard.pubkey);
    remainder += key;
  }

  if (remainder > 0n && remainder < DUST_LIMIT) {
    throw new PoolErrorException({
      message: 'Not enough liquidity to satisfy BTC dust remainder',
      type: PoolErrorType.InsufficientLiquidity,
      token: 'BTC',
      maxAmount: (totalSelectedBtc - remainder).toString(),
    });
  }

  return updated;
}