import { Address, Hex, HexBigInt, HexChainId } from "../schemas.js";
import { PrepareCallsCapabilities } from "../capabilities/index.js";
import { object, z } from "../zod.js";

const BaseParams = object({
  from: Address(),
  chainId: HexChainId(),
  fromToken: Address(),
  toToken: Address(),
  slippage: HexBigInt().optional(),
});

const SameChainParams = object({
  postCalls: z
    .array(
      object({
        to: Address(),
        data: Hex().optional(),
        value: HexBigInt().optional(),
      }),
    )
    .optional(),
  toChainId: z.never().optional(), // Must be undefined for same-chain swaps.
});

const CrossChainParams = object({
  postCalls: z.never().optional(), // Not supported for cross-chain swaps.
  toChainId: HexChainId(),
});

const FromAmountParams = object({ fromAmount: HexBigInt() });

const MinimumToAmountParams = object({ minimumToAmount: HexBigInt() });

const SmartWalletExecutionParams = object({
  returnRawCalls: z
    .literal(false)
    .meta({
      description: "Must be 'false' or undefined",
    })
    .optional(),
  capabilities: PrepareCallsCapabilities.optional(),
});

const EoaExecutionParams = object({
  returnRawCalls: z.literal(true).meta({
    description: "Must be 'true'",
  }),
});

export const SwapParams = {
  Base: BaseParams,
  Chain: {
    Same: SameChainParams,
    Cross: CrossChainParams,
  },
  AmountSpecification: {
    From: FromAmountParams,
    MinimumTo: MinimumToAmountParams,
  },
  Execution: {
    SmartWallet: SmartWalletExecutionParams,
    Eoa: EoaExecutionParams,
  },
  variant(opts: {
    chain: "Same" | "Cross";
    execution: "SmartWallet" | "Eoa";
    amount: "From" | "MinimumTo";
    description: string;
  }) {
    return z
      .intersection(
        z.intersection(this.Base, this.Chain[opts.chain]),
        z.intersection(
          this.Execution[opts.execution],
          this.AmountSpecification[opts.amount],
        ),
      )
      .describe(opts.description);
  },
};