import { z } from 'zod';

export const DEFAULT_HEARTBEAT_INTERVAL_MS = 15_000;

const pairOkServerMessageSchema = z.object({
  type: z.literal('pair.ok'),
  projectName: z.string().min(1),
  reconnectCredential: z.string().min(1),
  connectionEpoch: z.number().int().positive(),
});

const reconnectOkServerMessageSchema = z.object({
  type: z.literal('reconnect.ok'),
  projectName: z.string().min(1),
  connectionEpoch: z.number().int().positive(),
});

const taskDispatchServerMessageSchema = z.object({
  type: z.literal('task.dispatch'),
  taskId: z.string().min(1),
  prompt: z.string().min(1),
});

const taskResultClientMessageSchema = z.object({
  type: z.literal('task.result'),
  taskId: z.string().min(1),
  status: z.enum(['accepted', 'rejected']),
  reason: z.string().min(1).optional(),
});

const taskCompletedClientMessageSchema = z.object({
  type: z.literal('task.completed'),
  taskId: z.string().min(1),
  text: z.string(),
});

const taskFailedClientMessageSchema = z.object({
  type: z.literal('task.failed'),
  taskId: z.string().min(1),
  message: z.string().min(1),
});

const errorServerMessageSchema = z.object({
  type: z.literal('error'),
  code: z.string().min(1),
  message: z.string().min(1),
});

export const localAgentServerMessageSchema = z.discriminatedUnion('type', [
  pairOkServerMessageSchema,
  reconnectOkServerMessageSchema,
  taskDispatchServerMessageSchema,
  errorServerMessageSchema,
]);

export type LocalAgentServerMessage = z.infer<typeof localAgentServerMessageSchema>;

export const localAgentClientMessageSchema = z.discriminatedUnion('type', [
  taskResultClientMessageSchema,
  taskCompletedClientMessageSchema,
  taskFailedClientMessageSchema,
]);

export type LocalAgentClientMessage = z.infer<typeof localAgentClientMessageSchema>;