import {
  IdentifiableLiquidityPool,
  SwapInstruction,
  swapMessage,
  SwapParams,
} from '@saturnbtcio/pool-serde-sdk';
import {
  EstimateError,
  finalizeTransaction,
  getBitcoinNetwork,
  getInputTypeFromAddress,
  getOutputTypeFromAddress,
  PsbtBuilder,
  toVsize,
  toXOnly,
  TransactionSizeCalculator,
} from '@saturnbtcio/psbt';
import { base64, hex } from '@scure/base';
import { SigHash, Transaction } from '@scure/btc-signer';
import { sha256 } from '@scure/btc-signer/utils';
import { PoolErrorException, PoolErrorType } from '../../error/pool.error';
import { SaturnSdkConfig } from '../../saturn-sdk';
import { getCalculatorForRuneToBtcTx } from '../../util/calculator';
import { DUST_LIMIT } from '../../util/constants';
import { calculateFeeRuneToBtcExactInSwap, checkFeeRate } from '../../util/fee';
import { findBtc, findCollection } from '../../util/finder';
import { buildUtxoInfoFromInputs } from '../../util/utxo-info';
import { CollectionUtxo } from '../../wallet/wallet.dto';
import { BTC_TOKEN } from '../pool.dto';
import {
  RuneToBtcPsbtRequest,
  RuneToBtcSwapMessageRequest,
  SplitRunePsbtRequest,
  SplitUtxoResponse,
} from './zero-to-one.dto';
import {
  validateRuneToBtcPsbt,
  validateSplitRunePsbt,
} from './zero-to-one.validation';
import { validatePoolSdkData } from '../../util/validation';
import { getTxIdsFromPool, getTxIdsFromUtxos } from '../../util/mempool';
import { MempoolEntry } from '../../providers/bitcoin.provider';
import { SanitizedMessage, UtxoMetaData } from '@saturnbtcio/arch-sdk';
import { createProtocolPda } from '../../account/pda-finder';

export class ZeroToOneSwap {
  private readonly config: SaturnSdkConfig;

  constructor(config: SaturnSdkConfig) {
    this.config = config;
  }

  async buildSplitRunePsbt(request: SplitRunePsbtRequest) {
    const {
      runePublicKey,
      runeAddress,
      paymentWallet,
      paymentPublicKey,
      collectionUtxos,
      collection,
      amountIn,
      feeRate,
      runeToBtcTxVsize,
    } = request;

    const runeWallet = await this.config.bitcoinProvider.getWallet(runeAddress);

    const txBuilder = new PsbtBuilder({
      network: this.config.network,
    });

    const collectionUtxoMempoolStatuses =
      await this.config.bitcoinProvider.getMempoolInfo(
        getTxIdsFromUtxos(collectionUtxos),
      );

    // Add collection inputs
    let btcIn = 0n;
    let totalRuneAmount = 0n;
    for (const utxo of collectionUtxos) {
      txBuilder.addInput({
        utxo,
        owner: runeWallet.address,
        publicKey: runePublicKey,
        mempoolStatus: collectionUtxoMempoolStatuses.get(utxo.txid),
        sighashType: SigHash.ALL,
      });
      btcIn += utxo.value;
      totalRuneAmount +=
        utxo.collectionStatuses.find((c) => c.id === collection.id)?.amount ??
        0n;
    }

    // Send back the runes to the user.
    txBuilder.addOutput({
      address: runeWallet.address,
      value: DUST_LIMIT,
      runes: [],
    });

    const fee = BigInt(runeToBtcTxVsize) * feeRate;

    txBuilder.addOutput({
      address: runeWallet.address,
      value: fee,
      runes: [
        {
          amount: amountIn,
          id: collection.id,
        },
      ],
    });

    const choosenWallet = paymentWallet ?? runeWallet;
    const walletMempoolStatus =
      await this.config.bitcoinProvider.getMempoolInfo(
        getTxIdsFromUtxos(choosenWallet.utxos),
      );

    const usedUtxos = collectionUtxos.map((utxo) => ({
      txid: utxo.txid,
      vout: utxo.vout,
      status: utxo.status,
      value: utxo.value,
    }));

    const btcUtxos = findBtc(
      choosenWallet,
      walletMempoolStatus,
      usedUtxos,
      collectionUtxoMempoolStatuses,
      txBuilder.getOutputValue(),
      feeRate,
      txBuilder,
      SigHash.ALL,
      paymentPublicKey ?? runePublicKey,
      DUST_LIMIT,
    );

    try {
      const tx = txBuilder.buildAndAdjustChange({
        feeRate: BigInt(feeRate),
        dustLimit: DUST_LIMIT,
        changeAddress: paymentWallet?.address ?? runeWallet.address,
        additionalInputAmount: 0n,
      });

      const txId = hex.encode(sha256(sha256(tx.psbt.toBytes(true))).reverse());

      const utxo: CollectionUtxo = {
        txid: txId,
        vout: 1,
        value: fee,
        collectionStatuses: [
          {
            ...collection,
            amount: amountIn,
          },
        ],
        hasInscription: false,
        status: {
          confirmed: false,
        },
      };

      const utxoMempoolEntry: MempoolEntry = {
        fees: {
          ancestor: Number(btcUtxos.ancestorsFee + tx.fee),
          modified: 0,
          base: 0,
        },
        ancestorsCount: 0,
        descendantsCount: 0,
        ancestorsSize: btcUtxos.ancestorsSize + tx.vsize,
        descendantsSize: 0,
        depends: [
          ...getTxIdsFromUtxos(collectionUtxos),
          ...getTxIdsFromUtxos(btcUtxos.utxos),
        ],
        spentby: [],
      };

      return {
        tx,
        utxo,
        utxoMempoolEntry,
        totalAmount: totalRuneAmount,
        amountToReceive: amountIn,
      };
    } catch (err) {
      let amount = 0n;
      if (err instanceof EstimateError) {
        const error = err as EstimateError;
        if (error.error.type === 'outputs-spending-more-than-inputs') {
          amount = error.error.amount;
        } else if (error.error.type === 'not-enough-funds') {
          amount = error.error.amount + error.error.fee;
        }
      }

      throw new PoolErrorException({
        type: PoolErrorType.NotEnoughFunds,
        message: `Not enough bitcoin. You need at least ${amount.toString()} in your wallet at current fee rate.`,
        maxAmount: (btcUtxos.amount + btcIn).toString(),
        minAmount: amount.toString(),
        token: `${BTC_TOKEN.block}:${BTC_TOKEN.tx}`,
      });
    }
  }

  async buildRuneToBtcPsbt(request: RuneToBtcPsbtRequest) {
    const scureNetwork = getBitcoinNetwork(this.config.network);

    validatePoolSdkData<RuneToBtcPsbtRequest>(request, scureNetwork);

    const {
      runeAddress,
      runePublicKey,
      paymentAddress,
      paymentPublicKey,
      amountIn,
      feeRate,
      poolId,
    } = request;
    let { amountOut } = request;

    const runeWallet = await this.config.bitcoinProvider.getWallet(runeAddress);
    const paymentWallet = paymentAddress
      ? await this.config.bitcoinProvider.getWallet(paymentAddress)
      : undefined;

    await checkFeeRate(Number(feeRate), this.config.bitcoinProvider);

    const pool = await this.config.indexerProvider.getPoolById(poolId);

    if (!pool) {
      throw new PoolErrorException({
        message: `Pool ${poolId} does not exist`,
        type: PoolErrorType.PoolNotFound,
        poolId: poolId,
      });
    }

    const collection = await this.config.indexerProvider.getCollection(
      pool.config.token0,
    );

    if (!collection) {
      throw new PoolErrorException({
        message: `Collection ${pool.config.token0} does not exist`,
        type: PoolErrorType.InvalidToken,
        token: pool.config.token0,
      });
    }

    const btcAmount = BigInt(amountOut);
    if (btcAmount < DUST_LIMIT) {
      throw new PoolErrorException({
        message: `BTC amount is below the dust limit. Please increase the bitcoin amount to ${DUST_LIMIT} sats.`,
        type: PoolErrorType.InvalidAmountBelowMin,
        token: pool.config.token1,
        minAmount: DUST_LIMIT.toString(),
      });
    }

    const txBuilder = new PsbtBuilder({
      network: this.config.network,
    });

    const shardsMempoolStatus =
      await this.config.bitcoinProvider.getMempoolInfo(getTxIdsFromPool(pool));

    let inputType = getInputTypeFromAddress(runeWallet.address);
    let outputType = getOutputTypeFromAddress(
      paymentWallet?.address ?? runeWallet.address,
      scureNetwork,
    );

    const runeToBtcTxEstimatedWeight = getCalculatorForRuneToBtcTx(
      inputType,
      outputType,
      pool,
      amountOut,
      shardsMempoolStatus,
    ).estimateSize({
      network: scureNetwork,
    }).weight;

    let runeUtxo: CollectionUtxo | undefined;
    let splitUtxoTx: SplitUtxoResponse | undefined;
    let utxoMempoolInfo: MempoolEntry | undefined;
    if (request.splitRunePsbt) {
      const { utxo, inputs } = await validateSplitRunePsbt(
        runeWallet,
        paymentWallet,
        collection,
        amountIn,
        request.splitRunePsbt,
        scureNetwork,
        toVsize(runeToBtcTxEstimatedWeight),
        feeRate,
      );
      runeUtxo = utxo;
      splitUtxoTx = undefined;
    } else {
      const walletMempoolStatus =
        await this.config.bitcoinProvider.getMempoolInfo(
          getTxIdsFromUtxos(runeWallet.utxos),
        );

      // Try to find an existing UTXO that already matches the split requirements
      // value == required fee for the swap tx and runes == exact amountIn for the collection
      const requiredFee = BigInt(toVsize(runeToBtcTxEstimatedWeight)) * feeRate;
      const preSplitMatch = runeWallet.utxos.find((u) => {
        const statuses = u.collectionStatuses ?? [];
        if (u.value !== requiredFee) return false;
        if (u.hasInscription) return false;
        if (statuses.length !== 1) return false;
        const status = statuses[0];
        return status.id === collection.id && status.amount === amountIn;
      });

      if (preSplitMatch) {
        runeUtxo = preSplitMatch;
        utxoMempoolInfo = walletMempoolStatus.get(preSplitMatch.txid);
        splitUtxoTx = undefined;
      } else {
        const collectionUtxos = findCollection(
          runeWallet,
          walletMempoolStatus,
          collection,
          amountIn,
        );

        if (collectionUtxos.amount < amountIn) {
          throw new PoolErrorException({
            message: `Not enough runes. You need at least ${amountIn} in your wallet.`,
            type: PoolErrorType.NotEnoughFunds,
            token: collection.id,
            maxAmount: collectionUtxos.amount.toString(),
            minAmount: '0',
          });
        }

        const { tx, utxo, utxoMempoolEntry, totalAmount, amountToReceive } =
          await this.buildSplitRunePsbt({
            runeAddress,
            runePublicKey,
            paymentAddress,
            paymentPublicKey,
            paymentWallet,
            collectionUtxos: collectionUtxos.utxos,
            collection,
            amountIn,
            feeRate,
            runeToBtcTxVsize: toVsize(runeToBtcTxEstimatedWeight),
          });
        runeUtxo = utxo;
        utxoMempoolInfo = utxoMempoolEntry;
        splitUtxoTx = {
          tx,
          totalAmount: totalAmount.toString(),
          amountToReceive: amountToReceive.toString(),
        };
      }
    }

    txBuilder.addInput({
      utxo: runeUtxo,
      owner: runeWallet.address,
      publicKey: runePublicKey,
      mempoolStatus: utxoMempoolInfo,
      sighashType: SigHash.SINGLE_ANYONECANPAY,
    });

    txBuilder.addOutput({
      address: paymentWallet?.address ?? runeWallet.address,
      value: amountOut,
      runes: [],
    });

    if (splitUtxoTx) {
      const estimatedSplitUtxoTxByteSize =
        TransactionSizeCalculator.createFromTransaction(
          splitUtxoTx.tx.psbt as Transaction,
        ).estimateByteSize({ network: scureNetwork });
      if (estimatedSplitUtxoTxByteSize >= this.config.maxTxSize) {
        throw new PoolErrorException({
          message: `Transaction too big. Please consolidate your rune utxos or try again with a smaller amount.`,
          type: PoolErrorType.InvalidTxSize,
        });
      }
    }

    return {
      tx: txBuilder.build(),
      splitUtxoTx,
    };
  }

  async swapMessage(request: RuneToBtcSwapMessageRequest): Promise<{
    message: SanitizedMessage;
    utxosInfo: UtxoMetaData[];
    splitRuneInputs: string[];
  }> {
    const scureNetwork = getBitcoinNetwork(this.config.network);

    validatePoolSdkData<RuneToBtcSwapMessageRequest>(request, scureNetwork);

    const {
      runeAddress,
      runePublicKey,
      paymentAddress,
      feeRate,
      poolId,
      amountIn,
      amountOut,
      exactIn,
      signedPsbt,
      splitRunePsbt,
    } = request;

    const wallet = await this.config.bitcoinProvider.getWallet(runeAddress);
    const paymentWallet = paymentAddress
      ? await this.config.bitcoinProvider.getWallet(paymentAddress)
      : undefined;

    await checkFeeRate(Number(feeRate), this.config.bitcoinProvider);

    const pool: IdentifiableLiquidityPool | undefined =
      await this.config.indexerProvider.getPoolById(poolId);
    if (!pool) {
      throw new PoolErrorException({
        message: `Pool ${poolId} does not exist`,
        type: PoolErrorType.PoolNotFound,
        poolId: poolId,
      });
    }

    const collection = await this.config.indexerProvider.getCollection(
      pool.config.token0,
    );

    if (!collection) {
      throw new PoolErrorException({
        message: `Collection ${pool.config.token0} does not exist`,
        type: PoolErrorType.InvalidToken,
        token: pool.config.token0,
      });
    }

    let walletUtxos = [...wallet.utxos, ...(paymentWallet?.utxos ?? [])];
    let splitRuneInputs: Array<string> = [];

    const btcAmount = BigInt(amountOut);
    if (btcAmount < DUST_LIMIT) {
      throw new PoolErrorException({
        message: `BTC amount is below the dust limit. Please increase the bitcoin amount to ${DUST_LIMIT} sats.`,
        type: PoolErrorType.InvalidAmountBelowMin,
        token: pool.config.token1,
        minAmount: DUST_LIMIT.toString(),
      });
    }

    const splitRuneUtxo = await validateRuneToBtcPsbt(
      wallet,
      paymentWallet,
      collection,
      BigInt(amountIn),
      BigInt(amountOut),
      exactIn,
      pool,
      BigInt(feeRate),
      signedPsbt,
      this.config.bitcoinProvider,
      scureNetwork,
      DUST_LIMIT,
      splitRunePsbt,
    );
    walletUtxos = [...walletUtxos, splitRuneUtxo.utxo];
    splitRuneInputs = splitRuneUtxo.splitRuneInputs;

    // Get utxo info.
    const utxosInfo = buildUtxoInfoFromInputs(signedPsbt, walletUtxos);

    const transaction = Transaction.fromPSBT(base64.decode(signedPsbt));
    finalizeTransaction(transaction);

    const params: SwapParams = {
      exact_in: exactIn,
      // frontend_fee: {
      //   address: this.feeAddress,
      //   fee: 0,
      // },
      frontend_fee: null,
      transaction: transaction.hex,
      zero_to_one: true,
    };

    const instruction: SwapInstruction = {
      params,
    };

    const shards = pool.shards.map((shard) => shard.pubkey);
    const message = swapMessage(
      this.config.programAccount,
      hex.encode(toXOnly(hex.decode(runePublicKey))),
      hex.encode(toXOnly(hex.decode(request.feePayerPubkey))),
      this.config.mempoolInfoOracleAccount,
      this.config.feeRateOracleAccount,
      poolId,
      shards,
      shards,
      instruction,
      request.recentBlockhash,
      hex.encode(createProtocolPda(hex.decode(this.config.programAccount))[0]),
    );

    return {
      message,
      utxosInfo,
      splitRuneInputs,
    };
  }
}