import z from "zod";

export const JsonRpcVersionSchema = z.literal("2.0");

export const JsonRpcIdSchema = z.union([
  z.string().min(1),
  z.number().finite(),
]);

export const JsonRpcIdempotencyMetadataSchema = z
  .object({
    idempotencyKey: z.string().min(1).optional(),
    fingerprint: z.string().min(1).optional(),
  })
  .strict();

export const JsonRpcRequestSchema = z.object({
  jsonrpc: JsonRpcVersionSchema,
  id: JsonRpcIdSchema,
  method: z.string().min(1),
  params: z.unknown().optional(),
  meta: JsonRpcIdempotencyMetadataSchema.optional(),
});

export const JsonRpcNotificationSchema = z.object({
  jsonrpc: JsonRpcVersionSchema,
  method: z.string().min(1),
  params: z.unknown().optional(),
  meta: JsonRpcIdempotencyMetadataSchema.optional(),
});

export const JsonRpcErrorObjectSchema = z.object({
  code: z.number().int(),
  message: z.string(),
  data: z.unknown().optional(),
});

export const JsonRpcSuccessResponseSchema = z.object({
  jsonrpc: JsonRpcVersionSchema,
  id: JsonRpcIdSchema,
  result: z.unknown().optional(),
});

export const JsonRpcErrorResponseSchema = z.object({
  jsonrpc: JsonRpcVersionSchema,
  id: JsonRpcIdSchema,
  error: JsonRpcErrorObjectSchema,
});

export const JsonRpcResponseSchema = z.union([
  JsonRpcSuccessResponseSchema,
  JsonRpcErrorResponseSchema,
]);

export const JsonRpcMessageSchema = z.union([
  JsonRpcRequestSchema,
  JsonRpcNotificationSchema,
  JsonRpcResponseSchema,
]);

export const JsonRpcBatchSchema = z.array(JsonRpcMessageSchema).min(1);

export const JsonRpcPayloadSchema = z.union([
  JsonRpcMessageSchema,
  JsonRpcBatchSchema,
]);

export type JsonRpcId = z.output<typeof JsonRpcIdSchema>;
export type JsonRpcIdempotencyMetadata = z.output<
  typeof JsonRpcIdempotencyMetadataSchema
>;
export type JsonRpcRequest = z.output<typeof JsonRpcRequestSchema>;
export type JsonRpcNotification = z.output<typeof JsonRpcNotificationSchema>;
export type JsonRpcErrorObject = z.output<typeof JsonRpcErrorObjectSchema>;
export type JsonRpcSuccessResponse = z.output<
  typeof JsonRpcSuccessResponseSchema
>;
export type JsonRpcErrorResponse = z.output<typeof JsonRpcErrorResponseSchema>;
export type JsonRpcResponse = z.output<typeof JsonRpcResponseSchema>;
export type JsonRpcMessage = z.output<typeof JsonRpcMessageSchema>;
export type JsonRpcPayload = z.output<typeof JsonRpcPayloadSchema>;

export const parseJsonRpcPayload = (value: unknown): JsonRpcPayload | null => {
  const parsed = JsonRpcPayloadSchema.safeParse(value);
  return parsed.success ? parsed.data : null;
};

export const makeJsonRpcRequest = ({
  id,
  method,
  params,
  meta,
}: {
  id: JsonRpcId;
  method: string;
  params?: unknown;
  meta?: JsonRpcIdempotencyMetadata;
}): JsonRpcRequest => ({
  jsonrpc: "2.0",
  id,
  method,
  ...(params === undefined ? {} : { params }),
  ...(meta === undefined ? {} : { meta }),
});

export const makeJsonRpcNotification = ({
  method,
  params,
  meta,
}: {
  method: string;
  params?: unknown;
  meta?: JsonRpcIdempotencyMetadata;
}): JsonRpcNotification => ({
  jsonrpc: "2.0",
  method,
  ...(params === undefined ? {} : { params }),
  ...(meta === undefined ? {} : { meta }),
});

export const makeJsonRpcSuccessResponse = ({
  id,
  result,
}: {
  id: JsonRpcId;
  result?: unknown;
}): JsonRpcSuccessResponse => ({
  jsonrpc: "2.0",
  id,
  ...(result === undefined ? {} : { result }),
});

export const makeJsonRpcErrorResponse = ({
  id,
  code,
  message,
  data,
}: {
  id: JsonRpcId;
  code: number;
  message: string;
  data?: unknown;
}): JsonRpcErrorResponse => ({
  jsonrpc: "2.0",
  id,
  error: {
    code,
    message,
    ...(data === undefined ? {} : { data }),
  },
});