import { randomUUID } from "node:crypto";
import { encodeFunctionData, formatEther, getAddress, parseEther } from "viem";
import type { AgentConfig } from "./config.js";
import { assertGasReserve, createChainClients } from "./chain.js";
import {
  PONS_FACTORY,
  PONS_QUOTER_V2,
  PONS_SWAP_ROUTER,
  PONS_WETH,
  erc20Abi,
  ponsFactoryAbi,
  quoterV2Abi,
  swapRouterAbi,
} from "./contracts.js";
import { AgentState } from "./state.js";
import { decodeBuybackAcquired } from "./receipt-data.js";
import { sendTrackedTransaction } from "./broadcast.js";

export type BuybackResult = {
  mode: "simulated" | "confirmed";
  token: string;
  amountInEth: string;
  poolFee: number;
  quotedAmountOut: string;
  minimumAmountOut: string;
  priceImpactBps: number;
  approvalRequired: boolean;
  approvalTxHash?: string;
  txHash?: string;
  acquiredAmount?: string;
  executionId?: string;
};

type ChainClients = ReturnType<typeof createChainClients>;

const quote = async (
  client: ChainClients["publicClient"],
  tokenIn: `0x${string}`,
  tokenOut: `0x${string}`,
  amountIn: bigint,
  poolFee: number,
) => {
  const result = await client.simulateContract({
    address: PONS_QUOTER_V2,
    abi: quoterV2Abi,
    functionName: "quoteExactInputSingle",
    args: [{ tokenIn, tokenOut, amountIn, fee: poolFee, sqrtPriceLimitX96: 0n }],
  });
  return result.result[0];
};

export const executeBuyback = async (config: AgentConfig, tokenInput: string, amountInEth: string): Promise<BuybackResult> => {
  if (!config.account) throw new Error("Execution wallet is required for buybacks");
  const token = getAddress(tokenInput);
  const amountIn = parseEther(amountInEth);
  if (amountIn < config.minBuybackWei) {
    throw new Error(`Buyback is below the minimum ${formatEther(config.minBuybackWei)} WETH`);
  }

  const balanceState = new AgentState(config.stateDbPath);
  const available = balanceState.getBucketBalance(config.account.address, "buyback", PONS_WETH);
  balanceState.close();
  if (amountIn > available) {
    throw new Error(`Buyback exceeds allocated WETH. Requested ${formatEther(amountIn)}, available ${formatEther(available)}`);
  }

  const clients = createChainClients(config);
  const { publicClient, walletClient } = clients;
  const launched = await publicClient.readContract({
    address: PONS_FACTORY,
    abi: ponsFactoryAbi,
    functionName: "getLaunchedToken",
    args: [token],
  });
  if (!launched.exists) throw new Error("Buyback target is not a token launched by the active Pons factory");
  if (launched.pairedToken.toLowerCase() !== PONS_WETH.toLowerCase()) {
    throw new Error("Buyback target is not paired with Pons WETH");
  }

  const walletWeth = await publicClient.readContract({
    address: PONS_WETH,
    abi: erc20Abi,
    functionName: "balanceOf",
    args: [config.account.address],
  });
  if (walletWeth < amountIn) {
    throw new Error(`Wallet has ${formatEther(walletWeth)} WETH but buyback requires ${formatEther(amountIn)} WETH`);
  }

  const poolFee = Number(launched.poolFee);
  const quotedAmountOut = await quote(publicClient, PONS_WETH, token, amountIn, poolFee);
  const referenceAmount = amountIn / 20n > 0n ? amountIn / 20n : 1n;
  const referenceOut = await quote(publicClient, PONS_WETH, token, referenceAmount, poolFee);
  const linearOutput = referenceOut * amountIn / referenceAmount;
  const priceImpactBps = linearOutput > quotedAmountOut
    ? Number((linearOutput - quotedAmountOut) * 10_000n / linearOutput)
    : 0;
  if (priceImpactBps > config.maxPriceImpactBps) {
    throw new Error(`Estimated price impact ${priceImpactBps} bps exceeds ${config.maxPriceImpactBps} bps`);
  }

  const minimumAmountOut = quotedAmountOut * BigInt(10_000 - config.maxSlippageBps) / 10_000n;
  const swapParams = {
    tokenIn: PONS_WETH,
    tokenOut: token,
    fee: poolFee,
    recipient: config.account.address,
    amountIn,
    amountOutMinimum: minimumAmountOut,
    sqrtPriceLimitX96: 0n,
  } as const;
  const allowance = await publicClient.readContract({
    address: PONS_WETH,
    abi: erc20Abi,
    functionName: "allowance",
    args: [config.account.address, PONS_SWAP_ROUTER],
  });

  if (config.mode !== "live") {
    if (allowance >= amountIn) {
      await publicClient.simulateContract({
        account: config.account,
        address: PONS_SWAP_ROUTER,
        abi: swapRouterAbi,
        functionName: "exactInputSingle",
        args: [swapParams],
      });
    }
    return {
      mode: "simulated",
      token,
      amountInEth: formatEther(amountIn),
      poolFee,
      quotedAmountOut: quotedAmountOut.toString(),
      minimumAmountOut: minimumAmountOut.toString(),
      priceImpactBps,
      approvalRequired: allowance < amountIn,
    };
  }
  if (!walletClient) throw new Error("Wallet client is unavailable in live mode");

  const executionId = randomUUID();
  const reservationId = `${executionId}:buyback-reservation`;
  const plannedResult = {
    token,
    amountInEth: formatEther(amountIn),
    poolFee,
    quotedAmountOut: quotedAmountOut.toString(),
    minimumAmountOut: minimumAmountOut.toString(),
    priceImpactBps,
    approvalRequired: allowance < amountIn,
  };
  const reservationState = new AgentState(config.stateDbPath);
  const reserved = reservationState.startBuybackExecutionAndReserve(config.account.address, {
    id: executionId,
    target: token,
    amount: amountIn.toString(),
    result: plannedResult,
  }, { id: reservationId, asset: PONS_WETH });
  reservationState.close();
  if (!reserved) throw new Error("Buyback allocation was consumed by another execution");

  let approvalTxHash: `0x${string}` | undefined;
  let txHash: `0x${string}` | undefined;
  let preparedTxHash: `0x${string}` | undefined;
  let receiptSucceeded = false;
  let swapReverted = false;
  try {
    await assertGasReserve(publicClient, config.account.address, config.minGasReserveWei);
    if (allowance < amountIn) {
      const approval = await publicClient.simulateContract({
        account: config.account,
        address: PONS_WETH,
        abi: erc20Abi,
        functionName: "approve",
        args: [PONS_SWAP_ROUTER, amountIn],
      });
      approvalTxHash = await walletClient.writeContract(approval.request);
      const approvalReceipt = await publicClient.waitForTransactionReceipt({ hash: approvalTxHash, confirmations: 2, timeout: 120_000 });
      if (approvalReceipt.status !== "success") throw new Error(`WETH approval reverted: ${approvalTxHash}`);
    }

    const balanceBefore = await publicClient.readContract({ address: token, abi: erc20Abi, functionName: "balanceOf", args: [config.account.address] });
    await publicClient.simulateContract({
      account: config.account,
      address: PONS_SWAP_ROUTER,
      abi: swapRouterAbi,
      functionName: "exactInputSingle",
      args: [swapParams],
    });
    const data = encodeFunctionData({ abi: swapRouterAbi, functionName: "exactInputSingle", args: [swapParams] });
    txHash = await sendTrackedTransaction(config, clients, { to: PONS_SWAP_ROUTER, data }, (preparedHash) => {
      preparedTxHash = preparedHash;
      const submittedState = new AgentState(config.stateDbPath);
      submittedState.setLedgerEntryTxHash(config.account!.address, reservationId, preparedHash);
      submittedState.updateExecution(config.account!.address, executionId, "submitted", { txHash: preparedHash, result: { ...plannedResult, approvalTxHash } });
      submittedState.close();
    });

    const receipt = await publicClient.waitForTransactionReceipt({ hash: txHash, confirmations: 2, timeout: 120_000 });
    if (receipt.status !== "success") {
      swapReverted = true;
      throw new Error(`Buyback reverted: ${txHash}`);
    }
    receiptSucceeded = true;

    const balanceAfter = await publicClient.readContract({ address: token, abi: erc20Abi, functionName: "balanceOf", args: [config.account.address] });
    const transferredAmount = decodeBuybackAcquired(receipt, token, config.account.address);
    const acquiredAmount = transferredAmount > 0n ? transferredAmount : balanceAfter - balanceBefore;
    if (acquiredAmount < minimumAmountOut) {
      throw new Error(`Buyback output ${acquiredAmount} is below the recorded minimum ${minimumAmountOut}`);
    }

    const result = {
      mode: "confirmed" as const,
      token,
      amountInEth: formatEther(amountIn),
      poolFee,
      quotedAmountOut: quotedAmountOut.toString(),
      minimumAmountOut: minimumAmountOut.toString(),
      priceImpactBps,
      approvalRequired: allowance < amountIn,
      approvalTxHash,
      txHash,
      acquiredAmount: acquiredAmount.toString(),
      executionId,
    };
    const resultState = new AgentState(config.stateDbPath);
    resultState.updateExecution(config.account.address, executionId, "confirmed", { txHash, result });
    resultState.close();
    return result;
  } catch (error) {
    const failureState = new AgentState(config.stateDbPath);
    const recoverableHash = txHash ?? preparedTxHash;
    if (receiptSucceeded || (recoverableHash && !swapReverted)) {
      // A timeout after broadcast is deliberately frozen; the tx may still succeed.
      failureState.updateExecution(config.account.address, executionId, "uncertain", { txHash: recoverableHash, error: error instanceof Error ? error.message : String(error) });
    } else {
      failureState.releaseBucketFunds(config.account.address, reservationId, "buyback", PONS_WETH, amountIn.toString());
      failureState.updateExecution(config.account.address, executionId, "failed", { error: error instanceof Error ? error.message : String(error) });
    }
    failureState.close();
    throw error;
  }
};

export type ConversionResult = {
  mode: "simulated" | "confirmed";
  inputAsset: string;
  outputAsset: string;
  amountIn: string;
  amountOut: string;
  minimumAmountOut: string;
  priceImpactBps: number;
  txHash?: string;
  executionId?: string;
};

/** Convert the non-WETH half of a claimed fee allocation into WETH first. */
export const convertBuybackAssetToWeth = async (config: AgentConfig, assetInput: string, amountInput: bigint): Promise<ConversionResult> => {
  if (!config.account) throw new Error("Execution wallet is required for fee conversion");
  const inputAsset = getAddress(assetInput);
  if (inputAsset.toLowerCase() === PONS_WETH.toLowerCase()) throw new Error("WETH does not require conversion");
  if (amountInput <= 0n) throw new Error("Conversion amount must be positive");

  const clients = createChainClients(config);
  const { publicClient, walletClient } = clients;
  if (!walletClient) throw new Error("Wallet client is unavailable in live mode");
  const launched = await publicClient.readContract({
    address: PONS_FACTORY,
    abi: ponsFactoryAbi,
    functionName: "getLaunchedToken",
    args: [inputAsset],
  });
  if (!launched.exists || launched.pairedToken.toLowerCase() !== PONS_WETH.toLowerCase()) {
    throw new Error("Non-WETH fee allocation is not a verified WETH-paired Pons token");
  }

  const walletInput = await publicClient.readContract({ address: inputAsset, abi: erc20Abi, functionName: "balanceOf", args: [config.account.address] });
  if (walletInput < amountInput) throw new Error(`Wallet has insufficient ${inputAsset} balance for fee conversion`);
  const poolFee = Number(launched.poolFee);
  const quotedAmountOut = await quote(publicClient, inputAsset, PONS_WETH, amountInput, poolFee);
  const referenceAmount = amountInput / 20n > 0n ? amountInput / 20n : 1n;
  const referenceOut = await quote(publicClient, inputAsset, PONS_WETH, referenceAmount, poolFee);
  const linearOutput = referenceOut * amountInput / referenceAmount;
  const priceImpactBps = linearOutput > quotedAmountOut ? Number((linearOutput - quotedAmountOut) * 10_000n / linearOutput) : 0;
  if (priceImpactBps > config.maxPriceImpactBps) throw new Error(`Fee conversion price impact ${priceImpactBps} bps exceeds ${config.maxPriceImpactBps} bps`);
  const minimumAmountOut = quotedAmountOut * BigInt(10_000 - config.maxSlippageBps) / 10_000n;
  const swapParams = {
    tokenIn: inputAsset,
    tokenOut: PONS_WETH,
    fee: poolFee,
    recipient: config.account.address,
    amountIn: amountInput,
    amountOutMinimum: minimumAmountOut,
    sqrtPriceLimitX96: 0n,
  } as const;
  const allowance = await publicClient.readContract({ address: inputAsset, abi: erc20Abi, functionName: "allowance", args: [config.account.address, PONS_SWAP_ROUTER] });
  if (config.mode !== "live") {
    if (allowance >= amountInput) await publicClient.simulateContract({ account: config.account, address: PONS_SWAP_ROUTER, abi: swapRouterAbi, functionName: "exactInputSingle", args: [swapParams] });
    return { mode: "simulated", inputAsset, outputAsset: PONS_WETH, amountIn: amountInput.toString(), amountOut: quotedAmountOut.toString(), minimumAmountOut: minimumAmountOut.toString(), priceImpactBps };
  }

  const executionId = randomUUID();
  const reservationId = `${executionId}:buyback-reservation`;
  const plannedResult = { action: "convert", inputAsset, outputAsset: PONS_WETH, amountIn: amountInput.toString(), amountOut: quotedAmountOut.toString(), minimumAmountOut: minimumAmountOut.toString(), priceImpactBps };
  const reservationState = new AgentState(config.stateDbPath);
  const reserved = reservationState.startBuybackExecutionAndReserve(config.account.address, { id: executionId, target: inputAsset, amount: amountInput.toString(), result: plannedResult }, { id: reservationId, asset: inputAsset });
  reservationState.close();
  if (!reserved) throw new Error("Fee conversion allocation was consumed by another execution");

  let preparedTxHash: `0x${string}` | undefined;
  let txHash: `0x${string}` | undefined;
  let receiptSucceeded = false;
  let swapReverted = false;
  try {
    await assertGasReserve(publicClient, config.account.address, config.minGasReserveWei);
    if (allowance < amountInput) {
      const approval = await publicClient.simulateContract({ account: config.account, address: inputAsset, abi: erc20Abi, functionName: "approve", args: [PONS_SWAP_ROUTER, amountInput] });
      const approvalHash = await walletClient.writeContract(approval.request);
      const approvalReceipt = await publicClient.waitForTransactionReceipt({ hash: approvalHash, confirmations: 2, timeout: 120_000 });
      if (approvalReceipt.status !== "success") throw new Error(`Fee conversion approval reverted: ${approvalHash}`);
    }
    const balanceBefore = await publicClient.readContract({ address: PONS_WETH, abi: erc20Abi, functionName: "balanceOf", args: [config.account.address] });
    await publicClient.simulateContract({ account: config.account, address: PONS_SWAP_ROUTER, abi: swapRouterAbi, functionName: "exactInputSingle", args: [swapParams] });
    const data = encodeFunctionData({ abi: swapRouterAbi, functionName: "exactInputSingle", args: [swapParams] });
    txHash = await sendTrackedTransaction(config, clients, { to: PONS_SWAP_ROUTER, data }, (preparedHash) => {
      preparedTxHash = preparedHash;
      const submittedState = new AgentState(config.stateDbPath);
      submittedState.setLedgerEntryTxHash(config.account!.address, reservationId, preparedHash);
      submittedState.updateExecution(config.account!.address, executionId, "submitted", { txHash: preparedHash, result: plannedResult });
      submittedState.close();
    });
    const receipt = await publicClient.waitForTransactionReceipt({ hash: txHash, confirmations: 2, timeout: 120_000 });
    if (receipt.status !== "success") {
      swapReverted = true;
      throw new Error(`Fee conversion reverted: ${txHash}`);
    }
    receiptSucceeded = true;
    const balanceAfter = await publicClient.readContract({ address: PONS_WETH, abi: erc20Abi, functionName: "balanceOf", args: [config.account.address] });
    const transferredAmount = decodeBuybackAcquired(receipt, PONS_WETH, config.account.address);
    const amountOut = transferredAmount > 0n ? transferredAmount : balanceAfter - balanceBefore;
    if (amountOut < minimumAmountOut) throw new Error(`Fee conversion output ${amountOut} is below the recorded minimum ${minimumAmountOut}`);
    const result = { mode: "confirmed" as const, inputAsset, outputAsset: PONS_WETH, amountIn: amountInput.toString(), amountOut: amountOut.toString(), minimumAmountOut: minimumAmountOut.toString(), priceImpactBps, txHash, executionId };
    const confirmedState = new AgentState(config.stateDbPath);
    confirmedState.recordConversionExecution(config.account.address, executionId, { txHash, inputAsset, inputAmount: amountInput.toString(), outputAsset: PONS_WETH, outputAmount: amountOut.toString() }, result);
    confirmedState.close();
    return result;
  } catch (error) {
    const failureState = new AgentState(config.stateDbPath);
    const recoverableHash = txHash ?? preparedTxHash;
    if (receiptSucceeded || (recoverableHash && !swapReverted)) {
      failureState.updateExecution(config.account.address, executionId, "uncertain", { txHash: recoverableHash, error: error instanceof Error ? error.message : String(error) });
    } else {
      failureState.releaseBucketFunds(config.account.address, reservationId, "buyback", inputAsset, amountInput.toString());
      failureState.updateExecution(config.account.address, executionId, "failed", { txHash: recoverableHash, error: error instanceof Error ? error.message : String(error) });
    }
    failureState.close();
    throw error;
  }
};
