import { ScriptType } from "../types";
import { encodingLength } from "varuint-bitcoin";
import { estimateInputSize, getOutputSize } from "../common";

export function getTxOverhead(
  inputScripts: ScriptType[],
  outputCount: number,
): number {
  const inputCount = inputScripts.length;

  const hasWitnessInputs = inputScripts.some((type) =>
    [
      ScriptType.P2WPKH,
      ScriptType.P2WSH,
      ScriptType.P2TR,
      ScriptType.P2SH_P2WPKH,
      ScriptType.P2SH_P2WSH,
    ].includes(type),
  );

  const baseSize =
    4 + // nVersion
    encodingLength(inputCount) + // input count (varint)
    encodingLength(outputCount) + // output count (varint)
    4; // nLockTime

  const segwitMarkerAndFlag = hasWitnessInputs ? 0.5 : 0;

  return baseSize + segwitMarkerAndFlag;
}

export function estimateTxVBytesSimple(
  inputScript: ScriptType[] = [ScriptType.P2TR],
  outputScript: ScriptType[] = [ScriptType.P2TR],
  opReturnData: Buffer | string | null = null,
): number {
  const overhead = getTxOverhead(inputScript, outputScript.length);
  let totalInput = 0;
  for (const script of inputScript) {
    totalInput += estimateInputSize(script);
  }

  const outputSize = outputScript.reduce(
    (total, scriptType) => total + getOutputSize(scriptType),
    0,
  );

  let opReturnSize = 0;
  if (opReturnData !== null) {
    const dataBuffer =
      typeof opReturnData === "string"
        ? Buffer.from(opReturnData, "utf8")
        : opReturnData;
    opReturnSize =
      8 + 1 + encodingLength(dataBuffer.length) + dataBuffer.length; // (8)value + 1(OP_RETURN) + (N)length Bytes + (M)Data
  }

  const vBytes = Math.ceil(overhead + totalInput + outputSize + opReturnSize);
  return vBytes;
}

export function estimateTxFee(
  inputScript: ScriptType[] = [ScriptType.P2TR],
  outputScript: ScriptType[] = [ScriptType.P2TR],
  feeRate: number,
  opReturnData: Buffer | string | null = null,
): number {
  const vBytes = estimateTxVBytesSimple(
    inputScript,
    outputScript,
    opReturnData,
  );

  return Math.ceil(vBytes * feeRate);
}