import { Address, Hex, HexBigInt, HexNumber } from "../schemas.js";
import { strictObject, z } from "../zod.js";
import { getAddress, hexToBigInt, hexToNumber, type StateOverride } from "viem";

// https://geth.ethereum.org/docs/interacting-with-geth/rpc/objects#state-override-set

// State mapping: Record of slot -> value
const StateMapping = z.record(Hex(), Hex()).meta({
  description: "Mapping from storage slot (hex) to value (hex)",
});

// Base properties that can be used with any override variant
const baseProperties = {
  balance: HexBigInt().optional(),
  nonce: HexNumber().optional(),
  code: Hex().optional(),
};

export const AccountOverride = z.union([
  strictObject(
    {
      ...baseProperties,
      state: StateMapping,
    },
    { description: "Full account state" },
  ),
  strictObject(
    {
      ...baseProperties,
      stateDiff: StateMapping,
    },
    { description: "Account state diff" },
  ),
  strictObject(baseProperties, {
    description: "Account override without state changes",
  }),
]);

export const StateOverrideCapability = z
  .record(Address(), AccountOverride)
  .meta({
    description: "Mapping from account address to overrides",
  });

export type StateOverrideCapability = z.input<typeof StateOverrideCapability>;

export function toStateOverride(
  override: StateOverrideCapability,
): StateOverride {
  return Object.entries(override).map(([address, accountOverride]) => {
    const result: StateOverride[number] = {
      address: getAddress(address),
    };

    if (accountOverride.balance !== undefined) {
      result.balance = hexToBigInt(accountOverride.balance);
    }

    if (accountOverride.nonce !== undefined) {
      result.nonce = hexToNumber(accountOverride.nonce);
    }

    if (accountOverride.code !== undefined) {
      result.code = accountOverride.code;
    }

    // Convert state from RPC format (Record) to Viem format (Array)
    if ("state" in accountOverride && accountOverride.state !== undefined) {
      result.state = Object.entries(accountOverride.state).map(
        ([slot, value]) => ({
          slot: slot as Hex,
          value,
        }),
      );
    } else if (
      "stateDiff" in accountOverride &&
      accountOverride.stateDiff !== undefined
    ) {
      // Convert stateDiff from RPC format (Record) to Viem format (Array)
      result.stateDiff = Object.entries(accountOverride.stateDiff).map(
        ([slot, value]) => ({
          slot: slot as Hex,
          value,
        }),
      );
    }

    return result;
  });
}