import { Address as AddressSchema } from "../schemas.js";
import {
  SMA_7702_DELEGATION_ADDRESSES,
  SemiModularAccount7702Version,
} from "../sma7702.js";
import { object, withMetadata, z } from "../zod.js";
import { zeroAddress, type Address } from "viem";

export {
  SMA_7702_DELEGATION_ADDRESSES,
  SMA_7702_VERSIONS,
  SemiModularAccount7702Version,
  getVersionForDelegationAddress7702,
} from "../sma7702.js";
export type { Sma7702Version } from "../sma7702.js";

export const EIP_7702_ACCOUNT_TYPE = ["ModularAccountV2"] as const;

export type Supported7702AccountType = (typeof EIP_7702_ACCOUNT_TYPE)[number];

// An EIP-7702 authorization to the zero address clears an account's delegation.
export const UNDELEGATION_ADDRESS = zeroAddress;

// Undelegation leads so this stays the non-empty tuple `z.enum` requires,
// whatever versions the registry supplies.
export const SUPPORTED_DELEGATION_ADDRESSES = [
  UNDELEGATION_ADDRESS,
  ...Object.values(SMA_7702_DELEGATION_ADDRESSES),
] as const;

const SUPPORTED_DELEGATION_ADDRESSES_LOWER = new Set<string>(
  SUPPORTED_DELEGATION_ADDRESSES.map((address) => address.toLowerCase()),
);

/**
 * Case-insensitive: an address is the same address whatever its casing, so
 * callers may send either the checksummed or the lowercase form.
 */
export const isSupportedDelegationAddress7702 = (address: string): boolean =>
  SUPPORTED_DELEGATION_ADDRESSES_LOWER.has(address.toLowerCase());

const supportedDelegationAddressError = `delegation must be one of: ${SUPPORTED_DELEGATION_ADDRESSES.join(", ")} (any casing), or one of: ${EIP_7702_ACCOUNT_TYPE.join(", ")}`;

const SupportedDelegationAddress = () =>
  withMetadata(
    AddressSchema({ error: supportedDelegationAddressError }).refine(
      isSupportedDelegationAddress7702,
      { error: supportedDelegationAddressError },
    ),
    { excludeFromApiSpecGeneration: true },
  );

export const Eip7702AuthCapability = z.union([
  object(
    {
      account: AddressSchema().optional(),
      // Wallet applications must maintain a strict shortlist of well-known accounts.
      // https://github.com/ethereum/ERCs/pull/947/files#diff-dab085f963ca621595044bcbc0922705aa38b69f5bf8770dcb3d0496a633efcfR189
      delegation: z.union([
        z.enum(EIP_7702_ACCOUNT_TYPE),
        // Documents every canonical delegate while the refinement below accepts
        // the equivalent lowercase and uppercase forms.
        z.enum(SUPPORTED_DELEGATION_ADDRESSES),
        SupportedDelegationAddress(),
      ]),
      version: SemiModularAccount7702Version.optional().meta({
        description:
          "The SemiModularAccount7702 version to delegate to. When omitted, uses the version stored for this team and wallet. If no record exists, a recognized SemiModularAccount7702 delegation already on the wallet is preserved; a wallet that is undelegated, or delegated to an unrecognized implementation, is delegated to the latest supported version. Must not contradict 'delegation' when that is an address.",
      }),
    },
    {
      description: "Specify EIP-7702 delegation",
    },
  ),
  z.literal(true).meta({
    description:
      "Use implicit EIP-7702 delegation. Uses the version stored for this team and wallet; if no record exists, a recognized SemiModularAccount7702 delegation already on the wallet is preserved, while a wallet that is undelegated or delegated to an unrecognized implementation is delegated to the latest supported version. Note: EIP-7702 is enabled by default unless you first call `wallet_requestAccount` to get a smart contract address.",
  }),
]);

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

// Every released SemiModularAccount7702 version is a ModularAccountV2 — the
// account type is unversioned, so all of them map back to the same type.
// Keyed by lowercase address so lookups are casing-independent.
export const DelegationAddressToAccountType: Record<
  string,
  Supported7702AccountType
> = Object.fromEntries(
  Object.values(SMA_7702_DELEGATION_ADDRESSES).map((address) => [
    address.toLowerCase(),
    "ModularAccountV2",
  ]),
);

export const getAccountTypeForDelegationAddress7702 = (
  address: Address,
): Supported7702AccountType | undefined => {
  return DelegationAddressToAccountType[address.toLowerCase()];
};