// unilateral-exit.ts

import type { Logger } from "@lightsparkdev/core";
import { bytesToHex, hexToBytes } from "@noble/curves/utils";
import { ripemd160 } from "@noble/hashes/legacy";
import { sha256 } from "@noble/hashes/sha2";
import * as btc from "@scure/btc-signer";
import type * as psbt from "@scure/btc-signer/psbt";
import type { SparkServiceClient } from "../proto/spark.js";
import {
  TreeNode,
  TreeNodeStatus,
  treeNodeStatusToJSON,
} from "../proto/spark.js";
import { getTxFromRawTxHex, getTxId } from "../utils/bitcoin.js";
import { isTxBroadcast } from "../utils/mempool.js";
import { Network, NetworkToProto } from "../utils/network.js";

// Types
export interface LeafInfo {
  leafId: string;
  nodeTx: string; // raw tx hex
  refundTx: string; // raw tx hex
  // Add other fields as needed
}

export interface Utxo {
  txid: string;
  vout: number;
  value: bigint;
  script: string;
  publicKey: string;
}

export interface FeeRate {
  satPerVbyte: number;
}

export interface FeeBumpTxPackage {
  tx: string;
  feeBumpPsbt?: string;
}

export interface FeeBumpTxChain {
  leafId: string;
  txPackages: FeeBumpTxPackage[];
}

export interface TxChain {
  leafId: string;
  transactions: string[];
}

export interface BroadcastConfig {
  bitcoinCoreRpcUrl?: string;
  rpcUsername?: string;
  rpcPassword?: string;
  autoBroadcast?: boolean;
  network?: "MAINNET" | "REGTEST" | "TESTNET" | "SIGNET" | "LOCAL";
}

export interface BroadcastResult {
  success: boolean;
  txids?: string[];
  error?: string;
  broadcastedPackages?: number;
}

const TREE_NODE_STATUS_PREFIX = "TREE_NODE_STATUS_";

function getTreeNodeStatusString(status: TreeNodeStatus): string {
  const statusName = treeNodeStatusToJSON(status);
  if (!statusName.startsWith(TREE_NODE_STATUS_PREFIX)) {
    throw new Error(`Unexpected tree node status: ${statusName}`);
  }
  return statusName.slice(TREE_NODE_STATUS_PREFIX.length);
}

const EXIT_CHAIN_STATUSES = new Set([
  getTreeNodeStatusString(TreeNodeStatus.TREE_NODE_STATUS_AVAILABLE),
  getTreeNodeStatusString(TreeNodeStatus.TREE_NODE_STATUS_SPLITTED),
  getTreeNodeStatusString(TreeNodeStatus.TREE_NODE_STATUS_SPLIT_LOCKED),
  getTreeNodeStatusString(TreeNodeStatus.TREE_NODE_STATUS_ON_CHAIN),
]);

function formatErrorForLog(error: unknown): string {
  if (error instanceof Error) {
    return error.message;
  }
  if (typeof error === "string") {
    return error;
  }
  if (typeof error === "number" || typeof error === "boolean") {
    return String(error);
  }
  return "unknown error";
}

function warnUnilateralExit(logger: Logger | undefined, message: string) {
  logger?.warn(`constructUnilateralExitFeeBumpPackages: ${message}`);
}

export async function buildUnilateralExitChain(
  node: TreeNode,
  nodeMap: Map<string, TreeNode>,
  sparkClient?: SparkServiceClient,
  network: Network = Network.MAINNET,
): Promise<TreeNode[]> {
  const chain: TreeNode[] = [];
  let currentNode = node;

  while (currentNode) {
    // Only proceed with nodes that can still contribute to an exit chain.
    // ON_CHAIN is included here because the server marks a node ON_CHAIN when
    // either its raw tx or its direct tx confirms.
    if (!EXIT_CHAIN_STATUSES.has(currentNode.status)) {
      break;
    }
    chain.unshift(currentNode);

    if (!currentNode.parentNodeId) {
      break;
    }

    let parentNode = nodeMap.get(currentNode.parentNodeId);

    if (!parentNode && sparkClient) {
      try {
        const response = await sparkClient.query_nodes({
          source: {
            $case: "nodeIds",
            nodeIds: {
              nodeIds: [currentNode.parentNodeId],
            },
          },
          includeParents: true,
          network: NetworkToProto[network],
        });

        for (const returnedNode of Object.values(response.nodes)) {
          nodeMap.set(returnedNode.id, returnedNode);
        }

        parentNode = nodeMap.get(currentNode.parentNodeId);
        if (!parentNode) {
          throw new Error(
            `Parent node ${currentNode.parentNodeId} not returned by query_nodes. Exit chain is incomplete.`,
          );
        }
      } catch (error) {
        throw new Error(
          `Failed to query parent node ${currentNode.parentNodeId}: ${formatErrorForLog(error)}`,
          { cause: error },
        );
      }
    }

    if (parentNode) {
      currentNode = parentNode;
      continue;
    }
    break;
  }

  return chain;
}

export function isEphemeralAnchorOutput(
  script?: Uint8Array,
  amount?: bigint,
): boolean {
  return Boolean(
    amount === 0n &&
    script &&
    // Pattern 1: Bare OP_TRUE (single byte 0x51)
    ((script.length === 1 && script[0] === 0x51) ||
      // Pattern 2: Push OP_TRUE (two bytes 0x01 0x51) - MALFORMED but we detect it
      (script.length === 2 && script[0] === 0x01 && script[1] === 0x51) ||
      // Pattern 3: Bitcoin v29 ephemeral anchor script (7 bytes: 015152014e0173)
      (script.length === 7 &&
        script[0] === 0x01 &&
        script[1] === 0x51 &&
        script[2] === 0x52 &&
        script[3] === 0x01 &&
        script[4] === 0x4e &&
        script[5] === 0x01 &&
        script[6] === 0x73) ||
      // Pattern 4: Bitcoin ephemeral anchor OP_1 + push 2 bytes (4 bytes: 51024e73)
      (script.length === 4 &&
        script[0] === 0x51 &&
        script[1] === 0x02 &&
        script[2] === 0x4e &&
        script[3] === 0x73)),
  );
}

// Main function to generate unilateral exit tx chains for broadcasting CPFP transactions
export async function constructUnilateralExitFeeBumpPackages(
  nodeHexStrings: string[],
  utxos: Utxo[],
  feeRate: FeeRate,
  network: Network,
  sparkClient?: SparkServiceClient,
  logger?: Logger,
): Promise<FeeBumpTxChain[]> {
  const result: FeeBumpTxChain[] = [];

  // Sort UTXOs by value in descending order and make a copy we can modify
  const availableUtxos = [...utxos].sort((a, b) => {
    if (a.value > b.value) return -1;
    if (a.value < b.value) return 1;
    return 0;
  });

  // Convert hex strings to TreeNode objects with better error handling
  const nodes: TreeNode[] = [];
  for (let i = 0; i < nodeHexStrings.length; i++) {
    const hex = nodeHexStrings[i];
    try {
      // Validate that this looks like a proper hex string
      if (!hex || hex.length === 0) {
        throw new Error(`Node hex string at index ${i} is empty`);
      }

      // Check if this might be a raw transaction hex instead of TreeNode hex
      // Raw transaction hex typically starts with version (03000000 for version 3, required for
      // TRUC/ephemeral anchors)
      if (
        hex.startsWith("03000000") ||
        hex.startsWith("02000000") ||
        hex.startsWith("01000000")
      ) {
        throw new Error(
          `Node hex string at index ${i} appears to be a raw transaction hex, not a TreeNode protobuf. Use 'leafidtohex' command to convert node IDs to proper hex strings.`,
        );
      }

      const nodeBytes = hexToBytes(hex);
      const node = TreeNode.decode(nodeBytes);

      // Validate that the decoded node has required fields
      if (!node.id) {
        throw new Error(
          `Decoded TreeNode at index ${i} is missing required 'id' field`,
        );
      }
      if (!node.nodeTx || node.nodeTx.length === 0) {
        throw new Error(
          `Decoded TreeNode at index ${i} is missing required 'nodeTx' field`,
        );
      }

      nodes.push(node);
    } catch (decodeError) {
      throw new Error(
        `Failed to decode TreeNode hex string at index ${i}: ${formatErrorForLog(decodeError)}. Make sure you're providing TreeNode protobuf hex strings, not raw transaction hex. Use 'leafidtohex' command to get proper hex strings.`,
        { cause: decodeError },
      );
    }
  }

  // Create a map of nodes by ID for easy lookup
  const nodeMap = new Map<string, TreeNode>();
  for (const node of nodes) {
    nodeMap.set(node.id, node);
  }

  // Create a map of transactions that have already been broadcast to prevent
  // attempting to do so again.
  const broadcastTxs = new Map<string, boolean>();

  // For each provided node, build its complete chain to the root
  for (const node of nodes) {
    const txPackages: FeeBumpTxPackage[] = [];

    // Build the chain from this node to the root.
    // TODO(aakselrod): check whether
    // - two leaves are hanging off the same tree
    // - any ancestor nodes are already broadcast/confirmed
    const chain = await buildUnilateralExitChain(
      node,
      nodeMap,
      sparkClient,
      network,
    );

    // Now walk down the chain from root to leaf to build fee bump packages
    for (const chainNode of chain) {
      // Add node tx and its fee bump
      const nodeTxHex = bytesToHex(chainNode.nodeTx);

      // We skip tx's which have already been broadcasted, or we've seen in the past
      try {
        const txObj = getTxFromRawTxHex(nodeTxHex);
        const txid = getTxId(txObj);
        if (broadcastTxs.get(txid)) {
          // We already created a package for this node in another leaf.
          continue;
        }
        broadcastTxs.set(txid, true);
        const isBroadcast = await isTxBroadcast(txid, network);
        if (isBroadcast) {
          // This node has already been broadcast, so we don't need to do so.
          continue;
        }
      } catch (parseError) {
        warnUnilateralExit(
          logger,
          `unable to parse nodeTx for anchor check on node ${chainNode.id}: ${formatErrorForLog(parseError)}. This may indicate a corrupted transaction in the TreeNode; continuing with the original transaction hex, but fee bump may fail.`,
        );
      }

      const { feeBumpPsbt: nodeFeeBumpPsbt, usedUtxos } = constructFeeBumpTx(
        nodeTxHex,
        availableUtxos,
        feeRate,
        undefined,
        logger,
      );

      const feeBumpTx = btc.Transaction.fromPSBT(hexToBytes(nodeFeeBumpPsbt));

      // Get the fee bump transaction's output, if any
      const feeBumpOut: psbt.TransactionOutput | null =
        feeBumpTx.outputsLength === 1 ? feeBumpTx.getOutput(0) : null;
      let feeBumpOutPubKey: string | null = null;

      // Remove used UTXOs from the available list
      for (const usedUtxo of usedUtxos) {
        if (feeBumpOut && bytesToHex(feeBumpOut.script!) == usedUtxo.script) {
          feeBumpOutPubKey = usedUtxo.publicKey;
        }
        const index = availableUtxos.findIndex(
          (u) => u.txid === usedUtxo.txid && u.vout === usedUtxo.vout,
        );
        if (index !== -1) {
          availableUtxos.splice(index, 1);
        }
      }

      // If the fee bump TX has an output, it should have the same key as the
      // input. We can add the output to the beginning of the list.
      if (feeBumpOut)
        availableUtxos.unshift({
          txid: getTxId(feeBumpTx),
          vout: 0,
          value: feeBumpOut.amount!,
          script: bytesToHex(feeBumpOut.script!),
          publicKey: feeBumpOutPubKey!,
        });

      txPackages.push({ tx: nodeTxHex, feeBumpPsbt: nodeFeeBumpPsbt });

      // If this is the original node we started with, also add its refund tx
      if (chainNode.id === node.id) {
        const refundTxHex = bytesToHex(chainNode.refundTx);

        const refundFeeBump = constructFeeBumpTx(
          refundTxHex,
          availableUtxos,
          feeRate,
          undefined,
          logger,
        );

        const feeBumpTx = btc.Transaction.fromPSBT(
          hexToBytes(refundFeeBump.feeBumpPsbt),
        );

        const feeBumpOut: psbt.TransactionOutput | null =
          feeBumpTx.outputsLength === 1 ? feeBumpTx.getOutput(0) : null;
        let feeBumpOutPubKey: string | null = null;

        // Remove used UTXOs from the available list
        for (const usedUtxo of refundFeeBump.usedUtxos) {
          if (feeBumpOut && bytesToHex(feeBumpOut.script!) == usedUtxo.script) {
            feeBumpOutPubKey = usedUtxo.publicKey;
          }
          const index = availableUtxos.findIndex(
            (u) => u.txid === usedUtxo.txid && u.vout === usedUtxo.vout,
          );
          if (index !== -1) {
            availableUtxos.splice(index, 1);
          }
        }

        if (feeBumpOut)
          // Add to end instead of the beginning in case there are other available UTXOs we can use
          // first We don't want to wait on leaf 1 to be broadcasted to broadcast leaf 2 if we can
          // use a different available UTXO
          availableUtxos.push({
            txid: getTxId(feeBumpTx),
            vout: 0,
            value: feeBumpOut.amount!,
            script: bytesToHex(feeBumpOut.script!),
            publicKey: feeBumpOutPubKey!,
          });

        txPackages.push({
          tx: refundTxHex,
          feeBumpPsbt: refundFeeBump.feeBumpPsbt,
        });
      }
    }

    result.push({
      leafId: node.id,
      txPackages,
    });
  }

  return result;
}

// Helper function to create RIPEMD160(SHA256(data)) hash
export function hash160(data: Uint8Array): Uint8Array {
  // Proper implementation using RIPEMD160(SHA256(data))
  const sha256Hash = sha256(data);
  return ripemd160(sha256Hash);
}

// Helper function to calculate transaction vSize from hex
function calculateTransactionVSize(
  txHex: string,
  logger: Logger | undefined,
): number {
  try {
    const txBytes = hexToBytes(txHex);
    const tx = getTxFromRawTxHex(txHex);
    if (tx.vsize !== undefined) {
      return tx.vsize;
    }

    if (tx.weight !== undefined) {
      return Math.ceil(tx.weight / 4);
    }

    return txBytes.length;
  } catch (error) {
    logger?.warn(
      `calculateTransactionVSize: failed to calculate transaction vSize: ${formatErrorForLog(error)}; falling back to default estimate`,
    );
    // Fall back to default for typical transactions
    return 191;
  }
}

// Helper function to estimate CPFP fee bump transaction size
function estimateFeeBumpTxSize(
  numFundingUtxos: number,
  hasEphemeralAnchor: boolean = true,
): number {
  const TX_OVERHEAD = 10.5;

  const P2WPKH_INPUT_VBYTES = 68;

  const EPHEMERAL_ANCHOR_INPUT_VBYTES = 41;

  const P2WPKH_OUTPUT_VBYTES = 31;

  const inputVbytes =
    numFundingUtxos * P2WPKH_INPUT_VBYTES +
    (hasEphemeralAnchor ? EPHEMERAL_ANCHOR_INPUT_VBYTES : 0);

  const outputVbytes = P2WPKH_OUTPUT_VBYTES;

  return Math.ceil(TX_OVERHEAD + inputVbytes + outputVbytes);
}

// Helper function to select optimal UTXOs for fee payment
function selectUtxosForFee(
  utxos: Utxo[],
  parentTxSize: number,
  feeRate?: FeeRate,
): Utxo[] {
  if (utxos.length === 0) {
    throw new Error("No UTXOs available for selection");
  }

  // Sort UTXOs by value in descending order (largest first)
  const sortedUtxos = [...utxos].sort((a, b) => {
    if (a.value > b.value) return -1;
    if (a.value < b.value) return 1;
    return 0;
  });

  const selectedUtxos: Utxo[] = [];
  let totalValue = 0n;

  // If no fee rate provided, use all available UTXOs (fallback behavior)
  if (!feeRate?.satPerVbyte) {
    return sortedUtxos;
  }

  // Try to find the minimum number of UTXOs needed to cover the fee
  for (let i = 0; i < sortedUtxos.length; i++) {
    const utxo = sortedUtxos[i];
    if (!utxo) {
      continue;
    }
    selectedUtxos.push(utxo);
    totalValue += utxo.value;

    // Calculate child transaction size with current number of selected UTXOs
    const childTxSize = estimateFeeBumpTxSize(
      selectedUtxos.length, // number of funding UTXOs
      true, // has ephemeral anchor
    );

    // Calculate total fee needed for CPFP package
    const totalVbytes = parentTxSize + childTxSize;
    const requiredFee = BigInt(Math.ceil(totalVbytes * feeRate.satPerVbyte));

    // Minimum change amount (dust threshold)
    const minChange = 546n;

    // Check if we have enough to cover fee + minimum change
    if (totalValue >= requiredFee + minChange) {
      return selectedUtxos;
    }
  }

  return sortedUtxos;
}

// Helper function to construct a fee bump tx for a given tx using available UTXOs
export function constructFeeBumpTx(
  txHex: string,
  utxos: Utxo[],
  feeRate: FeeRate,
  previousFeeBumpTx?: string, // Optional previous fee bump tx to chain from
  logger?: Logger,
): { feeBumpPsbt: string; usedUtxos: Utxo[] } {
  // Validate inputs first
  if (!txHex || txHex.length === 0) {
    throw new Error("Transaction hex string is empty or undefined");
  }

  if (utxos.length === 0) {
    throw new Error("No UTXOs available for fee bump");
  }

  // Decode the parent tx using the utility function with error handling
  let parentTx: btc.Transaction;
  try {
    parentTx = getTxFromRawTxHex(txHex);
    if (!parentTx) {
      throw new Error("getTxFromRawTxHex returned null/undefined");
    }
  } catch (parseError) {
    throw new Error(
      `Failed to parse parent transaction hex: ${formatErrorForLog(parseError)}. Transaction hex: ${txHex}`,
      { cause: parseError },
    );
  }

  // Validate the parsed transaction has the expected structure
  try {
    const outputsLength = parentTx.outputsLength;
    const inputsLength = parentTx.inputsLength;
    if (typeof outputsLength !== "number" || outputsLength < 0) {
      throw new Error(
        "Invalid transaction: outputsLength is not a valid number",
      );
    }
    if (typeof inputsLength !== "number" || inputsLength < 0) {
      throw new Error(
        "Invalid transaction: inputsLength is not a valid number",
      );
    }
  } catch (validationError) {
    throw new Error(
      `Transaction validation failed: ${formatErrorForLog(validationError)}. This may indicate a corrupted or malformed transaction.`,
      { cause: validationError },
    );
  }

  // Find the ephemeral anchor output (should be the last output with 0 value and OP_TRUE script)
  const parentTxIdFromLib = parentTx.id; // Use the library's built-in ID property

  let ephemeralAnchorIndex = -1;

  for (let i = 0; i < parentTx.outputsLength; i++) {
    const output = parentTx.getOutput(i);

    // Check for ephemeral anchor: 0 value and OP_TRUE script patterns
    const isEphemeralAnchor = isEphemeralAnchorOutput(
      output?.script,
      output?.amount,
    );

    if (isEphemeralAnchor) {
      ephemeralAnchorIndex = i;
      break;
    }
  }
  if (ephemeralAnchorIndex === -1) {
    throw new Error(
      "No ephemeral anchor output found in parent transaction. Expected a 0-value output with OP_TRUE script (0x51), malformed OP_TRUE (0x0151), Bitcoin v29 ephemeral anchor script (015152014e0173), or Bitcoin OP_1 + push 2 bytes script (51024e73).",
    );
  }

  const ephemeralAnchorOutput = parentTx.getOutput(ephemeralAnchorIndex);
  if (!ephemeralAnchorOutput)
    throw new Error("No ephemeral anchor output found");
  if (!ephemeralAnchorOutput.script)
    throw new Error("No script found in ephemeral anchor output");
  if (utxos.length === 0) {
    throw new Error("No UTXOs available for fee bump");
  }

  // Calculate parent transaction size for CPFP fee calculation
  const parentTxSize = calculateTransactionVSize(txHex, logger);

  // Select optimal UTXOs based on fee requirements
  const selectedUtxos = selectUtxosForFee(utxos, parentTxSize, feeRate);

  // Create a new transaction using the builder pattern
  const builder = new btc.Transaction({
    version: 3,
    allowUnknown: true,
    allowLegacyWitnessUtxo: true,
  }); // ✅ set v3 in constructor

  // Track total value and process each selected funding UTXO
  let totalValue = 0n;
  const processedUtxos: Array<{
    utxo: Utxo;
    p2wpkhScript: Uint8Array;
  }> = [];

  for (let i = 0; i < selectedUtxos.length; i++) {
    const fundingUtxo = selectedUtxos[i];
    if (!fundingUtxo) {
      throw new Error(`UTXO at index ${i} is undefined`);
    }

    // Derive the public key from the private key and create the correct script
    const pubKeyHash = hash160(hexToBytes(fundingUtxo.publicKey));

    const scriptToUse = new Uint8Array([0x00, 0x14, ...pubKeyHash]); // OP_0 + 20-byte hash (P2WPKH)

    // Check if provided script is P2PKH and warn user
    const providedScript = hexToBytes(fundingUtxo.script);
    if (bytesToHex(scriptToUse) !== bytesToHex(providedScript)) {
      throw new Error(
        `❌  Derived script doesn't match provided script for UTXO ${i + 1}.`,
      );
    }

    // Add funding UTXO as segwit input using witnessUtxo (P2WPKH only)
    builder.addInput({
      txid: fundingUtxo.txid,
      index: fundingUtxo.vout,
      sequence: 0xffffffff,
      witnessUtxo: {
        script: scriptToUse, // Always P2WPKH
        amount: fundingUtxo.value,
      },
    });

    totalValue += fundingUtxo.value;
    processedUtxos.push({
      utxo: fundingUtxo,
      p2wpkhScript: scriptToUse,
    });
  }

  // Add ephemeral anchor output as the last input - use direct script spend
  builder.addInput({
    txid: parentTxIdFromLib,
    index: ephemeralAnchorIndex,
    sequence: 0xffffffff,
    witnessUtxo: {
      script: ephemeralAnchorOutput.script,
      amount: 0n,
    },
  });

  // Calculate child transaction size based on number of inputs and outputs
  // Use accurate CPFP-specific estimation instead of general estimation
  const childTxSize = estimateFeeBumpTxSize(
    selectedUtxos.length, // number of funding UTXOs
    true, // has ephemeral anchor
  );

  const totalVbytes = parentTxSize + childTxSize;

  // If no fee rate provided, fall back to fixed 1500 satoshis
  let fee: bigint;
  if (feeRate?.satPerVbyte) {
    // Calculate total fee needed for the entire package at target rate
    fee = BigInt(Math.ceil(totalVbytes * feeRate.satPerVbyte));
  } else {
    fee = 1500n;
  }

  // Minimum change amount (546 satoshis for a standard output)
  const remainingValue = totalValue - fee;

  if (remainingValue <= 0n) {
    throw new Error(
      `Insufficient funds for fee bump. Required fee: ${fee} sats, Available: ${totalValue} sats`,
    );
  }

  // Add output with remaining value using the first UTXO's P2WPKH script
  if (processedUtxos.length === 0) {
    throw new Error("No processed UTXOs available for change output");
  }

  const firstProcessedUtxo = processedUtxos[0];
  if (!firstProcessedUtxo) {
    throw new Error("First processed UTXO is undefined");
  }

  builder.addOutput({
    script: firstProcessedUtxo.p2wpkhScript,
    amount: remainingValue,
  });

  // Sign all funding UTXO inputs on the builder
  for (let i = 0; i < processedUtxos.length; i++) {
    const processed = processedUtxos[i];
    if (!processed) {
      throw new Error(`Processed UTXO at index ${i} is undefined`);
    }

    try {
      // Set proper witness script for P2WPKH input
      builder.updateInput(i, {
        witnessScript: processed.p2wpkhScript,
      });
    } catch (error) {
      throw new Error(
        `Failed to handle funding UTXO input ${i + 1}: ${formatErrorForLog(error)}`,
        { cause: error },
      );
    }
  }

  // Extract transaction bytes (funding inputs are finalized, ephemeral anchor has empty witness)
  let psbtHex: string;
  try {
    psbtHex = bytesToHex(builder.toPSBT());
  } catch (error) {
    throw new Error(
      `Failed to extract transaction: ${formatErrorForLog(error)}`,
      {
        cause: error,
      },
    );
  }

  return {
    feeBumpPsbt: psbtHex,
    usedUtxos: selectedUtxos,
  };
}