import { beforeEach, describe, expect, it } from 'vitest';
import { PoolErrorException } from '../../error/pool.error';
import {
  getShardsKey,
  selectBestShardsToRemoveFrom,
  sortShardsByPriority,
  splitRemainingAmountAmongShards,
  UpdateLiquidityBy,
} from '../shards';
import { MempoolEntry, MempoolInfoMap } from '../../providers/bitcoin.provider';

interface IdentifiableLiquidityPoolShard {
  poolPubkey: string;
  runeUtxo?: { utxo: any; runes: Array<{ amount: bigint }> };
  btcUtxos: Array<{ utxo: any; value: bigint }>;
  liquidity: bigint;
  protocolFeeOwed: bigint;
  lastBlockHeight: bigint;
  timesUpdated: bigint;
  kLast: bigint;
  kLastCounter: bigint;
  pubkey: string;
  utxo: string;
}

interface IdentifiableLiquidityPool {
  shards: IdentifiableLiquidityPoolShard[];
  config: {
    token_0: { block: bigint; tx: number };
    token_1: { block: bigint; tx: number };
  };
  liquidity: bigint;
  price: string;
  token0Amount: bigint;
  token1Amount: bigint;
}

const createTestShard = (
  params: Partial<IdentifiableLiquidityPoolShard>,
): IdentifiableLiquidityPoolShard => ({
  poolPubkey: params.poolPubkey || 'pool1',
  runeUtxo: params.runeUtxo,
  btcUtxos: params.btcUtxos || [],
  liquidity: params.liquidity ?? 0n,
  protocolFeeOwed: params.protocolFeeOwed ?? 0n,
  lastBlockHeight: params.lastBlockHeight ?? 0n,
  timesUpdated: params.timesUpdated ?? 0n,
  kLast: params.kLast ?? 0n,
  kLastCounter: params.kLastCounter ?? 0n,
  pubkey: params.pubkey || 'dummyPubkey',
  utxo: params.utxo || 'dummyTxId',
});

const createLiquidityShard = (
  liquidity: bigint,
): Partial<IdentifiableLiquidityPoolShard> => ({
  liquidity,
});

const createBtcAmountShard = (
  totalBtc: bigint,
): Partial<IdentifiableLiquidityPoolShard> => ({
  btcUtxos: [{ utxo: { txid: 'btcTx', vout: 0 }, value: totalBtc }],
});

const createRuneAmountShard = (
  runeAmount: bigint,
): Partial<IdentifiableLiquidityPoolShard> => ({
  runeUtxo: {
    utxo: { txid: 'runeTx', vout: 0 },
    runes: [{ amount: runeAmount }],
  },
});

const currentBlockHeight = 50n;

const dummyMempoolInfo: MempoolInfoMap = new Map();

describe('splitRemainingAmountAmongShards', () => {
  const usedUtxos: any[] = [];

  beforeEach(() => {
    // Clear usedUtxos if needed.
    usedUtxos.length = 0;
  });

  it('Returns an array with remainingAmount when no shards provided', () => {
    const result = splitRemainingAmountAmongShards(
      [],
      usedUtxos,
      1000n,
      UpdateLiquidityBy.Liquidity,
    );

    expect(result).toEqual([1000n]);
  });

  describe('UpdateLiquidityBy.Liquidity', () => {
    it('Distributes remainingAmount proportionally when total needed exceeds remainingAmount', () => {
      // Two shards: one with liquidity 1000n and one with 2000n.
      const shards = [createLiquidityShard(1000n), createLiquidityShard(2000n)];
      // Total current = 3000n, remainingAmount = 900n.
      // Total after distribution = 3900n, desired per shard = 1950n.
      // Needed: shard 1: 1950 - 1000 = 950n, shard 2: 1950 - 2000 = 0n.
      // Total needed = 950n > remainingAmount (900n), so distribute proportionally.
      // For shard1: assigned = (900 * 950/950) = 900n; for shard2: assigned = 0.
      const result = splitRemainingAmountAmongShards(
        shards as any,
        usedUtxos,
        900n,
        UpdateLiquidityBy.Liquidity,
      );
      expect(result).toEqual([900n, 0n]);
    });

    it('Distributes remainingAmount evenly when enough to cover needed amounts', () => {
      // Two shards: liquidity 1000n and 2000n, remainingAmount = 1000n.
      const shards = [createLiquidityShard(1000n), createLiquidityShard(2000n)];
      // Total current = 3000n, total after = 4000n, desired per shard = 2000n.
      // Needed: shard1: 2000 - 1000 = 1000n, shard2: 2000 - 2000 = 0n.
      // Total needed = 1000n, which equals remainingAmount.
      // Then, assignedAmounts = neededAmounts + leftover evenly.
      // Leftover = remainingAmount - totalNeeded = 0n, so assignedAmounts = [1000n, 0n].
      const result = splitRemainingAmountAmongShards(
        shards as any,
        usedUtxos,
        1000n,
        UpdateLiquidityBy.Liquidity,
      );
      expect(result).toEqual([1000n, 0n]);
    });
  });

  describe('UpdateLiquidityBy.BtcAmount', () => {
    it('Calculates current amount from btc_utxos and distributes remaining amount proportionally', () => {
      // Two shards with btc_utxos: first has 1500n, second has 2500n.
      const shards = [createBtcAmountShard(1500n), createBtcAmountShard(2500n)];
      // Total current = 1500n + 2500n = 4000n.
      // Let remainingAmount = 1000n, so total after = 5000n, desired per shard = 2500n.
      // Needed: shard1: 2500 - 1500 = 1000n; shard2: 2500 - 2500 = 0n.
      // Total needed = 1000n, equals remainingAmount.
      const result = splitRemainingAmountAmongShards(
        shards as any,
        usedUtxos,
        1000n,
        UpdateLiquidityBy.BtcAmount,
      );
      expect(result).toEqual([1000n, 0n]);
    });
  });

  describe('UpdateLiquidityBy.RuneAmount', () => {
    it('Calculates current amount from rune_utxo and distributes remaining amount correctly when sufficient', () => {
      // Two shards with rune amounts: first has 500n, second has 1500n.
      const shards = [
        createRuneAmountShard(500n),
        createRuneAmountShard(1500n),
      ];
      // Total current = 500n + 1500n = 2000n.
      // Let remainingAmount = 1000n, so total after = 3000n, desired per shard = 1500n.
      // Needed: shard1: 1500 - 500 = 1000n; shard2: 1500 - 1500 = 0n.
      // Total needed = 1000n, equals remainingAmount.
      const result = splitRemainingAmountAmongShards(
        shards as any,
        usedUtxos,
        1000n,
        UpdateLiquidityBy.RuneAmount,
      );
      expect(result).toEqual([1000n, 0n]);
    });
  });

  describe('Extra remaining amount distribution (totalNeeded <= remainingAmount)', () => {
    it('Distributes extra remaining amount evenly', () => {
      // Two shards: liquidity 1000n and 2000n.
      const shards = [createLiquidityShard(1000n), createLiquidityShard(2000n)];
      // Total current = 3000n, let remainingAmount = 1100n.
      // Total after = 4100n, desired per shard = 2050n.
      // Needed: shard1: 2050 - 1000 = 1050n; shard2: 2050 - 2000 = 50n.
      // Total needed = 1050 + 50 = 1100n exactly.
      // In this case, assignedAmounts should equal neededAmounts.
      const result = splitRemainingAmountAmongShards(
        shards as any,
        usedUtxos,
        1100n,
        UpdateLiquidityBy.Liquidity,
      );
      expect(result).toEqual([1050n, 50n]);
    });

    it('Distributes extra remaining amount evenly when there is leftover after fulfilling needs', () => {
      // Two shards: liquidity 1000n and 2000n.
      const shards = [createLiquidityShard(1000n), createLiquidityShard(2000n)];
      // Total current = 3000n, let remainingAmount = 1200n.
      // Total after = 4200n, desired per shard = 2100n.
      // Needed: shard1: 2100 - 1000 = 1100n; shard2: 2100 - 2000 = 100n.
      // Total needed = 1200n, equals remainingAmount.
      // So assignedAmounts should be [1100n, 100n].
      const result = splitRemainingAmountAmongShards(
        shards as any,
        usedUtxos,
        1200n,
        UpdateLiquidityBy.Liquidity,
      );
      expect(result).toEqual([1100n, 100n]);
    });
  });

  describe('Proportional distribution when remainingAmount < totalNeededAmount', () => {
    it('Distributes remaining amount proportionally', () => {
      // Two shards: liquidity 1000n and 2000n.
      const shards = [createLiquidityShard(1000n), createLiquidityShard(2000n)];
      // Total current = 3000n, let remainingAmount = 900n.
      // Total after = 3900n, desired per shard = 1950n.
      // Needed: shard1: 1950 - 1000 = 950n; shard2: 1950 - 2000 = 0n.
      // Total needed = 950n > remainingAmount (900n).
      // Proportional: For shard1, assigned = (900 * 950/950) = 900, shard2 = 0.
      const result = splitRemainingAmountAmongShards(
        shards as any,
        usedUtxos,
        900n,
        UpdateLiquidityBy.Liquidity,
      );
      expect(result).toEqual([900n, 0n]);
    });
  });
});

describe('selectBestShardsToRemoveFrom', () => {
  let testPool: IdentifiableLiquidityPool;
  beforeEach(() => {
    testPool = {
      shards: [
        createTestShard({
          liquidity: 300n,
          lastBlockHeight: 40n,
          timesUpdated: 1n,
          utxo: 'tx1',
        }),
        createTestShard({
          liquidity: 500n,
          lastBlockHeight: 40n,
          timesUpdated: 2n,
          utxo: 'tx2',
        }),
        createTestShard({
          liquidity: 200n,
          lastBlockHeight: 60n,
          timesUpdated: 1n,
          utxo: 'tx3',
        }),
      ],
      config: {
        token_0: { block: 100n, tx: 1 },
        token_1: { block: 200n, tx: 2 },
      },
      liquidity: 10000n,
      price: '0.5',
      token0Amount: 5000n,
      token1Amount: 5000n,
    };
  });

  it('Throws an error when total liquidity is insufficient', () => {
    // 300n + 500n + 200n = 1000n < 2000n
    expect(() =>
      selectBestShardsToRemoveFrom(
        testPool as any,
        2000n,
        UpdateLiquidityBy.Liquidity,
        dummyMempoolInfo,
      ),
    ).toThrow(PoolErrorException);
  });

  it('Selects shards until the required liquidity is met', () => {
    const selectedShards = selectBestShardsToRemoveFrom(
      testPool as any,
      600n,
      UpdateLiquidityBy.Liquidity,
      dummyMempoolInfo,
    );

    const totalSelected = selectedShards.reduce(
      (sum, shard) => sum + shard.liquidity,
      0n,
    );

    expect(totalSelected).toBeGreaterThanOrEqual(600n);

    selectedShards.forEach((shard) => {
      expect(testPool.shards).toContainEqual(shard);
    });

    expect(selectedShards.length).toBe(2);
  });

  it('Returns shards sorted in descending order by liquidity', () => {
    const selectedShards = selectBestShardsToRemoveFrom(
      testPool as any,
      600n,
      UpdateLiquidityBy.Liquidity,
      dummyMempoolInfo,
    );

    for (let i = 0; i < selectedShards.length - 1; i++) {
      expect(selectedShards[i].liquidity).toBeGreaterThanOrEqual(
        selectedShards[i + 1].liquidity,
      );
    }
  });

  it('Selects the minimal set of shards needed to reach the required liquidity', () => {
    const selectedShards = selectBestShardsToRemoveFrom(
      testPool as any,
      500n,
      UpdateLiquidityBy.Liquidity,
      dummyMempoolInfo,
    );

    expect(selectedShards).toHaveLength(1);
    const total = selectedShards.reduce(
      (sum, shard) => sum + shard.liquidity,
      0n,
    );
    expect(total).toBeGreaterThanOrEqual(300n);
  });

  it('Selects shards in order of increasing index and then sorts by key descending', () => {
    // Create a pool where multiple shards are available.
    const shardA = createTestShard({
      utxo: 'txA',
      liquidity: 300n,
      timesUpdated: 1n,
    });
    const shardB = createTestShard({
      utxo: 'txB',
      liquidity: 500n,
      timesUpdated: 1n,
    });
    const shardC = createTestShard({
      utxo: 'txC',
      liquidity: 200n,
      timesUpdated: 1n,
    });
    const testPool1 = { ...testPool, shards: [shardA, shardB, shardC] };

    const selectedShards = selectBestShardsToRemoveFrom(
      testPool1 as any,
      600n,
      UpdateLiquidityBy.Liquidity,
      dummyMempoolInfo,
    );

    // The function first collects indices based on order, then sorts them descending by key.
    // Selected shards to be sorted with the highest liquidity first.
    expect(selectedShards[0].utxo).toBe('txB');
    expect(selectedShards[1].utxo).toBe('txA');
  });

  it('Filters out shards with a zero key', () => {
    const shardA = createTestShard({ liquidity: 0n, utxo: 'txA' });
    const shardB = createTestShard({ liquidity: 800n, utxo: 'txB' });
    const testPool1 = { ...testPool, shards: [shardA, shardB] };
    const selected = selectBestShardsToRemoveFrom(
      testPool1 as any,
      500n,
      UpdateLiquidityBy.Liquidity,
      dummyMempoolInfo,
    );

    expect(selected).toHaveLength(1);
    expect(selected[0].utxo).toBe('txB');
  });
});

describe('sortShardsByPriority', () => {
  const currentBlockHeight = 100n;
  let mempoolInfo: MempoolInfoMap;

  beforeEach(() => {
    mempoolInfo = new Map<string, MempoolEntry>();
  });

  it('Prioritizes shards from previous blocks over current blocks', () => {
    const shardPrev = createTestShard({
      utxo: 'txPrev',
      lastBlockHeight: 50n,
      timesUpdated: 1n,
      liquidity: 1000n,
    });
    const shardCurrent = createTestShard({
      utxo: 'txCurrent',
      lastBlockHeight: 150n,
      timesUpdated: 1n,
      liquidity: 1000n,
    });
    const shards = [shardCurrent, shardPrev];

    sortShardsByPriority(
      shards as any,
      UpdateLiquidityBy.Liquidity,
      mempoolInfo,
    );

    expect(shards[0].utxo).toBe('txPrev');
  });

  it('Prioritizes shards with fewer descendants', () => {
    // Create two shards with equal block heights.
    const shardA = createTestShard({
      utxo: 'txA',
      lastBlockHeight: 80n,
      timesUpdated: 1n,
      liquidity: 1000n,
    });
    const shardB = createTestShard({
      utxo: 'txB',
      lastBlockHeight: 80n,
      timesUpdated: 1n,
      liquidity: 1000n,
    });

    mempoolInfo.set('txA', {
      descendantsCount: 2,
      ancestorsCount: 1,
      ancestorsSize: 1,
      descendantsSize: 0,
      depends: [],
      spentby: [],
      fees: { ancestor: 0, base: 0, modified: 0 },
    });
    mempoolInfo.set('txB', {
      descendantsCount: 5,
      ancestorsCount: 1,
      ancestorsSize: 1,
      descendantsSize: 0,
      depends: [],
      spentby: [],
      fees: { ancestor: 0, base: 0, modified: 0 },
    });
    const shards = [shardB, shardA];

    sortShardsByPriority(
      shards as any,
      UpdateLiquidityBy.Liquidity,
      mempoolInfo,
    );

    expect(shards[0].utxo).toBe('txA');
  });

  it('prioritizes shards with fewer ancestors when descendants are equal', () => {
    // Equal block height and descendant count
    const shardA = createTestShard({
      utxo: 'txA',
      lastBlockHeight: 80n,
      timesUpdated: 1n,
      liquidity: 1000n,
    });
    const shardB = createTestShard({
      utxo: 'txB',
      lastBlockHeight: 80n,
      timesUpdated: 1n,
      liquidity: 1000n,
    });

    mempoolInfo.set('txA', {
      descendantsCount: 3,
      ancestorsCount: 1,
      ancestorsSize: 1,
      descendantsSize: 0,
      depends: [],
      spentby: [],
      fees: { ancestor: 0, base: 0, modified: 0 },
    });
    mempoolInfo.set('txB', {
      descendantsCount: 3,
      ancestorsCount: 4,
      ancestorsSize: 1,
      descendantsSize: 0,
      depends: [],
      spentby: [],
      fees: { ancestor: 0, base: 0, modified: 0 },
    });
    const shards = [shardB, shardA];

    sortShardsByPriority(
      shards as any,
      UpdateLiquidityBy.Liquidity,
      mempoolInfo,
    );

    expect(shards[0].utxo).toBe('txA');
  });

  it('Prioritizes shards with fewer times_updated when previous criteria are equal', () => {
    const shardA = createTestShard({
      utxo: 'txA',
      lastBlockHeight: 80n,
      timesUpdated: 1n,
      liquidity: 1000n,
    });
    const shardB = createTestShard({
      utxo: 'txB',
      lastBlockHeight: 80n,
      timesUpdated: 3n,
      liquidity: 1000n,
    });
    const shards = [shardB, shardA];

    sortShardsByPriority(
      shards as any,
      UpdateLiquidityBy.Liquidity,
      mempoolInfo,
    );
    expect(shards[0].utxo).toBe('txA');
  });

  it('Sorts by shards key (largest first) when all other criteria are equal', () => {
    const shardA = createTestShard({
      utxo: 'txA',
      lastBlockHeight: 80n,
      timesUpdated: 1n,
      liquidity: 3000n,
    });
    const shardB = createTestShard({
      utxo: 'txB',
      lastBlockHeight: 80n,
      timesUpdated: 1n,
      liquidity: 5000n,
    });
    const shards = [shardA, shardB];

    sortShardsByPriority(
      shards as any,
      UpdateLiquidityBy.Liquidity,
      mempoolInfo,
    );

    expect(shards[0].utxo).toBe('txB');
  });
});

describe('getShardsKey', () => {
  it('Returns liquidity for UpdateLiquidityBy.Liquidity', () => {
    const shard = createTestShard({ liquidity: 5000n });
    const key = getShardsKey(shard as any, UpdateLiquidityBy.Liquidity);
    expect(key).toBe(5000n);
  });

  it('Returns sum of btc_utxos for UpdateLiquidityBy.BtcAmount', () => {
    const shard = createTestShard({
      btcUtxos: [
        { utxo: { txid: 'tx1', vout: 0 }, value: 1000n },
        { utxo: { txid: 'tx2', vout: 0 }, value: 2000n },
      ],
    });
    const key = getShardsKey(shard as any, UpdateLiquidityBy.BtcAmount);
    expect(key).toBe(3000n);
  });

  it('Returns rune amount for UpdateLiquidityBy.RuneAmount', () => {
    const shard = createTestShard({
      runeUtxo: { utxo: { txid: 'tx3', vout: 0 }, runes: [{ amount: 4000n }] },
    });
    const key = getShardsKey(shard as any, UpdateLiquidityBy.RuneAmount);

    expect(key).toBe(4000n);
  });

  it('Returns 0n for UpdateLiquidityBy.RuneAmount if no rune_utxo is present', () => {
    const shard = createTestShard({});
    const key = getShardsKey(shard as any, UpdateLiquidityBy.RuneAmount);
    expect(key).toBe(0n);
  });
});