import type { Input, PrivateKey} from "libnexa-ts";
import { BufferUtils, Hash, Opcode, Output, OutputSighashType, Script, ScriptFactory, SighashType, Transaction, TxSigner } from "libnexa-ts";
import type { KeyManager } from "./key";
import type { RostrumService } from "./rostrum";
import type { AddressKey, AssetMovement, MovementDetails, TransactionDetails } from "../types";
import { tokenAmountToAssetAmount, tokenIdToHex } from "../utils";
import type { AssetService } from "./asset";

interface TransactionContext {
  address: string;
  privateKey: PrivateKey;
  txDetails: TransactionDetails;
  myInputs: Map<number, AssetMovement>;
  myOutputs: Map<number, AssetMovement>;
  involvedAssets: Set<string>;
  inputsToSign: Set<number>;
  coveredOutputs: Set<number>;
  allMyOutputsCovered: boolean;
}

export class TransactionTransformer {

  private readonly rostrumService: RostrumService;
  private readonly keyManager: KeyManager;
  private readonly assetService: AssetService;

  public constructor(rostrumService: RostrumService, keyManager: KeyManager, assetService: AssetService) {
    this.rostrumService = rostrumService;
    this.keyManager = keyManager;
    this.assetService = assetService;
  }

  public async transformRawTransaction(hex: string, addressKey: AddressKey): Promise<TransactionDetails> {
    const context = this.createContext(hex, addressKey);

    await this.processInputs(context);
    this.processOutputs(context);
    this.signInputs(context);
    await this.processAssetMovements(context);

    return context.txDetails;
  }

  private createContext(hex: string, addressKey: AddressKey): TransactionContext {
    try {
      return {
        address: addressKey.address,
        privateKey: this.keyManager.getKey(addressKey.keyPath).privateKey,
        txDetails: {
          tx: new Transaction(hex),
          send: [],
          receive: []
        },
        myInputs: new Map<number, AssetMovement>(),
        myOutputs: new Map<number, AssetMovement>(),
        involvedAssets: new Set<string>(),
        inputsToSign: new Set<number>(),
        coveredOutputs: new Set<number>(),
        allMyOutputsCovered: false,
      };
    } catch {
      throw new Error("Invalid transaction format.");
    }
  }

  private async processInputs(context: TransactionContext): Promise<void> {
    const promises = [];
    for (let i = 0; i < context.txDetails.tx.inputs.length; i++) {
      const p = this.processInput(context, i);
      promises.push(p);
    }
    await Promise.all(promises);
  }

  private async processInput(context: TransactionContext, index: number): Promise<void> {
    const input = context.txDetails.tx.inputs[index];
    const utxo = await this.rostrumService.getUtxo(BufferUtils.bufferToHex(input.outpoint));
    if (utxo.status == "spent") {
      throw new Error("Input UTXO is already spent.");
    }

    input.output = new Output(utxo.amount, utxo.scriptpubkey);
    if (input.output.address == context.address) {
      const assetId = utxo.token_id_hex || "NEXA";
      context.myInputs.set(index, {
        address: context.address,
        nexaAmount: utxo.amount.toString(),
        assetId: assetId,
        assetAmount: tokenAmountToAssetAmount(utxo.group_quantity) || utxo.amount.toString()
      });
      context.involvedAssets.add(assetId);
      
      if (input.scriptSig.isEmpty()) {
        context.inputsToSign.add(index);
      }
    }

    if (input.scriptSig.findPlaceholder() > 0) {
      context.inputsToSign.add(index);
    }
  }

  private processOutputs(context: TransactionContext): void {
    for (let i = 0; i < context.txDetails.tx.outputs.length; i++) {
      const output = context.txDetails.tx.outputs[i].toObject();
      if (output.address !== context.address) {
        continue;
      }

      const assetId = output.groupId ? tokenIdToHex(output.groupId) : "NEXA";
      const movement = {
        address: output.address,
        nexaAmount: output.value.toString(),
        assetId: assetId,
        assetAmount: tokenAmountToAssetAmount(output.groupAmount) || output.value.toString()
      };

      context.involvedAssets.add(assetId);
      context.myOutputs.set(i, movement);
    }
  }

  private signInputs(context: TransactionContext): void {
    if (context.inputsToSign.size == 0) {
      throw new Error("No inputs to sign.");
    }

    for (const index of context.inputsToSign) {
      const input = context.txDetails.tx.inputs[index];
      const subscript = this.validateAndGetSubscript(input);

      if (input.scriptSig.isEmpty()) { // only p2pkt can be empty
        context.allMyOutputsCovered = true;
        const txSig = TxSigner.sign(context.txDetails.tx, index, SighashType.ALL, subscript, context.privateKey).toTxFormat();
        const constraint = Script.empty().add(context.privateKey.publicKey.toBuffer());
        input.scriptSig = ScriptFactory.buildScriptTemplateIn(undefined, constraint, Script.empty().add(txSig))
      } else {
        const placeholderIndex = input.scriptSig.findPlaceholder();
        const placeholder = input.scriptSig.chunks[placeholderIndex];
        const sigHashBuf = placeholder.buf!.subarray(64);
        const sigHashType = SighashType.fromBuffer(sigHashBuf);
        this.processSignatureCoverage(context, sigHashType);

        const txSig = TxSigner.sign(context.txDetails.tx, index, sigHashType, subscript, context.privateKey).toTxFormat(sigHashBuf);
        input.scriptSig.replaceChunk(placeholderIndex, Script.empty().add(txSig));
      }
    }

    if (!context.allMyOutputsCovered) {
      throw new Error("Some of interested outputs are not covered by signatures.");
    }
  }

  private validateAndGetSubscript(input: Input): Script {
    if (input.output!.scriptPubKey.isPublicKeyTemplateOut()) {
      if (input.scriptSig.isEmpty() || input.scriptSig.isPublicKeyTemplateIn()) {
        return Script.empty().add(Opcode.OP_FROMALTSTACK).add(Opcode.OP_CHECKSIGVERIFY);
      }
      throw new Error("Invalid input script type.");
    }
    
    if (input.output!.scriptPubKey.isScriptTemplateOut()) {
      if (!input.scriptSig.isScriptTemplateIn()) {
        throw new Error("Unsupported input script type.");
      }

      const templateHash = input.output!.scriptPubKey.getTemplateHash() as Uint8Array;
      const templateBuf = input.scriptSig.chunks[0].buf!;
      if (BufferUtils.equals(templateHash, Hash.sha256ripemd160(templateBuf)) || BufferUtils.equals(templateHash, Hash.sha256sha256(templateBuf))) {
        return Script.fromBuffer(templateBuf);
      }

      throw new Error("Invalid input script template.");
    }
    
    throw new Error("Unsupported prevout script type.");
  }

  private processSignatureCoverage(context: TransactionContext, sigHashType: SighashType): void {
    if (context.allMyOutputsCovered) {
      return;
    }

    if (sigHashType.outType == OutputSighashType.TWO) {
      context.coveredOutputs.add(sigHashType.outData[0]);
      context.coveredOutputs.add(sigHashType.outData[1]);
      this.checkSignatureCoverage(context);
    } else if (sigHashType.outType == OutputSighashType.FIRSTN) {
      const n = sigHashType.outData[0];
      for (let i = 0; i < n; i++) {
        context.coveredOutputs.add(i);
      }
      this.checkSignatureCoverage(context);
    } else {
      context.allMyOutputsCovered = true;
    }
  }

  private checkSignatureCoverage(context: TransactionContext): void {
    if (context.allMyOutputsCovered) {
      return;
    }

    for (const i of context.myOutputs.keys()) {
      if (!context.coveredOutputs.has(i)) {
        return;
      }
    }
    context.allMyOutputsCovered = true;
  }

  private async processAssetMovements(context: TransactionContext): Promise<void> {
    const myInputs = Array.from(context.myInputs.values());
    const myOutputs = Array.from(context.myOutputs.values());

    const promises = [];
    for (const assetId of context.involvedAssets) {
      const p = this.processAssetMovement(context, assetId, myInputs, myOutputs);
      promises.push(p);
    }
    await Promise.all(promises);
  }

  private async processAssetMovement(context: TransactionContext, assetId: string, myInputs: AssetMovement[], myOutputs: AssetMovement[]): Promise<void> {
    const sentAmount = myInputs
      .reduce((sum, m) => m.assetId === assetId ? sum + BigInt(m.assetAmount) : sum, 0n);

    const receivedAmount = myOutputs
      .reduce((sum, m) => m.assetId === assetId ? sum + BigInt(m.assetAmount) : sum, 0n);

    if (sentAmount > 0n || receivedAmount > 0n) {
      const movement: MovementDetails = { amount: receivedAmount - sentAmount };
      if (assetId !== "NEXA") {
        movement.asset = await this.assetService.getAssetInfo(assetId);
      }

      if (movement.amount > 0n) {
        context.txDetails.receive.push(movement);
      } else if (movement.amount < 0n) {
        context.txDetails.send.push(movement);
      }
    }
  }
}