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

const createPermission = <TData extends z.ZodType, TType extends string>(
  typeString: TType,
  data: TData,
  description: string,
) => {
  return object(
    {
      type: z.literal(typeString),
      data,
    },
    { description, title: `${typeString}` },
  );
};

const NativeTokenTransfer = createPermission(
  "native-token-transfer",
  object({
    allowance: HexBigInt(),
  }),
  "allows transfer of native tokens from the account",
);

const Erc20TokenTransfer = createPermission(
  "erc20-token-transfer",
  object({
    allowance: HexBigInt(),
    address: Address(),
  }),
  "allows transfer or approval of erc20 tokens from the account",
);

const GasLimit = createPermission(
  "gas-limit",
  object({
    limit: HexBigInt(),
  }),
  "allows the key to spend gas for UOs",
);

const ContractAccess = createPermission(
  "contract-access",
  object({
    address: Address(),
  }),
  "grants access to all functions in a contract",
);

const AccountFunction = createPermission(
  "account-functions",
  object({
    functions: z.array(Hex()).min(1),
  }),
  "grants access to functions in the account",
);

const AllExternalFunctionAccess = createPermission(
  "functions-on-all-contracts",
  object({
    functions: z.array(Hex()).min(1),
  }),
  "access to a function selector in any address or contract",
);

const ContractFunctionAccess = createPermission(
  "functions-on-contract",
  object({
    address: Address(),
    functions: z.array(Hex()).min(1),
  }),
  "grants access to specified functions on a specific contract",
);

const Root = createPermission(
  "root",
  z.never().optional(),
  "grants full access to everything",
);

export const Permission = z.discriminatedUnion("type", [
  // 7715 permissions
  NativeTokenTransfer,
  Erc20TokenTransfer,
  GasLimit,
  // additional permissions
  ContractAccess,
  AccountFunction,
  AllExternalFunctionAccess,
  ContractFunctionAccess,
  Root,
]);

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

export const PermissionsArray = z.array(Permission).min(1);

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

export const PermissionsContext = object(
  {
    context: Hex(),
  },
  {
    description: "Permissions context",
  },
);

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

export const RemotePermission = object(
  {
    sessionId: Hex(),
    signature: Hex(),
  },
  {
    description: "Remote permission",
  },
);

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

export const PermissionsData = object({
  chainId: HexChainId(),
  expirySec: z
    .int()
    .nonnegative()
    .max(9999999999)
    .meta({
      description: "Unix timestamp in seconds (0 for no expiry)",
    })
    .optional(),
  key: KeySigner,
  permissions: PermissionsArray,
});

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

export const PermissionsCapability = z.union([
  PermissionsContext,
  RemotePermission,
  // TODO: support permissions object here instead of only context-encoded permissions
  // PermissionsData,
]);

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