import { Address, HexBigInt, HexNumber } from "../schemas.js";
import { strictObject, object, union, z } from "../zod.js";

const PolicyId = z.uuidv4({
  error: "Must be a valid UUID string",
});

const TokenGasFields = {
  tokenAddress: Address({
    description: "The address of the token to pay gas with",
  }),
  maxTokenAmount: HexBigInt({
    description:
      "The maximum amount of tokens to allow during ERC20 gas payment",
  }).optional(),
};

// Restructure as a Union at the top level with strict objects to enforce
// mutual exclusivity between preOpSettings and postOpSettings.
const Erc20PaymasterSettings = union(
  {
    description: "ERC20 Paymaster Settings",
    errorMessage:
      "ERC20 settings must specify only one of: preOpSettings, postOpSettings, or neither (for manual approvals). Cannot specify both preOpSettings and postOpSettings",
  },
  [
    strictObject(
      {
        ...TokenGasFields,
        preOpSettings: object(
          {
            autoPermit: object(
              {
                below: HexBigInt({
                  description:
                    "The amount below which a Permit signature will be injected",
                }),
                amount: HexBigInt({
                  description:
                    "The amount of tokens to Permit to the paymaster for gas payment",
                }),
                durationSeconds: HexNumber({
                  description:
                    "The duration of the Permit signature in seconds",
                }).optional(),
              },
              {
                description:
                  "Settings for automatically injecting a Permit signature for ERC20 gas payment",
              },
            ),
          },
          {
            description:
              "ERC20 Paymaster Settings for pre-op policies (autoPermit)",
          },
        ),
      },
      {
        description: "ERC20 Paymaster Settings with preOpSettings (autoPermit)",
      },
    ),
    strictObject(
      {
        ...TokenGasFields,
        preOpSettings: object(
          {
            permitDetails: object(
              {
                deadline: HexNumber({
                  description: "The deadline in the signed Permit object",
                }),
                value: HexBigInt({
                  description: "The value in the signed Permit object",
                }),
              },
              {
                description: "Details of the returned Permit signature",
              },
            ),
          },
          {
            description:
              "ERC20 Paymaster Settings for pre-op policies (permitDetails)",
          },
        ),
      },
      {
        description:
          "ERC20 Paymaster Settings with preOpSettings (permitDetails)",
      },
    ),
    strictObject(
      {
        ...TokenGasFields,
        postOpSettings: object(
          {
            autoApprove: z
              .union([
                z.literal(true).meta({
                  description:
                    "Automatically inject approval for the exact amount needed for this operation",
                }),
                object(
                  {
                    below: HexBigInt({
                      description:
                        "The amount below which an approval will be injected",
                    }),
                    amount: HexBigInt({
                      description:
                        "The amount of tokens to approve to the paymaster for gas payment",
                    }),
                  },
                  {
                    description:
                      "Settings for automatically injecting an approval for ERC20 gas payment with specific thresholds",
                  },
                ),
              ])
              .optional(),
            balanceCheck: z
              .literal(false)
              .meta({
                description:
                  "By default, the sender's ERC-20 balance is checked at prepare time and the request is rejected if it can't cover the maximum gas cost. Set to false when the User Operation itself credits the sender enough tokens to cover gas. If the sender's balance is still insufficient when the paymaster attempts to pull payment, the User Operation will fail on-chain.",
              })
              .optional(),
          },
          {
            description: "ERC20 Paymaster Settings for post-op policies",
          },
        ),
      },
      {
        description: "ERC20 Paymaster Settings with postOpSettings",
      },
    ),
    strictObject(
      {
        ...TokenGasFields,
      },
      {
        description:
          "No auto approval (requires manual ERC20 approvals to paymaster contract)",
      },
    ),
  ],
);

export const SinglePolicyId = object(
  {
    policyId: PolicyId.meta({
      description: "The policy ID to use",
    }),
  },
  {
    description: "Single policy ID",
  },
);

export const MultiplePolicyIds = object(
  {
    policyIds: z.array(PolicyId).meta({
      description: "The policy IDs to use",
    }),
  },
  {
    description: "Multiple policy IDs",
  },
);

export const WebhookData = z.string().meta({
  description:
    "Additional data you can include in the request, such as proof of humanity, if you have enabled custom rules in your Gas Manager policy.",
});

export const PaymasterCapability = z.intersection(
  z.union([SinglePolicyId, MultiplePolicyIds]),
  object({
    onlyEstimation: z
      .boolean()
      .meta({
        description:
          "If true, will only estimate fees. Skips signature request and gas sponsorship",
      })
      .optional(),
    erc20: Erc20PaymasterSettings.optional(),
    webhookData: WebhookData.optional(),
  }),
);

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