import type { ContractFactory, Signer } from "ethers";
import { concat, getAddress, keccak256, toUtf8Bytes } from "ethers";
import { artifacts, ethers } from "hardhat";
import {
  Create2Options,
  deployCreate2InitCode,
  prepareCreate2Call,
  PreparedCreate2Call,
  PostDeployCall,
  toOpenZeppelinCreate2Config,
} from "./create2";
import type { ExternalCall } from "./external";

const IMPLEMENTATION_SLOT =
  "0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc";

export interface ProxyOptions {
  kind?: string;
  initializer?: string | false;
  unsafeAllow?: string[];
  constructorArgs?: unknown[];
}

export interface ProxyDeployment {
  proxyAddress: string;
  implementationAddress: string;
  predictedAddress?: string;
  alreadyDeployed?: boolean;
  firstUpgradeImplementationAddress?: string;
  firstUpgradeTxHash?: string;
  txHash?: string;
}

export interface FirstUpgradeOptions {
  contractName: string;
  call?: {
    fn: string;
    args: unknown[];
  };
  constructorArgs?: unknown[];
  unsafeAllow?: string[];
}

export interface FirstUpgradePrepared {
  implementationAddress: string;
  implementationTxHash?: string;
  postDeployCalls: PostDeployCall[];
}

interface DeployableContract {
  waitForDeployment: () => Promise<unknown>;
  getAddress: () => Promise<string>;
  deploymentTransaction?: () => { hash: string } | null;
}

interface UpgradesLike {
  deployProxy: (
    factory: ContractFactory,
    args: unknown[],
    options: Record<string, unknown>,
  ) => Promise<DeployableContract>;
  deployImplementation?: (
    factory: ContractFactory,
    options: Record<string, unknown>,
  ) => Promise<string | DeployableContract>;
  validateImplementation?: (
    factory: ContractFactory,
    options: Record<string, unknown>,
  ) => Promise<void>;
  validateUpgrade?: (
    proxyAddress: string,
    factory: ContractFactory,
    options: Record<string, unknown>,
  ) => Promise<void>;
  forceImport?: (
    proxyAddress: string,
    factory: ContractFactory,
    options: Record<string, unknown>,
  ) => Promise<unknown>;
  erc1967: {
    getImplementationAddress: (proxyAddress: string) => Promise<string>;
    getAdminAddress?: (proxyAddress: string) => Promise<string>;
  };
}

export interface PreparedUUPSCalldataDeployment {
  implementationSalt: string;
  implementationAddress: string;
  proxyAddress: string;
  calls: Array<Omit<ExternalCall, "executor">>;
}

export interface PreparedExternalProxyUpgrade {
  implementationSalt: string;
  implementationAddress: string;
  calls: Array<Omit<ExternalCall, "executor">>;
}

export interface ValidatedUpgradeImplementation {
  oldImplementationAddress: string;
  implementationAddress: string;
  implementationTxHash?: string;
}

const getUpgrades = async (): Promise<UpgradesLike> => {
  const hardhat = await import("hardhat") as typeof import("hardhat") & {
    upgrades?: UpgradesLike;
  };

  if (!hardhat.upgrades) {
    throw new Error(
      "OpenZeppelin upgrades plugin is required for proxy workflow steps. " +
        "Install @openzeppelin/hardhat-upgrades and import it in hardhat.config.ts.",
    );
  }

  return hardhat.upgrades;
};

const normalizeDeploymentAddress = async (
  deployment: string | DeployableContract,
): Promise<{ address: string; txHash?: string }> => {
  if (typeof deployment === "string") {
    return { address: deployment };
  }

  await deployment.waitForDeployment();
  return {
    address: await deployment.getAddress(),
    txHash: deployment.deploymentTransaction?.()?.hash,
  };
};

const buildProxyOptions = (
  proxy: ProxyOptions,
  create2?: Create2Options,
): Record<string, unknown> => {
  const options: Record<string, unknown> = {
    initializer: proxy.initializer === undefined ? "initialize" : proxy.initializer,
    kind: proxy.kind || "uups",
  };

  if (proxy.unsafeAllow) {
    options.unsafeAllow = proxy.unsafeAllow;
  }
  if (proxy.constructorArgs) {
    options.constructorArgs = proxy.constructorArgs;
  }
  if (create2) {
    options.create2Factory = toOpenZeppelinCreate2Config(create2);
  }

  return options;
};

const implementationSalt = (
  executionId: string,
  contractName: string,
  scope: string,
): string => {
  return keccak256(toUtf8Bytes("gateflow:" + executionId + ":" + contractName + ":" + scope));
};

const contractFactoryWithoutSigner = async (
  contractName: string,
): Promise<{ factory: ContractFactory; initCodePrefix: string }> => {
  const artifact = await artifacts.readArtifact(contractName);
  return {
    factory: new ethers.ContractFactory(artifact.abi, artifact.bytecode),
    initCodePrefix: artifact.bytecode,
  };
};

const prepareImplementationCreate2 = async (options: {
  executionId: string;
  stepId: string;
  contractName: string;
  constructorArgs: unknown[];
  unsafeAllow?: string[];
  create2: Create2Options;
  scope: string;
  validate: (factory: ContractFactory, options: Record<string, unknown>) => Promise<void>;
}): Promise<{
  salt: string;
  call: PreparedCreate2Call;
  factory: ContractFactory;
}> => {
  const { factory, initCodePrefix } = await contractFactoryWithoutSigner(options.contractName);
  await options.validate(factory, {
    kind: "uups",
    constructorArgs: options.constructorArgs,
    unsafeAllow: options.unsafeAllow,
  });
  const encodedArgs = factory.interface.encodeDeploy(options.constructorArgs);
  const initCode = concat([initCodePrefix, encodedArgs]);
  const salt = implementationSalt(
    options.executionId,
    options.contractName,
    options.scope,
  );
  const call = prepareCreate2Call(
    options.contractName + " implementation",
    initCode,
    {
      factory: options.create2.factory,
      salt,
      leadingZeroBytes: undefined,
      expectedAddress: undefined,
    },
  );
  return { salt, call, factory };
};

const preparedCreate2ExternalCall = (
  stepId: string,
  prepared: PreparedCreate2Call,
  reason: string,
): Omit<ExternalCall, "executor"> => ({
  transactionId: stepId,
  stepId,
  to: prepared.target,
  value: prepared.value,
  data: prepared.data,
  operation: 0,
  method: "deploy",
  args: [prepared.salt, prepared.initCode],
  reason,
});

export const prepareUUPSCalldataDeployment = async (options: {
  executionId: string;
  stepId: string;
  contractName: string;
  firstUpgrade: FirstUpgradeOptions;
  proxy: ProxyOptions;
  create2: Create2Options;
}): Promise<PreparedUUPSCalldataDeployment> => {
  if (!options.create2.uupsPlaceholderImplementation) {
    throw new Error("uups-placeholder-proxy requires target.uups.placeholderImplementation");
  }
  const upgrades = await getUpgrades();
  if (!upgrades.validateImplementation) {
    throw new Error("OpenZeppelin upgrades validateImplementation is required for calldata mode");
  }
  const implementation = await prepareImplementationCreate2({
    executionId: options.executionId,
    stepId: options.stepId,
    contractName: options.contractName,
    constructorArgs: options.firstUpgrade.constructorArgs || [],
    unsafeAllow: options.firstUpgrade.unsafeAllow,
    create2: options.create2,
    scope: options.create2.salt,
    validate: (factory, validateOptions) => (
      upgrades.validateImplementation!(factory, validateOptions)
    ),
  });
  const placeholderImplementation = getAddress(options.create2.uupsPlaceholderImplementation);
  const proxyArtifact = await artifacts.readArtifact(
    "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol:ERC1967Proxy",
  );
  const proxyConstructorArgs = ethers.AbiCoder.defaultAbiCoder().encode(
    ["address", "bytes"],
    [placeholderImplementation, "0x"],
  );
  const proxyInitCode = concat([proxyArtifact.bytecode, proxyConstructorArgs]);
  const proxyCall = prepareCreate2Call(
    "UUPS placeholder proxy",
    proxyInitCode,
    options.create2,
  );
  const initData = options.firstUpgrade.call
    ? implementation.factory.interface.encodeFunctionData(
        options.firstUpgrade.call.fn,
        options.firstUpgrade.call.args,
      )
    : "0x";
  const upgradeInterface = new ethers.Interface([
    "function upgradeToAndCall(address newImplementation, bytes data) payable",
  ]);
  const upgradeData = upgradeInterface.encodeFunctionData("upgradeToAndCall", [
    implementation.call.predictedAddress,
    initData,
  ]);
  return {
    implementationSalt: implementation.salt,
    implementationAddress: implementation.call.predictedAddress,
    proxyAddress: proxyCall.predictedAddress,
    calls: [
      preparedCreate2ExternalCall(
        options.stepId + "-implementation",
        implementation.call,
        "deploy UUPS implementation for explicit calldata mode",
      ),
      preparedCreate2ExternalCall(
        options.stepId + "-proxy",
        proxyCall,
        "deploy UUPS placeholder proxy for explicit calldata mode",
      ),
      {
        transactionId: options.stepId + "-first-upgrade",
        stepId: options.stepId + "-first-upgrade",
        to: proxyCall.predictedAddress,
        value: "0",
        data: upgradeData,
        operation: 0,
        method: "upgradeToAndCall",
        args: [implementation.call.predictedAddress, initData],
        reason: "activate UUPS implementation for explicit calldata mode",
      },
    ],
  };
};

export const prepareExternalProxyUpgrade = async (options: {
  executionId: string;
  stepId: string;
  proxyAddress: string;
  contractName: string;
  constructorArgs: unknown[];
  unsafeAllow?: string[];
  create2: Create2Options;
  call?: { fn: string; args: unknown[] };
}): Promise<PreparedExternalProxyUpgrade> => {
  const upgrades = await getUpgrades();
  if (!upgrades.validateUpgrade) {
    throw new Error("OpenZeppelin upgrades validateUpgrade is required for calldata mode");
  }
  const implementation = await prepareImplementationCreate2({
    executionId: options.executionId,
    stepId: options.stepId,
    contractName: options.contractName,
    constructorArgs: options.constructorArgs,
    unsafeAllow: options.unsafeAllow,
    create2: options.create2,
    scope: options.proxyAddress,
    validate: (factory, validateOptions) => (
      upgrades.validateUpgrade!(options.proxyAddress, factory, validateOptions)
    ),
  });
  const initData = options.call
    ? implementation.factory.interface.encodeFunctionData(options.call.fn, options.call.args)
    : "0x";
  const upgradeInterface = new ethers.Interface([
    "function upgradeToAndCall(address newImplementation, bytes data) payable",
  ]);
  return {
    implementationSalt: implementation.salt,
    implementationAddress: implementation.call.predictedAddress,
    calls: [
      preparedCreate2ExternalCall(
        options.stepId + "-implementation",
        implementation.call,
        "deploy upgrade implementation for explicit calldata mode",
      ),
      {
        transactionId: options.stepId,
        stepId: options.stepId,
        to: options.proxyAddress,
        value: "0",
        data: upgradeInterface.encodeFunctionData("upgradeToAndCall", [
          implementation.call.predictedAddress,
          initData,
        ]),
        operation: 0,
        method: "upgradeToAndCall",
        args: [implementation.call.predictedAddress, initData],
        reason: "upgrade proxy for explicit calldata mode",
      },
    ],
  };
};

export const getImplementationAddress = async (proxyAddress: string): Promise<string> => {
  const upgrades = await getUpgrades();

  try {
    return await upgrades.erc1967.getImplementationAddress(proxyAddress);
  } catch {
    const rawImplementationAddress = await ethers.provider.getStorage(proxyAddress, IMPLEMENTATION_SLOT);
    return `0x${rawImplementationAddress.slice(-40)}`;
  }
};

export const getProxyAdminAddress = async (proxyAddress: string): Promise<string | undefined> => {
  const upgrades = await getUpgrades();
  if (!upgrades.erc1967.getAdminAddress) return undefined;

  try {
    return await upgrades.erc1967.getAdminAddress(proxyAddress);
  } catch {
    return undefined;
  }
};

export const deployImplementationContract = async (
  contractName: string,
  constructorArgs: unknown[],
  proxy: Pick<ProxyOptions, "kind" | "unsafeAllow">,
  signer: Signer,
): Promise<{ address: string; txHash?: string }> => {
  const upgrades = await getUpgrades();
  const factory = await ethers.getContractFactory(contractName, signer);
  const options: Record<string, unknown> = {
    kind: proxy.kind || "uups",
    constructorArgs,
  };

  if (proxy.unsafeAllow) {
    options.unsafeAllow = proxy.unsafeAllow;
  }

  if (upgrades.deployImplementation) {
    return normalizeDeploymentAddress(await upgrades.deployImplementation(factory, options));
  }

  const deployed = await factory.deploy(...constructorArgs) as unknown as DeployableContract;
  return normalizeDeploymentAddress(deployed);
};

export const validateAndDeployUpgradeImplementation = async (options: {
  proxyAddress: string;
  oldContractName?: string;
  newContractName: string;
  constructorArgs: unknown[];
  unsafeAllow?: string[];
  proxyKind?: string;
  signer: Signer;
}): Promise<ValidatedUpgradeImplementation> => {
  const upgrades = await getUpgrades();
  if (!upgrades.validateUpgrade) {
    throw new Error("OpenZeppelin upgrades validateUpgrade is required for upgrade simulation");
  }
  const validateOptions: Record<string, unknown> = {
    kind: options.proxyKind || "uups",
    constructorArgs: options.constructorArgs,
  };
  if (options.unsafeAllow) validateOptions.unsafeAllow = options.unsafeAllow;

  const oldImplementationAddress = await getImplementationAddress(options.proxyAddress);
  if (upgrades.forceImport) {
    const oldFactory = await ethers.getContractFactory(
      options.oldContractName || options.newContractName,
      options.signer,
    );
    await upgrades.forceImport(options.proxyAddress, oldFactory, validateOptions);
  }

  const factory = await ethers.getContractFactory(options.newContractName, options.signer);
  await upgrades.validateUpgrade(options.proxyAddress, factory, validateOptions);
  const deployment = upgrades.deployImplementation
    ? await normalizeDeploymentAddress(
        await upgrades.deployImplementation(factory, validateOptions),
      )
    : await normalizeDeploymentAddress(
        await factory.deploy(...options.constructorArgs) as unknown as DeployableContract,
      );
  if (await ethers.provider.getCode(deployment.address) === "0x") {
    throw new Error(
      `Upgrade simulation implementation was not deployed on the active fork: ${deployment.address}`,
    );
  }
  return {
    oldImplementationAddress,
    implementationAddress: deployment.address,
    implementationTxHash: deployment.txHash,
  };
};

export const prepareFirstUpgrade = async (
  firstUpgrade: FirstUpgradeOptions,
  proxy: Pick<ProxyOptions, "kind">,
  signer: Signer,
): Promise<FirstUpgradePrepared> => {
  const implementation = await deployImplementationContract(
    firstUpgrade.contractName,
    firstUpgrade.constructorArgs || [],
    {
      kind: proxy.kind,
      unsafeAllow: firstUpgrade.unsafeAllow,
    },
    signer,
  );
  const factory = await ethers.getContractFactory(firstUpgrade.contractName, signer);
  const callData = firstUpgrade.call
    ? factory.interface.encodeFunctionData(firstUpgrade.call.fn, firstUpgrade.call.args)
    : "0x";
  const upgradeInterface = new ethers.Interface([
    "function upgradeToAndCall(address newImplementation, bytes data) payable",
  ]);

  return {
    implementationAddress: implementation.address,
    implementationTxHash: implementation.txHash,
    postDeployCalls: [
      {
        target: "$proxy",
        value: "0",
        data: upgradeInterface.encodeFunctionData("upgradeToAndCall", [
          implementation.address,
          callData,
        ]),
      },
    ],
  };
};

export const deployProxyContract = async (
  contractName: string,
  args: unknown[],
  proxy: ProxyOptions,
  create2?: Create2Options,
  signer?: Signer,
): Promise<ProxyDeployment> => {
  const upgrades = await getUpgrades();
  const factory = await ethers.getContractFactory(contractName, signer);
  const contract = await upgrades.deployProxy(factory, args, buildProxyOptions(proxy, create2));
  await contract.waitForDeployment();

  const proxyAddress = await contract.getAddress();
  const implementationAddress = await getImplementationAddress(proxyAddress);

  return {
    proxyAddress,
    implementationAddress,
    txHash: contract.deploymentTransaction?.()?.hash,
  };
};

export const deployUUPSPlaceholderProxyContract = async (
  create2: Create2Options,
  signer: Signer,
): Promise<ProxyDeployment> => {
  if (!create2.uupsPlaceholderImplementation) {
    throw new Error("uups-placeholder-proxy requires target.uups.placeholderImplementation");
  }

  const placeholderImplementation = getAddress(create2.uupsPlaceholderImplementation);
  const proxyArtifact = await artifacts.readArtifact(
    "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol:ERC1967Proxy",
  );
  const constructorArgs = ethers.AbiCoder.defaultAbiCoder().encode(
    ["address", "bytes"],
    [placeholderImplementation, "0x"],
  );
  const proxyInitCode = concat([proxyArtifact.bytecode, constructorArgs]);
  const deployment = await deployCreate2InitCode(
    "UUPS placeholder proxy",
    proxyInitCode,
    create2,
    signer,
  );
  const implementationAddress = await getImplementationAddress(deployment.address);

  return {
    proxyAddress: deployment.address,
    implementationAddress,
    predictedAddress: deployment.predictedAddress,
    alreadyDeployed: deployment.alreadyDeployed,
    txHash: deployment.txHash,
  };
};