import {
  addAssets,
  Assets,
  Data,
  fromHex,
  LucidEvolution,
  OutRef,
  TxBuilder,
  UTxO,
} from '@lucid-evolution/lucid';
import {
  fromSystemParamsScriptRef,
  SystemParams,
} from '../../types/system-params';
import { matchSingle } from '../../utils/utils';
import {
  adaAssetClass,
  AssetClass,
  estimateUtxoMinLovelace,
  mkAssetsOf,
  mkLovelacesOf,
  negateAssets,
} from '@3rd-eye-labs/cardano-offchain-common';
import { serialiseTreasuryDatum, serialiseTreasuryRedeemer } from './types-new';
import { array as A, function as F } from 'fp-ts';
import { parseExecuteDatumOrThrow } from '../execute/types-new';
import { getInlineDatumOrThrow } from '../../utils/lucid-utils';
import { createValueFromWithdrawal } from '../gov/helpers';
import { mkTreasuryAddr } from './helpers';

export async function treasuryFeeTx(
  assetToPay: AssetClass,
  amountToPay: bigint,
  extraLovelaces: bigint,
  lucid: LucidEvolution,
  sysParams: SystemParams,
  tx: TxBuilder,
  actionOref: OutRef,
  /**
   * `undefined` in case using direct treasury payment.
   */
  treasuryOref: OutRef | undefined,
): Promise<UTxO | null> {
  const [asset, amt, extraLov] = ((): [AssetClass, bigint, bigint] => {
    if (amountToPay <= 0 && extraLovelaces > 0) {
      return [adaAssetClass, extraLovelaces, 0n];
    }

    return [assetToPay, amountToPay, extraLovelaces];
  })();

  // If no fee is to be collected, do not include treasury collection.
  if (amt <= 0n) return null;

  let treasuryRefScriptUtxo = null;

  if (treasuryOref !== undefined) {
    const treasuryUtxo = matchSingle(
      await lucid.utxosByOutRef([treasuryOref]),
      (_) => new Error('Expected a single treasury UTXO'),
    );

    treasuryRefScriptUtxo = matchSingle(
      await lucid.utxosByOutRef([
        fromSystemParamsScriptRef(
          sysParams.scriptReferences.treasuryValidatorRef,
        ),
      ]),
      (_) => new Error('Expected a single treasury Ref Script UTXO'),
    );

    tx.readFrom([treasuryRefScriptUtxo])
      .collectFrom(
        [treasuryUtxo],
        serialiseTreasuryRedeemer({
          Collect: {
            assetToCollect: asset,
            amountToCollect: amt,
            extraLovelaces: extraLov,
            actionInputOref: {
              txHash: fromHex(actionOref.txHash),
              outputIndex: BigInt(actionOref.outputIndex),
            },
          },
        }),
      )
      .pay.ToContract(
        mkTreasuryAddr(lucid, sysParams),
        {
          kind: 'inline',
          value: serialiseTreasuryDatum({
            treasuryInputOref: {
              txHash: fromHex(treasuryOref.txHash),
              outputIndex: BigInt(treasuryOref.outputIndex),
            },
            actionInputOref: {
              txHash: fromHex(actionOref.txHash),
              outputIndex: BigInt(actionOref.outputIndex),
            },
          }),
        },
        addAssets(
          treasuryUtxo.assets,
          mkAssetsOf(asset, amt),
          mkLovelacesOf(extraLov),
        ),
      );
  } else {
    tx.pay.ToContract(
      mkTreasuryAddr(lucid, sysParams),
      {
        kind: 'inline',
        value: serialiseTreasuryDatum({
          treasuryInputOref: null,
          actionInputOref: {
            txHash: fromHex(actionOref.txHash),
            outputIndex: BigInt(actionOref.outputIndex),
          },
        }),
      },
      addAssets(mkAssetsOf(asset, amt), mkLovelacesOf(extraLov)),
    );
  }

  return treasuryRefScriptUtxo;
}

export async function treasuryCollect(
  assetToPay: AssetClass,
  amountToPay: bigint,
  extraLovelaces: bigint,
  lucid: LucidEvolution,
  sysParams: SystemParams,
  actionOref: OutRef,
  treasuryOref: OutRef,
): Promise<TxBuilder> {
  const treasuryUtxo = matchSingle(
    await lucid.utxosByOutRef([treasuryOref]),
    (_) => new Error('Expected a single treasury UTXO'),
  );

  const treasuryRefScriptUtxo = matchSingle(
    await lucid.utxosByOutRef([
      fromSystemParamsScriptRef(
        sysParams.scriptReferences.treasuryValidatorRef,
      ),
    ]),
    (_) => new Error('Expected a single treasury Ref Script UTXO'),
  );

  return lucid
    .newTx()
    .readFrom([treasuryRefScriptUtxo])
    .collectFrom(
      [treasuryUtxo],
      serialiseTreasuryRedeemer({
        Collect: {
          assetToCollect: assetToPay,
          amountToCollect: amountToPay,
          extraLovelaces: extraLovelaces,
          actionInputOref: {
            txHash: fromHex(actionOref.txHash),
            outputIndex: BigInt(actionOref.outputIndex),
          },
        },
      }),
    )
    .pay.ToContract(
      mkTreasuryAddr(lucid, sysParams),
      {
        kind: 'inline',
        value: serialiseTreasuryDatum({
          treasuryInputOref: {
            txHash: fromHex(treasuryOref.txHash),
            outputIndex: BigInt(treasuryOref.outputIndex),
          },
          actionInputOref: {
            txHash: fromHex(actionOref.txHash),
            outputIndex: BigInt(actionOref.outputIndex),
          },
        }),
      },
      addAssets(
        treasuryUtxo.assets,
        mkAssetsOf(assetToPay, amountToPay),
        mkLovelacesOf(extraLovelaces),
      ),
    );
}

export async function treasuryMerge(
  treasuryOutRefs: OutRef[],
  lucid: LucidEvolution,
  sysParams: SystemParams,
): Promise<TxBuilder> {
  const treasuryScriptRefUtxo = matchSingle(
    await lucid.utxosByOutRef([
      fromSystemParamsScriptRef(
        sysParams.scriptReferences.treasuryValidatorRef,
      ),
    ]),
    (_) => new Error('Expected a single treasury Ref Script UTXO'),
  );

  const treasuryUtxos = await lucid.utxosByOutRef(treasuryOutRefs);

  const totalAssets = F.pipe(
    treasuryUtxos,
    A.reduce<UTxO, Assets>({}, (acc, utxo) => addAssets(acc, utxo.assets)),
  );

  return lucid
    .newTx()
    .readFrom([treasuryScriptRefUtxo])
    .collectFrom(treasuryUtxos, serialiseTreasuryRedeemer('Merge'))
    .pay.ToContract(
      mkTreasuryAddr(lucid, sysParams),
      { kind: 'inline', value: Data.void() },
      totalAssets,
    );
}

export async function treasurySplit(
  treasuryOutRef: OutRef,
  lucid: LucidEvolution,
  sysParams: SystemParams,
): Promise<TxBuilder> {
  const treasuryScriptRefUtxo = matchSingle(
    await lucid.utxosByOutRef([
      fromSystemParamsScriptRef(
        sysParams.scriptReferences.treasuryValidatorRef,
      ),
    ]),
    (_) => new Error('Expected a single treasury Ref Script UTXO'),
  );

  const treasuryUtxo = matchSingle(
    await lucid.utxosByOutRef([treasuryOutRef]),
    (_) => new Error('Expected a single treasury UTXO'),
  );

  const [adaAsset, ...nonAdaAssets] = Object.keys(treasuryUtxo.assets);

  const tx = lucid
    .newTx()
    .collectFrom([treasuryUtxo], serialiseTreasuryRedeemer('Split'))
    .readFrom([treasuryScriptRefUtxo]);

  let paidLovelaces = 0n;

  for (const asset of A.reverse(nonAdaAssets.sort())) {
    const utxoMinLovelace = estimateUtxoMinLovelace(
      lucid.config().protocolParameters!,
      mkTreasuryAddr(lucid, sysParams),
      { [asset]: treasuryUtxo.assets[asset] },
      { InlineDatum: { datum: Data.void() } },
    );

    paidLovelaces += utxoMinLovelace;

    tx.pay.ToContract(
      mkTreasuryAddr(lucid, sysParams),
      { kind: 'inline', value: Data.void() },
      addAssets(
        { [asset]: treasuryUtxo.assets[asset] },
        mkLovelacesOf(utxoMinLovelace),
      ),
    );
  }

  tx.pay.ToContract(
    mkTreasuryAddr(lucid, sysParams),
    { kind: 'inline', value: Data.void() },
    { [adaAsset]: treasuryUtxo.assets[adaAsset] - paidLovelaces },
  );

  return tx;
}

export async function treasuryPrepareWithdrawal(
  treasuryOutRefs: OutRef[],
  upgradeOutRef: OutRef,
  lucid: LucidEvolution,
  sysParams: SystemParams,
): Promise<TxBuilder> {
  const treasuryScriptRefUtxo = matchSingle(
    await lucid.utxosByOutRef([
      fromSystemParamsScriptRef(
        sysParams.scriptReferences.treasuryValidatorRef,
      ),
    ]),
    (_) => new Error('Expected a single treasury Ref Script UTXO'),
  );

  const treasuryUtxos = await lucid.utxosByOutRef(treasuryOutRefs);

  const upgradeUtxo = matchSingle(
    await lucid.utxosByOutRef([upgradeOutRef]),
    (_) => new Error('Expected a single upgrade UTXO'),
  );

  const treasuryAddress = mkTreasuryAddr(lucid, sysParams);

  const totalAssets = F.pipe(
    treasuryUtxos,
    A.reduce<UTxO, Assets>({}, (acc, utxo) => addAssets(acc, utxo.assets)),
  );
  const executeDatum = parseExecuteDatumOrThrow(
    getInlineDatumOrThrow(upgradeUtxo),
  );
  if (!executeDatum.treasuryWithdrawal)
    throw new Error('Expected a treasury withdrawal in the execute datum');
  const withdrawalVal = createValueFromWithdrawal(
    executeDatum.treasuryWithdrawal,
  );
  const change = addAssets(totalAssets, negateAssets(withdrawalVal));

  const tx = lucid
    .newTx()
    .collectFrom(treasuryUtxos, serialiseTreasuryRedeemer('PrepareWithdraw'))
    .readFrom([treasuryScriptRefUtxo, upgradeUtxo])
    .pay.ToContract(
      treasuryAddress,
      { kind: 'inline', value: Data.void() },
      withdrawalVal,
    )
    .pay.ToContract(
      treasuryAddress,
      { kind: 'inline', value: Data.void() },
      change,
    );

  return tx;
}