import { Assets, TxBuilder } from '@lucid-evolution/lucid';
import { LucidContext } from '../test-helpers';
import { createUtxoAtTreasury } from '../endpoints/treasury';
import { SystemParams, treasuryMerge, treasurySplit } from '../../src';
import { findAllTreasuryUtxos } from './treasury-queries';
import { assert } from 'vitest';

export async function treasuryMergeTx(
  context: LucidContext,
  sysParams: SystemParams,
  treasuryOutputs: Assets[],
): Promise<TxBuilder> {
  for (const output of treasuryOutputs) {
    await createUtxoAtTreasury(output, sysParams, context);
  }

  const treasuryUtxos = await findAllTreasuryUtxos(context.lucid, sysParams);

  // Simpler approach: for each expected output in treasuryOutputs, find one UTxO with exactly matching assets and take it
  // Find UTxOs for each expected output, but remove it from the candidate list as we match
  const availableUtxos = [...treasuryUtxos];
  const matchedTreasuryUtxos = treasuryOutputs
    .map((expectedAssets) => {
      const idx = availableUtxos.findIndex((utxo) => {
        const utxoAssets = utxo.assets;
        const keys1 = Object.keys(utxoAssets).sort();
        const keys2 = Object.keys(expectedAssets).sort();
        if (keys1.length !== keys2.length) return false;
        return keys1.every(
          (k, i) => k === keys2[i] && utxoAssets[k] === expectedAssets[k],
        );
      });
      if (idx !== -1) {
        return availableUtxos.splice(idx, 1)[0];
      } else {
        return undefined;
      }
    })
    .filter((utxo) => utxo !== undefined);

  assert(
    matchedTreasuryUtxos.length === treasuryOutputs.length,
    'Expected all treasury outputs to be matched',
  );

  return await treasuryMerge(matchedTreasuryUtxos, context.lucid, sysParams);
}

export async function treasurySplitTx(
  context: LucidContext,
  sysParams: SystemParams,
  treasuryOutput: Assets,
): Promise<TxBuilder> {
  await createUtxoAtTreasury(treasuryOutput, sysParams, context);

  const treasuryUtxos = await findAllTreasuryUtxos(context.lucid, sysParams);

  const matchedTreasuryUtxo = treasuryUtxos.find((utxo) =>
    Object.keys(treasuryOutput).every(
      (key) => utxo.assets[key] === treasuryOutput[key],
    ),
  );

  assert(matchedTreasuryUtxo !== undefined, 'Expected a single treasury UTXO');

  return await treasurySplit(matchedTreasuryUtxo, context.lucid, sysParams);
}