import z from "zod";

const SolTypeZod = z.enum([
  "uint8",
  "uint16",
  "uint32",
  "uint64",
  "uint128",
  "uint256",
  "int128",
  "int256",
  "address",
  "bool",
  "bytes",
  "bytes4",
  "bytes32",
  "string",
]);

export const ResourcePortZod = z.object({
  kind: z.literal("resource"),
  name: z.string(),
  accepts: z.enum(["erc20", "native", "any"]),
  mode: z.enum(["linear", "copy"]),
  optional: z.boolean().optional(),
});

export const HANDLE_UNITS = [
  "raw",
  "wad",
  "ray",
  "bps",
  "token-decimals",
] as const;

// NOTE: the uint members below must stay in sync with the `UintN` union in
// `ts/packages/lifi_ir1_ts/src/types/ValueHandle.ts`. When adding a new
// uintN width, update both locations.
export const STATIC_SOL_TYPES = [
  "uint256",
  "uint128",
  "uint64",
  "uint32",
  "uint16",
  "uint8",
  "address",
  "bool",
  "bytes32",
] as const;

export const HandleUnitsZod = z.enum(HANDLE_UNITS);

export const StaticSolTypeZod = z.enum(STATIC_SOL_TYPES);

const STATIC_SOL_TYPE_SET: ReadonlySet<string> = new Set(STATIC_SOL_TYPES);

export const isStaticSolType = (value: string): value is StaticSolType =>
  STATIC_SOL_TYPE_SET.has(value);

export const HandlePortZod = z.object({
  kind: z.literal("handle"),
  name: z.string(),
  type: SolTypeZod,
  mode: z.enum(["linear", "copy"]),
  expose: z.boolean().optional(),
  units: HandleUnitsZod.optional(),
});

export const InputPortZod = z.discriminatedUnion("kind", [
  ResourcePortZod,
  HandlePortZod,
]);

export const ResourceOutputPortZod = z.object({
  kind: z.literal("resource_output"),
  name: z.string(),
  mode: z.enum(["linear", "copy"]),
  availability: z.enum(["now", "future"]).optional(),
  providesMinimum: z.boolean().optional(),
  omitIfZero: z.boolean().optional(),
  deliveryAddressInput: z.string().optional(),
});

export const OutputPortZod = z.discriminatedUnion("kind", [
  ResourceOutputPortZod,
  HandlePortZod,
]);

export const ManifestOperationZod = z.object({
  id: z.string(),
  description: z.string().optional(),
  inputs: z.array(InputPortZod),
  outputs: z.array(OutputPortZod),
  configSchema: z.unknown().optional(),
});

const PortKindEnum = z.enum(["resource", "resource_output", "handle"]);

const SelectorMatchZod = z.object({
  kind: z.union([PortKindEnum, z.array(PortKindEnum)]),
  mode: z.enum(["linear", "copy"]).optional(),
  type: z.enum(["erc20", "native", "any"]).optional(),
});

const ConfigSelectionZod = z.object({
  kind: z.literal("config"),
  configKey: z.string(),
  cardinality: z.enum(["one", "many"]),
});

const AllMatchingSelectionZod = z.object({
  kind: z.literal("all_matching"),
});

const SelectorSelectionZod = z.discriminatedUnion("kind", [
  ConfigSelectionZod,
  AllMatchingSelectionZod,
]);

export const GuardSelectorZod = z.object({
  binding: z.string(),
  source: z.enum(["inputs", "outputs"]),
  match: SelectorMatchZod,
  selection: SelectorSelectionZod,
});

export const GuardCompatibilityZod = z.object({
  selectors: z.array(GuardSelectorZod),
});

export const ManifestGuardZod = z.object({
  kind: z.string(),
  description: z.string().optional(),
  configSchema: z.unknown().optional(),
  compatibility: GuardCompatibilityZod.optional(),
});

export const ManifestMaterialiserZod = z.object({
  kind: z.string(),
  description: z.string().optional(),
  accepts: z.enum(["resource", "handle", "any"]),
  configSchema: z.unknown().optional(),
});

export const ManifestPreconditionZod = z.object({
  type: z.string(),
  description: z.string().optional(),
  configSchema: z.unknown().optional(),
});

export const ComposeManifestZod = z.object({
  manifestVersion: z.number(),
  manifestHash: z.string(),
  flowSchema: z.object({}).passthrough(),
  operations: z.array(ManifestOperationZod),
  guards: z.array(ManifestGuardZod),
  materialisers: z.array(ManifestMaterialiserZod),
  preconditions: z.array(ManifestPreconditionZod).optional(),
});

export type ResourcePort = z.infer<typeof ResourcePortZod>;
export type HandlePort = z.infer<typeof HandlePortZod>;
export type HandleUnits = z.infer<typeof HandleUnitsZod>;
export type StaticSolType = z.infer<typeof StaticSolTypeZod>;
export type OpInputPort = z.infer<typeof InputPortZod>;
export type ResourceOutputPort = z.infer<typeof ResourceOutputPortZod>;
export type OpOutputPort = z.infer<typeof OutputPortZod>;
export type ManifestOperation = z.infer<typeof ManifestOperationZod>;
export type GuardSelector = z.infer<typeof GuardSelectorZod>;
export type GuardCompatibility = z.infer<typeof GuardCompatibilityZod>;
export type ManifestGuard = z.infer<typeof ManifestGuardZod>;
export type ManifestMaterialiser = z.infer<typeof ManifestMaterialiserZod>;
export type ManifestPrecondition = z.infer<typeof ManifestPreconditionZod>;
export type ComposeManifest = z.infer<typeof ComposeManifestZod>;