import { ChainArtifacts, ChainDefinition, findContract, getArtifacts, renderTrace, TraceEntry } from '@usecannon/builder';
import { bold, gray, green, red, yellow } from 'chalk';
import Debug from 'debug';
import _ from 'lodash';
import * as viem from 'viem';
import { readDeployRecursive } from '../package';
import { getProvider, runRpc } from '../rpc';
import { CliSettings } from '../settings';
import { logSpinner } from '../util/console';
import { ProviderAction, resolveProvider } from '../util/provider';
import { ANVIL_FIRST_ADDRESS } from '../constants';

const debug = Debug('cannon:cli:trace');

export async function trace({
  packageRef,
  data,
  chainId,
  cliSettings,
  from,
  to,
  value,
  block,
  json = false,
}: {
  packageRef: string;
  data: viem.Hex;
  chainId: number;
  cliSettings: CliSettings;
  from?: viem.Address;
  to?: viem.Address;
  block?: string;
  value?: bigint;
  json: boolean;
}) {
  // data can be:
  // 1. on-chain transaction hash
  // 2. calldata (will automatically detect contract to execute on)
  //   * in which case, will figure out what to execute
  //   * additional option to override the contract
  //
  // in any case, cannon will run the transaction in anvil. afterwards, it
  // will call `trace_transaction`, and decode as much data from the trace
  // as possible, the same way that an error occurs

  // if chain id is not specified, get it from the provider
  const { provider } = await resolveProvider({ action: ProviderAction.ReadProvider, cliSettings, chainId });

  const deployInfos = await readDeployRecursive(packageRef, chainId);

  const artifacts: ChainArtifacts = {};

  for (const di of deployInfos) {
    _.merge(artifacts, getArtifacts(new ChainDefinition(di.def), di.state));
  }

  if (viem.isHash(data)) {
    const txHash = data as viem.Hash;

    try {
      const txData = await provider.getTransaction({ hash: txHash });
      const txReceipt = await provider.getTransactionReceipt({ hash: txHash });

      // this is a transaction hash
      logSpinner(gray('Detected transaction hash'));

      data = txData.input;
      value = value || txData.value;
      block = block || txReceipt.blockNumber.toString();
      from = from || txData.from;
      if (!to && txData.to) to = txData.to;
    } catch (err) {
      throw new Error('could not get transaction information. The transaction may not exist?');
    }
  } else if (!to) {
    const r = findContract(artifacts, ({ abi }) => {
      try {
        viem.decodeFunctionData({ abi, data });
        return true;
      } catch (_) {
        // intentionally empty
      }

      return false;
    });
    if (r !== null) {
      to = r.contract.address;
      logSpinner(gray(`Inferred contract for call: ${r.name}`));
    } else {
      logSpinner(
        yellow(
          'Could not find a contract for this call. Are you sure the call can be traced on a contract on this cannon package? Pass `--to` to set manually if necessary'
        )
      );
    }
  }

  // create an anvil server
  let rpc;
  if (block) {
    // subtract one second because 1 second is added when the block is mined
    const blockInfo = await provider.getBlock(
      (block || 'latest').match(/^[0-9]*$/) ? { blockNumber: BigInt(block) } : { blockTag: block as viem.BlockTag }
    );
    const timestamp = blockInfo.timestamp - BigInt(1);
    rpc = await runRpc(
      {
        port: 0,
        forkBlockNumber: !block || block === 'latest' ? undefined : (blockInfo.number! - BigInt(1)).toString(),
        timestamp,
        chainId,
      },
      { forkProvider: provider as any }
    );
  } else {
    rpc = await runRpc({ port: 0, chainId }, { forkProvider: provider as any });
  }
  const simulateProvider = getProvider(rpc)!;

  const fullTxn = {
    from: from || viem.zeroAddress,
    to: to || viem.zeroAddress,
    data,
    value,
    // set the gas limit very high to make sure the txn does not try to estimate
    gasLimit: 10000000,
  };

  debug('full txn to execute', fullTxn);

  // now we should be able to run the transaction. force it through
  let txnHash: viem.Hash;
  try {
    const signer = (fullTxn.from || ANVIL_FIRST_ADDRESS) as viem.Address;
    if (fullTxn.from) {
      await simulateProvider.impersonateAccount({ address: fullTxn.from });
      await simulateProvider.setBalance({ address: fullTxn.from, value: viem.parseEther('10000') });
    }

    logSpinner(gray('Simulating transaction (be patient! this could take a while...)'));
    const pushedTxn = await simulateProvider.sendTransaction({ account: signer, chain: simulateProvider.chain, ...fullTxn });

    try {
      await simulateProvider.waitForTransactionReceipt({ hash: pushedTxn });
    } catch {
      // intentionally empty
    }
    txnHash = pushedTxn;
  } catch (err: any) {
    throw new Error(`failed to simulate txn: ${err.toString()}`);
  }

  // once we have forced through the transaction, call `trace_transaction`
  const traces: TraceEntry[] = await simulateProvider.request({ method: 'trace_transaction' as any, params: [txnHash] });

  if (!json) {
    const traceText = renderTrace(artifacts, traces);

    logSpinner(traceText);

    const receipt = await simulateProvider.getTransactionReceipt({ hash: txnHash });
    const totalGasUsed = computeGasUsed(traces, fullTxn).toLocaleString();
    logSpinner();
    if (receipt.status == 'success') {
      logSpinner(
        green(
          bold(
            `Transaction completes successfully with return value: ${
              traces[0].result?.output ?? 'unknown'
            } (${totalGasUsed} gas)`
          )
        )
      );
    } else {
      logSpinner(
        red(bold(`Transaction completes with error: ${traces[0].result?.output ?? 'unknown'} (${totalGasUsed} gas)`))
      );
    }
  } else {
    logSpinner(JSON.stringify(traces, null, 2));
  }
}

function computeGasUsed(traces: TraceEntry[], txn: viem.TransactionRequest): number {
  // total gas required for the transaction is whatever was used by the actual txn
  // + 21000 (base transaction cost)
  // + cost of the txn calldata
  const txnData = viem.hexToBytes(txn.data || '0x');
  const zeroDataCount = txnData.filter((d) => d === 0).length;
  const nonZeroDataCount = txnData.length - zeroDataCount;
  return parseInt(traces[0].result?.gasUsed ?? '0') + 21000 + 4 * zeroDataCount + 16 * nonZeroDataCount;
}