import {
  DEFAULT_ARCH_P2TR_INPUT,
  DEFAULT_P2TR_OUTPUT,
  EstimateError,
  getBitcoinNetwork,
  getOutputTypeFromAddress,
  Psbt,
  PsbtBuilder,
  toVsize,
  toXOnly,
  TransactionSizeCalculator,
} from '@saturnbtcio/psbt';
import { ArchWalletSdkConfig } from '../arch-wallet-sdk';
import { validateArchWalletData } from '../util/validation';
import {
  BTC_TOKEN,
  WithdrawFundsMessageRequest,
  DepositBitcoinToArchControlledAddressPsbtRequest,
  DepositRunesToArchControlledAddressPsbtRequest,
} from './arch-wallet.dto';
import { hex } from '@scure/base';
import { Collection, CollectionUtxo, Wallet } from '../wallet/wallet.dto';
import {
  PubkeyUtil,
  SanitizedMessage,
  UtxoMetaData,
} from '@saturnbtcio/arch-sdk';
import {
  ArchWalletErrorException,
  ArchWalletErrorType,
} from '../error/arch-wallet.error';
import {
  WithdrawFundsInstruction,
  withdrawArchWalletFundsMessage,
} from '@saturnbtcio/arch-wallet-serde-sdk';
import { findBtc, findCollection } from '../util/finder';
import { getTxIdsFromUtxos } from '../util/mempool';
import { SigHash } from '@scure/btc-signer';
import { DUST_LIMIT } from '../util/constants';

export class ArchWalletManager {
  private readonly config: ArchWalletSdkConfig;

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

  async depositBtcFundsToArchControlledAddress(
    request: DepositBitcoinToArchControlledAddressPsbtRequest,
  ): Promise<Psbt<CollectionUtxo>> {
    const wallet = await this.config.bitcoinProvider.getWallet(
      request.paymentAddress ?? request.runeAddress,
    );

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

    const archAddress = await this.config.archProvider.getAccountAddress(
      PubkeyUtil.fromHex(
        hex.encode(toXOnly(hex.decode(request.runePublicKey))),
      ),
    );

    txBuilder.addOutput({
      address: archAddress,
      value: request.amount,
      runes: [],
    });

    const walletMempoolStatus =
      await this.config.bitcoinProvider.getMempoolInfo(
        getTxIdsFromUtxos(wallet.utxos),
      );

    const btcUtxos = findBtc(
      wallet,
      walletMempoolStatus,
      [],
      new Map(),
      request.amount,
      request.feeRate,
      txBuilder,
      SigHash.ALL,
      request.paymentPublicKey ?? request.runePublicKey,
      DUST_LIMIT,
    );

    try {
      const tx = txBuilder.buildAndAdjustChange({
        feeRate: BigInt(request.feeRate),
        dustLimit: DUST_LIMIT,
        changeAddress: request.paymentAddress ?? request.runeAddress,
        additionalInputAmount: 0n,
      });

      return tx;
    } 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 ArchWalletErrorException({
        type: ArchWalletErrorType.NotEnoughFunds,
        message: `Not enough bitcoin. You need at least ${amount} in your wallet at current fee rate.`,
        maxAmount: btcUtxos.amount.toString(),
        minAmount: amount.toString(),
        token: `${BTC_TOKEN.block}:${BTC_TOKEN.tx}`,
      });
    }
  }

  async depositRunesToArchControlledAddress(
    request: DepositRunesToArchControlledAddressPsbtRequest,
  ): Promise<Psbt<CollectionUtxo>> {
    const wallet = await this.config.bitcoinProvider.getWallet(
      request.runeAddress,
    );

    const walletMempoolStatus =
      await this.config.bitcoinProvider.getMempoolInfo(
        getTxIdsFromUtxos(wallet.utxos),
      );

    const collection: Collection = {
      id: request.runeId,
      type: 'rune',
    };

    const collectionUtxos = findCollection(
      wallet,
      walletMempoolStatus,
      collection,
      request.amount,
    );

    if (collectionUtxos.amount < request.amount) {
      throw new ArchWalletErrorException({
        type: ArchWalletErrorType.NotEnoughFunds,
        message: `Not enough runes. You need at least ${request.amount} in your wallet.`,
        maxAmount: collectionUtxos.amount.toString(),
        minAmount: request.amount.toString(),
        token: request.runeId,
      });
    }

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

    const archAddress = await this.config.archProvider.getAccountAddress(
      PubkeyUtil.fromHex(
        hex.encode(toXOnly(hex.decode(request.runePublicKey))),
      ),
    );

    // Add collection inputs
    for (const utxo of collectionUtxos.utxos) {
      txBuilder.addInput({
        utxo,
        owner: wallet.address,
        publicKey: request.runePublicKey,
        mempoolStatus: walletMempoolStatus.get(utxo.txid),
        sighashType: SigHash.ALL,
      });
    }

    // Return remaining funds if there are other runes or if we're sending less than total
    if (
      collectionUtxos.amount < collectionUtxos.totalAmount ||
      collectionUtxos.containsOtherRunes
    ) {
      txBuilder.addOutput({
        address: request.paymentAddress ?? request.runeAddress,
        value: DUST_LIMIT,
        runes: [],
      });
    }

    // Add output to arch address with runes
    txBuilder.addOutput({
      address: archAddress,
      value: DUST_LIMIT,
      runes: [
        {
          amount: collectionUtxos.amount,
          id: request.runeId,
        },
      ],
    });

    // Calculate total output value (dust limits for outputs)
    const totalOutputValue = txBuilder.outputs.reduce(
      (sum: bigint, output) => sum + output.value,
      0n,
    );

    // Chosen wallet for fees.
    const chosenWallet = request.paymentAddress
      ? await this.config.bitcoinProvider.getWallet(request.paymentAddress)
      : wallet;

    const chosenWalletMempoolStatus = request.paymentAddress
      ? await this.config.bitcoinProvider.getMempoolInfo(
          getTxIdsFromUtxos(chosenWallet.utxos),
        )
      : walletMempoolStatus;

    // Find BTC to cover fees and outputs
    const btcUtxos = findBtc(
      chosenWallet,
      chosenWalletMempoolStatus,
      collectionUtxos.utxos,
      walletMempoolStatus,
      totalOutputValue,
      request.feeRate,
      txBuilder,
      SigHash.ALL,
      request.paymentPublicKey ?? request.runePublicKey,
      DUST_LIMIT,
    );

    try {
      const tx = txBuilder.buildAndAdjustChange({
        feeRate: BigInt(request.feeRate),
        dustLimit: DUST_LIMIT,
        changeAddress: request.paymentAddress ?? request.runeAddress,
        additionalInputAmount: 0n,
      });

      return tx;
    } 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 ArchWalletErrorException({
        type: ArchWalletErrorType.NotEnoughFunds,
        message: `Not enough bitcoin. You need at least ${amount} in your wallet at current fee rate.`,
        maxAmount: btcUtxos.amount.toString(),
        minAmount: amount.toString(),
        token: `${BTC_TOKEN.block}:${BTC_TOKEN.tx}`,
      });
    }
  }

  async withdrawFundsMessage(
    request: WithdrawFundsMessageRequest,
  ): Promise<SanitizedMessage> {
    const scureNetwork = getBitcoinNetwork(this.config.network);

    // Validate the request
    validateArchWalletData<WithdrawFundsMessageRequest>(request, scureNetwork);

    // Fetch the UTXOs that will be used to process the recover.
    const archWallet = await this.getArchWallet(request.runePublicKey);
    const utxos = await this.getUtxosFromArchWallet(archWallet, request);

    // Validate that there are enough funds in the user's wallet to process to the recover.
    this.validateEnoughFunds(request, utxos);

    // Create the recover funds instruction
    const withdrawFundsInstruction: WithdrawFundsInstruction = {
      utxos,
      runeAddress: request.runeAddress,
      paymentAddress: request.paymentAddress,
    };

    const message = withdrawArchWalletFundsMessage(
      this.config.programAddress,
      hex.encode(toXOnly(hex.decode(request.runePublicKey))),
      hex.encode(toXOnly(hex.decode(request.feePayerPubkey))),
      this.config.mempoolInfoOracleAccount,
      this.config.feeRateOracleAccount,
      withdrawFundsInstruction,
      request.recentBlockhash,
    );

    return message;
  }

  private validateEnoughFunds(
    request: WithdrawFundsMessageRequest,
    utxos: CollectionUtxo[],
  ) {
    const txSize = this.estimateRecoverFundsTxSize(
      request.runeAddress,
      request.paymentAddress,
      utxos,
      false,
    );

    const vsize = toVsize(txSize.weight);
    const fee = BigInt(vsize) * request.feeRate;

    const totalFunds = utxos.reduce((acc, utxo) => {
      return acc + utxo.value;
    }, BigInt(0));

    if (totalFunds < fee) {
      const newTxSize = this.estimateRecoverFundsTxSize(
        request.runeAddress,
        request.paymentAddress,
        utxos,
        true,
      );

      const newFee = BigInt(toVsize(newTxSize.weight)) * request.feeRate;

      throw new ArchWalletErrorException({
        type: ArchWalletErrorType.NotEnoughFunds,
        message: `Not enough bitcoin in the arch wallet to pay for the bitcoin transaction. Send an additional amount of ${(newFee - totalFunds).toString()} sats to the arch wallet.`,
        minAmount: newFee.toString(),
        maxAmount: Number.MAX_SAFE_INTEGER.toString(),
        token: `${BTC_TOKEN.block}:${BTC_TOKEN.tx}`,
      });
    }
  }

  private estimateRecoverFundsTxSize(
    runeAddress: string,
    paymentAddress: string | null,
    utxos: CollectionUtxo[],
    additionalInputUtxo: boolean,
  ) {
    const scureNetwork = getBitcoinNetwork(this.config.network);
    const txSizeCalculator = new TransactionSizeCalculator();

    // Fee payer
    txSizeCalculator.addInput(DEFAULT_ARCH_P2TR_INPUT);
    txSizeCalculator.addOutput(DEFAULT_P2TR_OUTPUT);

    for (const utxo of utxos) {
      txSizeCalculator.addInput(DEFAULT_ARCH_P2TR_INPUT);
    }

    if (additionalInputUtxo) {
      txSizeCalculator.addInput(DEFAULT_ARCH_P2TR_INPUT);
    }

    const areThereRunes = utxos.some((utxo) =>
      utxo.collectionStatuses.some((status) => status.type !== 'btc'),
    );

    let btcOutput = getOutputTypeFromAddress(
      paymentAddress ?? runeAddress,
      scureNetwork,
    );

    txSizeCalculator.addOutput(btcOutput);

    if (areThereRunes) {
      const runeOutput = getOutputTypeFromAddress(runeAddress, scureNetwork);
      txSizeCalculator.addOutput(runeOutput);
    }

    return txSizeCalculator.estimateSize({
      network: scureNetwork,
    });
  }

  private async getUtxosFromArchWallet(
    archWallet: Wallet,
    request: WithdrawFundsMessageRequest,
  ): Promise<CollectionUtxo[]> {
    switch (request.type) {
      case 'selected-utxos':
        return this.getSelectedUtxosFromArchWallet(request.utxos, archWallet);
      case 'all-utxos':
        return archWallet.utxos;
    }
  }

  private getSelectedUtxosFromArchWallet(
    utxos: UtxoMetaData[],
    wallet: Wallet,
  ) {
    let collectionUtxos: CollectionUtxo[] = [];
    for (const utxo of utxos) {
      const utxoInWallet = wallet.utxos.find(
        (utxoInWallet) =>
          utxoInWallet.txid === utxo.txid && utxoInWallet.vout === utxo.vout,
      );

      if (!utxoInWallet) {
        throw new ArchWalletErrorException({
          type: ArchWalletErrorType.InvalidUtxo,
          message: 'Some of the selected UTXOs are not in the arch wallet',
          utxos: [`${utxo.txid}:${utxo.vout}`],
        });
      }

      collectionUtxos.push(utxoInWallet);
    }

    return collectionUtxos;
  }

  private async getArchWallet(publicKey: string): Promise<Wallet> {
    const pubkey = hex.decode(publicKey);
    const archAddress =
      await this.config.archProvider.getAccountAddress(pubkey);
    const archWallet = await this.config.bitcoinProvider.getWallet(archAddress);
    return archWallet;
  }
}