import { isAddress } from "ethers";
import { ethers } from "hardhat";
import path from "path";
import {
  assertSelectedSimulationSignerAuthorized,
  decideControlledMulticallAuthority,
  decideAuthority,
} from "../lib/authority";
import {
  actualContractName,
  contractInfoKey,
  deploymentInstanceId,
  loadContractInfo,
  resolveDeployment,
} from "../lib/deployments";
import type { DeploymentChange } from "../lib/deploymentRecords";
import { ExternalExecutionRequired } from "../lib/external";
import { runGateToolValidate } from "../lib/gateTool";
import type { GateToolUpgradeOutput } from "../lib/gateTool";
import { recordParameterLockImplementationSalts } from "../lib/lock";
import { resolveValue, WorkflowStep } from "../lib/parameters";
import {
  prepareControlledMulticallTransaction,
  sendPreparedControlledMulticallTransaction,
} from "../lib/create2";
import {
  prepareExternalProxyUpgrade,
  ProxyOptions,
  validateAndDeployUpgradeImplementation,
} from "../lib/proxy";
import { writeJsonFile } from "../lib/config";
import type { WorkflowContext } from "../run";
import { impersonateWorkflowSimulationSigner } from "../lib/signer";

const resolveInContext = (
  context: WorkflowContext,
  value: unknown,
): unknown => {
  return resolveValue(value, {
    targetId: context.targetId,
    target: context.target,
    params: context.targetParams,
    state: context.state,
  });
};

const resolveObject = (
  context: WorkflowContext,
  value: unknown,
): Record<string, unknown> => {
  const resolved = resolveInContext(context, value || {});
  if (!resolved || typeof resolved !== "object" || Array.isArray(resolved)) {
    return {};
  }
  return resolved as Record<string, unknown>;
};

const resolveArray = (
  context: WorkflowContext,
  value: unknown,
  label: string,
): unknown[] => {
  const resolved = resolveInContext(context, value || []);
  if (!Array.isArray(resolved)) {
    throw new Error(`${label} must resolve to an array`);
  }
  return resolved;
};

const stringField = (
  source: Record<string, unknown>,
  key: string,
): string | undefined => {
  return typeof source[key] === "string" && source[key] ? String(source[key]) : undefined;
};

const proxyReference = (
  context: WorkflowContext,
  step: WorkflowStep,
): {
  proxyAddress: string;
  saveAs: string;
  contractInfoKey: string;
  proxy: ProxyOptions;
  proxyAdmin?: string;
} => {
  const proxyObject = resolveObject(context, step.proxy);
  const rawProxy =
    stringField(proxyObject, "address") ||
    stringField(proxyObject, "name") ||
    stringField(proxyObject, "contract") ||
    (typeof step.proxy === "string" ? String(resolveInContext(context, step.proxy)) : undefined) ||
    (typeof step.address === "string" ? String(resolveInContext(context, step.address)) : undefined);

  if (!rawProxy) {
    throw new Error(`Step ${step.id} requires proxy address or deployment name`);
  }

  const registryDeployment = isAddress(rawProxy)
    ? undefined
    : resolveDeployment(
        loadContractInfo(context.root, context.config),
        context.targetId,
        rawProxy,
      );
  if (!isAddress(rawProxy) && !registryDeployment) {
    throw new Error(`Missing contractInfo address: ${context.targetId}.${rawProxy}`);
  }

  const proxyAddress = isAddress(rawProxy) ? rawProxy : registryDeployment!.address;
  const defaultSaveAs = isAddress(rawProxy)
    ? actualContractName(String(step.contract))
    : registryDeployment!.key.split(":", 1)[0];
  const saveAs = deploymentInstanceId(String(step.saveAs || defaultSaveAs));
  const registryKey = registryDeployment?.key || contractInfoKey(saveAs, String(step.contract));
  const initializer = proxyObject.initializer === false
    ? false
    : stringField(proxyObject, "initializer");

  return {
    proxyAddress,
    saveAs,
    contractInfoKey: registryKey,
    proxy: {
      kind: stringField(proxyObject, "kind") || (typeof step.proxyKind === "string" ? step.proxyKind : "uups"),
      initializer,
      unsafeAllow: Array.isArray(proxyObject.unsafeAllow)
        ? proxyObject.unsafeAllow.map((item) => String(item))
        : undefined,
      constructorArgs: proxyObject.constructorArgs === undefined
        ? undefined
        : resolveArray(context, proxyObject.constructorArgs, `Step ${step.id} proxy.constructorArgs`),
    },
    proxyAdmin: stringField(proxyObject, "proxyAdmin") || (typeof step.proxyAdmin === "string" ? step.proxyAdmin : undefined),
  };
};

const upgradeCall = (
  context: WorkflowContext,
  step: WorkflowStep,
): { fn: string; args: unknown[] } | undefined => {
  const call = resolveObject(context, step.call);
  const fn = stringField(call, "fn") || stringField(call, "method");
  if (!fn) return undefined;

  const args = resolveArray(context, call.args || [], `Step ${step.id} call.args`);
  return {
    fn,
    args,
  };
};

export const runUpgradeStep = async (
  context: WorkflowContext,
  step: WorkflowStep,
): Promise<void> => {
  const contractName = String(step.contract);
  const proxy = proxyReference(context, step);
  const call = upgradeCall(context, step);
  const actualContract = actualContractName(contractName);
  const preparedChange: DeploymentChange = {
    stepId: step.id,
    kind: "upgrade",
    instanceId: proxy.saveAs,
    actualContract,
    contractInfoKey: proxy.contractInfoKey,
    strategy: "gate-tool",
    proxyKind: proxy.proxy.kind || "uups",
    proxyAddress: proxy.proxyAddress,
  };
  if (context.executionPlanner) {
    if (!context.signer || !context.signerAddress) {
      throw new Error(`Upgrade step ${step.id} requires a planned signer context`);
    }
    const recorder = context.mode === "execute" ? context.deploymentRecorder : undefined;
    if (context.mode === "execute" && !recorder) {
      throw new Error(`Upgrade step ${step.id} requires an execution deployment recorder`);
    }
    const targetCreate2 = context.target.create2;
    const targetMulticall = context.target.multicall;
    if (!targetCreate2?.factory || !targetMulticall?.controlledMulticall) {
      throw new Error(
        `Upgrade step ${step.id} requires target create2 factory and controlledMulticall`,
      );
    }
    const constructorArgs = resolveArray(
      context,
      step.constructorArgs || proxy.proxy.constructorArgs || [],
      `Step ${step.id} constructorArgs`,
    );
    const prepared = await prepareExternalProxyUpgrade({
      executionId: context.executionId,
      stepId: step.id,
      proxyAddress: proxy.proxyAddress,
      contractName,
      constructorArgs,
      unsafeAllow: proxy.proxy.unsafeAllow,
      create2: {
        factory: targetCreate2.factory,
        salt: "0x" + "00".repeat(32),
      },
      call,
    });
    const plannedChange: DeploymentChange = {
      ...preparedChange,
      implementationAddress: prepared.implementationAddress,
      implementationSalt: prepared.implementationSalt,
    };
    if (context.parameterLockExecutionId) {
      recordParameterLockImplementationSalts(
        context.taskDir,
        context.parameterLockExecutionId,
        { [step.id + ".implementation"]: prepared.implementationSalt },
      );
    }
    recorder?.prepare(plannedChange);

    const implementationCall = prepared.calls[0];
    const multicallTransaction = prepareControlledMulticallTransaction(
      targetMulticall.controlledMulticall,
      [{
        target: implementationCall.to,
        value: implementationCall.value,
        data: implementationCall.data,
      }],
    );
    const multicall = await ethers.getContractAt(
      "Multicall3",
      targetMulticall.controlledMulticall,
      context.signer,
    );
    const deploymentAuthority = await decideControlledMulticallAuthority(
      multicall as unknown as {
        callers: (caller: string) => Promise<boolean>;
        factoryCallers: (factoryAddress: string, caller: string) => Promise<boolean>;
        owner?: () => Promise<unknown>;
      },
      targetCreate2.factory,
      context.target,
      context.signerAddress,
    );
    const deploymentExecutor = deploymentAuthority.mode === "execute"
      ? context.signerAddress
      : deploymentAuthority.executor;
    if (!deploymentExecutor) {
      throw new Error(`Upgrade step ${step.id} cannot resolve implementation executor`);
    }
    const implementationTransaction = context.executionPlanner.add({
      stepId: step.id + "-implementation",
      executor: deploymentExecutor,
      to: multicallTransaction.target,
      value: multicallTransaction.value,
      data: multicallTransaction.data,
      operation: 0,
    });
    if (implementationTransaction.executionAction === "replay") {
      // The immutable transaction was already submitted or finalized in an earlier segment.
    } else if (implementationTransaction.route === "external" && context.mode !== "simulate") {
      context.externalCalls.push({
        transactionId: implementationTransaction.transactionId,
        stepId: implementationTransaction.stepId,
        to: implementationTransaction.to,
        value: implementationTransaction.value,
        data: implementationTransaction.data,
        operation: implementationTransaction.operation,
        method: multicallTransaction.method,
        args: multicallTransaction.calls,
        executor: implementationTransaction.executor,
        reason: deploymentAuthority.reason,
      });
    } else {
      const selectedSigner = implementationTransaction.route === "external"
        ? await impersonateWorkflowSimulationSigner(implementationTransaction.executor)
        : context.signer;
      const transaction = await sendPreparedControlledMulticallTransaction(
        multicallTransaction,
        selectedSigner,
      );
      if (context.mode === "execute") {
        context.operatorReceipts.push({
          transactionId: implementationTransaction.transactionId,
          transactionHash: transaction.hash,
        });
      }
    }

    const upgradeCall = prepared.calls[1];
    const proxyContract = await ethers.getContractAt(
      String(step.abi || contractName),
      proxy.proxyAddress,
      context.signer,
    );
    const upgradeAuthority = await decideAuthority(
      proxyContract,
      context.target,
      context.signerAddress,
    );
    const upgradeExecutor = upgradeAuthority.mode === "execute"
      ? context.signerAddress
      : upgradeAuthority.executor;
    if (!upgradeExecutor) {
      throw new Error(`Upgrade step ${step.id} cannot resolve proxy executor`);
    }
    const upgradeTransaction = context.executionPlanner.add({
      stepId: step.id,
      executor: upgradeExecutor,
      to: upgradeCall.to,
      value: upgradeCall.value,
      data: upgradeCall.data,
      operation: 0,
    });
    let upgradeTxHash: string | undefined;
    if (upgradeTransaction.executionAction === "replay") {
      upgradeTxHash = upgradeTransaction.transactionHash;
    } else if (upgradeTransaction.route === "external" && context.mode !== "simulate") {
      context.externalCalls.push({
        ...upgradeCall,
        transactionId: upgradeTransaction.transactionId,
        stepId: upgradeTransaction.stepId,
        to: upgradeTransaction.to,
        value: upgradeTransaction.value,
        data: upgradeTransaction.data,
        operation: upgradeTransaction.operation,
        executor: upgradeTransaction.executor,
        reason: upgradeAuthority.reason,
      });
    } else {
      const selectedSigner = upgradeTransaction.route === "external"
        ? await impersonateWorkflowSimulationSigner(upgradeTransaction.executor)
        : context.signer;
      const transaction = await selectedSigner.sendTransaction({
        to: upgradeTransaction.to,
        value: upgradeTransaction.value,
        data: upgradeTransaction.data,
      });
      await transaction.wait();
      upgradeTxHash = transaction.hash;
      if (context.mode === "execute") {
        context.operatorReceipts.push({
          transactionId: upgradeTransaction.transactionId,
          transactionHash: transaction.hash,
        });
      }
    }
    if (
      context.mode === "execute"
      && implementationTransaction.route === "operator"
      && upgradeTransaction.route === "operator"
    ) {
      recorder?.apply({ ...plannedChange, txHash: upgradeTxHash });
    }
    context.state.set(`${proxy.saveAs}Upgrade`, {
      proxyAddress: proxy.proxyAddress,
      implementationAddress: prepared.implementationAddress,
      implementationSalt: prepared.implementationSalt,
      txHash: upgradeTxHash,
      external: upgradeTransaction.route === "external",
    });
    context.reporter.result("execution", {
      stepId: step.id,
      kind: step.kind,
      mode: upgradeTransaction.route,
      implementationTransactionId: implementationTransaction.transactionId,
      upgradeTransactionId: upgradeTransaction.transactionId,
      contract: contractName,
      actualContract,
      saveAs: proxy.saveAs,
      contractInfoKey: proxy.contractInfoKey,
      proxy: proxy.proxyAddress,
      implementationAddress: prepared.implementationAddress,
      txHash: upgradeTxHash,
    });
    return;
  }
  if (!context.signer || !context.signerAddress) {
    throw new Error(`Upgrade step ${step.id} requires a signer`);
  }
  const factory = await ethers.getContractFactory(contractName, context.signer);
  const proxyContract = await ethers.getContractAt(String(step.abi || contractName), proxy.proxyAddress, context.signer);
  const authority = await decideAuthority(proxyContract, context.target, context.signerAddress);
  if (context.mode === "simulate" && !context.configuredSimulationCaller) {
    assertSelectedSimulationSignerAuthorized(authority, step.id);
  }
  const recorder = context.mode === "execute" ? context.deploymentRecorder : undefined;
  if (context.mode === "execute" && !recorder) {
    throw new Error(`Upgrade step ${step.id} requires an execution deployment recorder`);
  }
  recorder?.prepare(preparedChange);
  const constructorArgs = resolveArray(
    context,
    step.constructorArgs || proxy.proxy.constructorArgs || [],
    `Step ${step.id} constructorArgs`,
  );
  const oldContractName =
    typeof step.old === "string"
      ? step.old
      : typeof step.oldContract === "string"
        ? step.oldContract
        : undefined;
  const gateToolOutputPath = path.join(
    context.taskDir,
    "results",
    context.targetId,
    `${step.id}-gate-tool-validate.json`,
  );
  let gateToolOutput: GateToolUpgradeOutput;
  if (context.mode === "simulate") {
    const deployment = await validateAndDeployUpgradeImplementation({
      proxyAddress: proxy.proxyAddress,
      oldContractName,
      newContractName: contractName,
      constructorArgs,
      unsafeAllow: proxy.proxy.unsafeAllow,
      proxyKind: proxy.proxy.kind,
      signer: context.signer,
    });
    const upgradeInterface = new ethers.Interface([
      "function upgradeToAndCall(address newImplementation, bytes data) payable",
    ]);
    gateToolOutput = {
      proxyAddress: proxy.proxyAddress,
      oldImplementation: deployment.oldImplementationAddress,
      newImplementation: deployment.implementationAddress,
      deployment: {
        method: "hardhat-upgrades-simulation",
        txHash: deployment.implementationTxHash || null,
      },
      upgradeToAndCall: {
        target: proxy.proxyAddress,
        calldata: upgradeInterface.encodeFunctionData("upgradeToAndCall", [
          deployment.implementationAddress,
          "0x",
        ]),
      },
      timestamp: new Date().toISOString(),
      network: context.targetId,
    };
    writeJsonFile(gateToolOutputPath, gateToolOutput);
  } else {
    gateToolOutput = runGateToolValidate({
      root: context.root,
      network: context.targetId,
      proxyAddress: proxy.proxyAddress,
      oldContractName,
      newContractName: contractName,
      outputPath: gateToolOutputPath,
      constructorArgs,
      unsafeAllow: proxy.proxy.unsafeAllow,
    });
  }
  const validatedChange: DeploymentChange = {
    ...preparedChange,
    proxyAddress: gateToolOutput.proxyAddress,
    implementationAddress: gateToolOutput.newImplementation,
    implementationTxHash: gateToolOutput.deployment?.txHash || undefined,
    factory: gateToolOutput.deployment?.factoryAddress || undefined,
    salt: gateToolOutput.deployment?.factorySalt || undefined,
    gateTool: {
      outputPath: path.relative(context.root, gateToolOutputPath),
      data: gateToolOutput as unknown as Record<string, unknown>,
    },
  };
  recorder?.prepare(validatedChange);
  let data = gateToolOutput.upgradeToAndCall.calldata;

  if (call) {
    const initData = factory.interface.encodeFunctionData(call.fn, call.args);
    const upgradeInterface = new ethers.Interface([
      "function upgradeToAndCall(address newImplementation, bytes data) payable",
    ]);
    data = upgradeInterface.encodeFunctionData("upgradeToAndCall", [
      gateToolOutput.newImplementation,
      initData,
    ]);
  }

  if (authority.mode === "execute" || context.configuredSimulationCaller) {
    const tx = await context.signer.sendTransaction({
      to: gateToolOutput.upgradeToAndCall.target,
      data,
    });
    await tx.wait();

    recorder?.apply({
      ...validatedChange,
      txHash: tx.hash,
    });
    context.state.set(`${proxy.saveAs}Upgrade`, {
      proxyAddress: gateToolOutput.proxyAddress,
      implementationAddress: gateToolOutput.newImplementation,
      txHash: tx.hash,
      gateToolOutput: gateToolOutputPath,
    });
    context.reporter.result("execution", {
      stepId: step.id,
      kind: step.kind,
      mode: "execute",
      contract: contractName,
      actualContract,
      saveAs: proxy.saveAs,
      contractInfoKey: proxy.contractInfoKey,
      proxy: gateToolOutput.proxyAddress,
      implementationAddress: gateToolOutput.newImplementation,
      txHash: tx.hash,
      gateToolOutput: gateToolOutputPath,
    });
    return;
  }

  const calldataOutput = {
    transactionId: step.id,
    stepId: step.id,
    to: gateToolOutput.upgradeToAndCall.target,
    value: "0",
    data,
    operation: 0 as const,
    method: "upgradeToAndCall",
    args: [gateToolOutput.newImplementation, call ? factory.interface.encodeFunctionData(call.fn, call.args) : "0x"],
    executor: authority.executor,
    reason: authority.reason,
  };

  if (context.mode === "execute") {
    const externalExecutor = calldataOutput.executor;
    if (!externalExecutor) {
      throw new Error(`Upgrade step ${step.id} cannot resolve an external executor`);
    }
    throw new ExternalExecutionRequired(step.id, [{
      ...calldataOutput,
      executor: externalExecutor,
    }]);
  }

  throw new Error(`Upgrade step ${step.id} requires a locked mixed-authority execution plan`);
};