import { PubkeyUtil } from '@saturnbtcio/arch-sdk';
import { Runestone } from '@saturnbtcio/ordinals-lib';
import { IdentifiableLiquidityPool } from '@saturnbtcio/pool-serde-sdk';
import {
  containsOpReturn,
  DEFAULT_ARCH_P2TR_INPUT,
  estimateInput,
  finalizeTransaction,
  toVsize,
  toXOnly,
  TransactionSizeCalculator,
} from '@saturnbtcio/psbt';
import { base64, hex } from '@scure/base';
import {
  Address,
  NETWORK,
  OutScript,
  SigHash,
  Transaction,
} from '@scure/btc-signer';
import { PoolErrorException, PoolErrorType } from '../../error/pool.error';
import { IArchProvider } from '../../providers/arch.provider';
import {
  IBitcoinProvider,
  MempoolEntry,
  MempoolInfoMap,
} from '../../providers/bitcoin.provider';
import {
  checkFeeRate,
  getFeeForSwapBtcToRune,
  getTotalFeeAndSizeFromAncestorsInputs,
  getUtxosFromTxInputs,
  checkFee,
} from '../../util/fee';
import { validateAmountsAreWithinThreshold } from '../../util/shards';
import { verifyUtxosInPsbtNotUsedAndStillOwnedByUser } from '../../util/validation';
import { Collection, CollectionUtxo, Wallet } from '../../wallet/wallet.dto';
import { getTxIdsFromUtxos } from '../../util/mempool';

export const validateBtcToRunePsbt = async (
  runeWallet: Wallet,
  paymentsWallet: Wallet | undefined,
  collection: Collection,
  pool: IdentifiableLiquidityPool,
  amountIn: bigint,
  amountOut: bigint,
  feeRate: bigint,
  psbtBase64: string,
  feePayerPubkey: string,
  shardPubkeys: string[],
  network: typeof NETWORK,
  dustLimit: bigint,
  archProvider: IArchProvider,
  bitcoinProvider: IBitcoinProvider,
  programAddress: string,
  mergeUtxoPsbtBase64?: string,
): Promise<{ additionalUtxos: CollectionUtxo[] }> => {
  const tx = Transaction.fromPSBT(base64.decode(psbtBase64));

  assertAllInputsHaveSighashAnyoneCanPay(tx);

  let utxos: CollectionUtxo[];
  let overrideUtxosMempoolInfo: MempoolInfoMap | undefined = undefined;
  let additionalUtxos: CollectionUtxo[] = [];
  if (mergeUtxoPsbtBase64) {
    const { btcUtxo, utxosMempoolInfo } = await validateBtcMergeUtxoPsbt(
      runeWallet,
      paymentsWallet,
      mergeUtxoPsbtBase64,
      dustLimit,
      bitcoinProvider,
      network,
    );
    utxos = [btcUtxo];
    overrideUtxosMempoolInfo = utxosMempoolInfo;
    additionalUtxos = [btcUtxo];
  } else {
    utxos = await getVerifiedUtxosForWallet(
      psbtBase64,
      paymentsWallet ?? runeWallet,
    );
  }

  const { shards, totalRuneAmountInShards } = resolveShardsOrThrow(
    pool,
    shardPubkeys,
  );

  const runestone = parseRunestoneOrThrow(psbtBase64);

  const { programRuneOutputsIndexes, userEdict } =
    validateRunestoneEdictsAndAmounts(
      runestone,
      collection.id,
      amountOut,
      totalRuneAmountInShards,
    );

  validateUserRuneOutputAddress(
    tx,
    userEdict.output,
    runeWallet.address,
    network,
  );

  validateShardPubkeysAndOutputsCount(tx, shardPubkeys);

  const { feePayerAddress, shardAddresses } =
    await resolveFeePayerAndShardAddresses(
      archProvider,
      feePayerPubkey,
      shardPubkeys,
    );
  const startIndex = validateFeePayerAndShardOutputs(
    tx,
    feePayerAddress,
    shardAddresses,
    dustLimit,
    network,
  );

  const updatedProgramRuneOutputsIndexes = validatePointerAndDefaultRuneOutput(
    tx,
    runestone,
    programAddress,
    dustLimit,
    network,
    programRuneOutputsIndexes,
  );

  const { totalAmountBtcOut } = validateLpAndBtcOutputs(
    tx,
    programAddress,
    amountIn,
    feeRate,
    dustLimit,
    updatedProgramRuneOutputsIndexes,
    startIndex,
    network,
  );

  finalizeSignedTransactionOrThrow(tx);

  await computeAndValidateEffectiveFeeRate(
    tx,
    utxos,
    shardPubkeys.length,
    shardPubkeys.length,
    dustLimit,
    bitcoinProvider,
    network,
    totalAmountBtcOut,
    overrideUtxosMempoolInfo,
  );
  return { additionalUtxos };
};

function assertAllInputsHaveSighashAnyoneCanPay(tx: Transaction) {
  for (let i = 0; i < tx.inputsLength; i++) {
    const input = tx.getInput(i);
    if (input.sighashType !== SigHash.ALL_ANYONECANPAY) {
      throw new PoolErrorException({
        message: `Invalid psbt. Tx should have all anyone can pay sighash type.`,
        type: PoolErrorType.InvalidPsbt,
      });
    }
  }
}

async function getVerifiedUtxosForWallet(
  psbtBase64: string,
  wallet: Wallet,
): Promise<CollectionUtxo[]> {
  const inputs = await verifyUtxosInPsbtNotUsedAndStillOwnedByUser(psbtBase64, [
    wallet,
  ]);

  const utxos = inputs.map((input) => {
    const utxo = wallet.utxos.find((u) => `${u.txid}:${u.vout}` === input);
    if (!utxo) {
      throw new PoolErrorException({
        message: `Invalid psbt. Utxo not found.`,
        type: PoolErrorType.InvalidPsbt,
      });
    }
    return utxo;
  });

  return utxos;
}

function resolveShardsOrThrow(
  pool: IdentifiableLiquidityPool,
  shardPubkeys: string[],
) {
  const shards = shardPubkeys.map((pubkey) => {
    const shard = pool.shards.find((s) => s.pubkey === pubkey);
    if (!shard) {
      throw new PoolErrorException({
        message: `Invalid psbt. Shard not found.`,
        type: PoolErrorType.InvalidPsbt,
      });
    }
    return shard;
  });

  const totalRuneAmountInShards = shards.reduce((acc, shard) => {
    if (!shard.runeUtxo) {
      throw new Error('Shard has no rune utxo');
    }
    if (shard.runeUtxo.runes.length === 0) {
      return acc;
    }
    return acc + shard.runeUtxo.runes[0].amount;
  }, 0n);

  return { shards, totalRuneAmountInShards };
}

function parseRunestoneOrThrow(psbtBase64: string) {
  const runestone = Runestone.decipher(psbtBase64);
  if (!runestone) {
    throw new PoolErrorException({
      message: `Invalid psbt. Runestone not found.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }
  if (runestone.edicts.length === 0) {
    throw new PoolErrorException({
      message: `Invalid psbt. Runestone should have at least one edict.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }
  return runestone;
}

export function validateRunestoneEdictsAndAmounts(
  runestone: Runestone,
  collectionId: string,
  amountOut: bigint,
  totalRuneAmountInShards: bigint,
) {
  const programEdicts =
    runestone.edicts.length > 1
      ? runestone.edicts.slice(0, runestone.edicts.length - 1)
      : [];

  const programRuneAmounts = programEdicts.map((edict: any) => edict.amount);
  const totalProgramRuneAmountInEdicts = programRuneAmounts.reduce(
    (acc: bigint, amount: bigint) => acc + amount,
    0n,
  );

  if (totalRuneAmountInShards < totalProgramRuneAmountInEdicts) {
    throw new PoolErrorException({
      message: `Invalid psbt. Runestone should not move more than ${totalRuneAmountInShards} rune.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }

  if (
    totalRuneAmountInShards - totalProgramRuneAmountInEdicts - amountOut <
    0n
  ) {
    throw new PoolErrorException({
      message: `Invalid psbt. Runestone should not move more than ${totalRuneAmountInShards} rune.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }

  programRuneAmounts.push(
    totalRuneAmountInShards - totalProgramRuneAmountInEdicts - amountOut,
  );

  validateAmountsAreWithinThreshold(programRuneAmounts, 50n);

  const programRuneOutputsIndexes = programEdicts.map(
    (edict: any) => edict.output,
  );

  const userEdict = runestone.edicts[runestone.edicts.length - 1];

  if (userEdict.amount !== amountOut) {
    throw new PoolErrorException({
      message: `Invalid psbt. Runestone should move ${amountOut} rune.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }

  if (`${userEdict.id.block}:${userEdict.id.tx}` !== collectionId) {
    throw new PoolErrorException({
      message: `Invalid psbt. Runestone should move ${collectionId} rune.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }

  return { programRuneOutputsIndexes, userEdict };
}

function validateUserRuneOutputAddress(
  tx: Transaction,
  userOutputIndex: number,
  expectedAddress: string,
  network: typeof NETWORK,
) {
  const userRuneOutput = tx.getOutput(userOutputIndex);
  const userRuneAddress = Address(network).encode(
    OutScript.decode(userRuneOutput.script!),
  );

  if (userRuneAddress !== expectedAddress) {
    throw new PoolErrorException({
      message: `Invalid psbt. Output should be to your wallet.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }
}

function validateShardPubkeysAndOutputsCount(
  tx: Transaction,
  shardPubkeys: string[],
) {
  if (tx.outputsLength <= shardPubkeys.length) {
    throw new PoolErrorException({
      message: `Invalid psbt. There should be a minimum of ${shardPubkeys.length} outputs. Currently there are ${tx.outputsLength}.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }

  if (shardPubkeys.length === 0) {
    throw new PoolErrorException({
      message: `Invalid psbt. No shard pubkeys provided.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }
}

async function resolveFeePayerAndShardAddresses(
  archProvider: IArchProvider,
  feePayerPubkey: string,
  shardPubkeys: string[],
) {
  const allPubkeys = [
    hex.encode(toXOnly(hex.decode(feePayerPubkey))),
    ...shardPubkeys,
  ];
  const addresses = await Promise.all(
    allPubkeys.map((pk) =>
      archProvider.getAccountAddress(PubkeyUtil.fromHex(pk)),
    ),
  );
  const [feePayerAddress, ...shardAddresses] = addresses;
  return { feePayerAddress, shardAddresses };
}

function validateFeePayerAndShardOutputs(
  tx: Transaction,
  feePayerAddress: string,
  shardAddresses: string[],
  dustLimit: bigint,
  network: typeof NETWORK,
) {
  let shardOutputsFound = 0;
  let feePayerFound = false;

  for (let i = 0; i < tx.outputsLength; i++) {
    if (feePayerFound && shardOutputsFound === shardAddresses.length) {
      break;
    }

    const output = tx.getOutput(i);
    const address = Address(network).encode(OutScript.decode(output.script!));

    if (address === feePayerAddress) {
      if (output.amount !== dustLimit) {
        throw new PoolErrorException({
          message: `Invalid psbt. Fee payer should have ${dustLimit} sats.`,
          type: PoolErrorType.InvalidPsbt,
        });
      }
      feePayerFound = true;
      continue;
    }

    if (shardAddresses.includes(address)) {
      if (output.amount !== dustLimit) {
        throw new PoolErrorException({
          message: `Invalid psbt. Shard should have ${dustLimit} sats.`,
          type: PoolErrorType.InvalidPsbt,
        });
      }
      shardOutputsFound++;
    }
  }

  if (!feePayerFound) {
    throw new PoolErrorException({
      message: `Invalid psbt. Fee payer output not found.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }

  if (shardOutputsFound !== shardAddresses.length) {
    throw new PoolErrorException({
      message: `Invalid psbt. Shard output not found.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }

  return 1 + shardAddresses.length;
}

function validatePointerAndDefaultRuneOutput(
  tx: Transaction,
  runestone: any,
  programAddress: string,
  dustLimit: bigint,
  network: typeof NETWORK,
  programRuneOutputsIndexes: number[],
) {
  if (!runestone.pointer) {
    throw new PoolErrorException({
      message: `Invalid psbt. Runestone should have pointer.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }
  programRuneOutputsIndexes.push(runestone.pointer);

  const defaultRuneOutput = tx.getOutput(runestone.pointer);

  if (defaultRuneOutput.amount !== dustLimit) {
    throw new PoolErrorException({
      message: `Invalid psbt. Output should be ${dustLimit} rune.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }

  const defaultRuneAddress = Address(network).encode(
    OutScript.decode(defaultRuneOutput.script!),
  );

  if (defaultRuneAddress !== programAddress) {
    throw new PoolErrorException({
      message: `Invalid psbt. Output should be to lp.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }

  return programRuneOutputsIndexes;
}

function validateLpAndBtcOutputs(
  tx: Transaction,
  programAddress: string,
  amountIn: bigint,
  feeRate: bigint,
  dustLimit: bigint,
  programRuneOutputsIndexes: number[],
  startIndex: number,
  network: typeof NETWORK,
) {
  let lpOutputExists = false;
  let totalAmountBtcOut = 0n;
  for (let i = startIndex; i < tx.outputsLength; i++) {
    const output = tx.getOutput(i);
    const amount = output.amount ?? 0n;
    totalAmountBtcOut += amount;

    const script = output.script;
    if (!script || containsOpReturn(script)) {
      continue;
    }

    const address = Address(network).encode(OutScript.decode(script));

    if (programRuneOutputsIndexes.includes(i)) {
      if (amount !== dustLimit) {
        throw new PoolErrorException({
          message: `Invalid psbt. Outpus should have ${dustLimit} btc for a rune output.`,
          type: PoolErrorType.InvalidPsbt,
        });
      }
    } else if (address === programAddress) {
      const p2trInputSize = estimateInput(DEFAULT_ARCH_P2TR_INPUT);
      const p2trInputCost = BigInt(
        Math.floor(toVsize(p2trInputSize.weight) * Number(feeRate)),
      );

      if (amountIn + p2trInputCost !== output.amount) {
        throw new PoolErrorException({
          message: `Invalid psbt. Output to lp should be ${amountIn} btc.`,
          type: PoolErrorType.InvalidPsbt,
        });
      }

      lpOutputExists = true;
    }
  }

  if (!lpOutputExists) {
    throw new PoolErrorException({
      message: `Invalid psbt. Lp btc output not found.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }

  return { totalAmountBtcOut };
}

function finalizeSignedTransactionOrThrow(tx: Transaction) {
  try {
    finalizeTransaction(tx);
  } catch {
    throw new PoolErrorException({
      message: `Invalid psbt. Tx has not been signed.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }
}

async function computeAndValidateEffectiveFeeRate(
  tx: Transaction,
  utxos: CollectionUtxo[],
  shardsCount: number,
  shardOutputs: number,
  dustLimit: bigint,
  bitcoinProvider: IBitcoinProvider,
  network: typeof NETWORK,
  totalAmountBtcOut: bigint,
  overrideUtxosMempoolStatuses?: MempoolInfoMap,
) {
  const additionalBtcIn = BigInt(shardsCount) * dustLimit * 2n;
  const totalAmountBtcIn =
    utxos.reduce((acc, utxo) => acc + utxo.value, 0n) + additionalBtcIn;

  const txSizeCalculator = TransactionSizeCalculator.createFromTransaction(tx);
  getFeeForSwapBtcToRune(txSizeCalculator, shardOutputs);

  const utxosFromInputs = getUtxosFromTxInputs(tx, utxos);
  const utxosMempoolStatuses = overrideUtxosMempoolStatuses
    ? overrideUtxosMempoolStatuses
    : await bitcoinProvider.getMempoolInfo(getTxIdsFromUtxos(utxosFromInputs));
  const { totalSize, totalFee } = getTotalFeeAndSizeFromAncestorsInputs(
    utxosFromInputs,
    utxosMempoolStatuses,
  );

  const fee = totalAmountBtcIn - totalAmountBtcOut + totalFee;
  const size =
    toVsize(
      txSizeCalculator.estimateSize({
        network: network,
      }).weight,
    ) + totalSize;
  const txFeeRate = fee / BigInt(Math.ceil(size));
  await checkFeeRate(Number(txFeeRate), bitcoinProvider);
}

async function validateBtcMergeUtxoPsbt(
  runeWallet: Wallet,
  paymentsWallet: Wallet | undefined,
  psbtBase64: string,
  dustLimit: bigint,
  bitcoinProvider: IBitcoinProvider,
  network: typeof NETWORK,
): Promise<{ btcUtxo: CollectionUtxo; utxosMempoolInfo: MempoolInfoMap }> {
  const psbt = Transaction.fromPSBT(base64.decode(psbtBase64));

  // Ensure inputs belong to the user (rune or payment wallet)
  const inputs = await verifyUtxosInPsbtNotUsedAndStillOwnedByUser(psbtBase64, [
    paymentsWallet ?? runeWallet,
  ]);

  // Merge PSBT should not contain runestones
  const maybeRunestone = Runestone.decipher(psbtBase64);
  if (maybeRunestone) {
    throw new PoolErrorException({
      message: `Invalid merge psbt. It should not contain a runestone`,
      type: PoolErrorType.InvalidPsbt,
    });
  }

  // Expect exactly one output back to the same chosen wallet
  if (psbt.outputsLength !== 1) {
    throw new PoolErrorException({
      message: `Invalid merge psbt. Expected 1 output`,
      type: PoolErrorType.InvalidPsbt,
    });
  }

  const output = psbt.getOutput(0);
  const address = Address(network).encode(OutScript.decode(output.script!));
  const expectedAddress = (paymentsWallet ?? runeWallet).address;
  if (address !== expectedAddress) {
    throw new PoolErrorException({
      message: `Invalid merge psbt. Output not to expected address`,
      type: PoolErrorType.InvalidPsbt,
    });
  }

  if ((output.amount ?? 0n) < dustLimit) {
    throw new PoolErrorException({
      message: `Invalid merge psbt. Output below dust`,
      type: PoolErrorType.InvalidPsbt,
    });
  }

  try {
    psbt.finalize();
  } catch (err) {
    throw new PoolErrorException({
      message: `Invalid merge psbt. Psbt has not been signed.`,
      type: PoolErrorType.InvalidPsbt,
    });
  }

  // Build synthetic mempool info for the merged utxo's parent tx (this merge)
  const { ancestorsCount, fee, size } = await checkFee(
    psbt,
    inputs.map((i) => {
      const [txid, voutStr] = i.split(':');
      const vout = Number(voutStr);
      const src = [...runeWallet.utxos, ...(paymentsWallet?.utxos ?? [])].find(
        (u) => u.txid === txid && u.vout === vout,
      );
      if (!src) {
        throw new PoolErrorException({
          message: `Invalid merge psbt. Input not found in wallet`,
          type: PoolErrorType.InvalidPsbt,
        });
      }
      return src;
    }),
    bitcoinProvider,
  );

  const btcUtxoMempoolInfo: MempoolEntry = {
    fees: { ancestor: Number(fee), modified: 0, base: 0 },
    ancestorsCount: ancestorsCount + 1,
    descendantsCount: 0,
    ancestorsSize: size,
    descendantsSize: 0,
    depends: [],
    spentby: [],
  };
  const utxosMempoolInfo: MempoolInfoMap = new Map();
  utxosMempoolInfo.set(psbt.id, btcUtxoMempoolInfo);

  const btcUtxo: CollectionUtxo = {
    txid: psbt.id,
    vout: 0,
    value: output.amount ?? 0n,
    collectionStatuses: [],
    hasInscription: false,
    status: { confirmed: false },
  };

  return { btcUtxo, utxosMempoolInfo };
}