import { z } from "zod";

const id = z.string().trim().min(1).max(200);
const safeUnknown = z.unknown();

const base = {
  eventId: id,
  sessionId: id,
  runId: id.optional(),
  timestamp: z.string().datetime({ offset: true }),
  sequence: z.int().positive(),
};

function event<const T extends string, P extends z.ZodType>(type: T, payload: P) {
  return z.object({ ...base, type: z.literal(type), payload }).strict();
}

const sessionStarted = event("session.started", z.object({
  cwd: z.string().trim().min(1).max(1000),
}).strict());
const sessionEnded = event("session.ended", z.object({
  reason: z.enum(["deleted", "shutdown", "error"]),
}).strict());
const userMessage = event("user.message", z.object({
  messageId: id,
  text: z.string().max(100_000),
  attachmentIds: z.array(id).max(4),
}).strict());
const assistantDelta = event("assistant.delta", z.object({
  messageId: id,
  delta: z.string(),
}).strict());
const assistantCompleted = event("assistant.completed", z.object({
  messageId: id,
}).strict());
const skillRequested = event("skill.requested", z.object({
  name: z.string().trim().min(1).max(64),
  source: z.string().trim().min(1).max(200),
}).strict());
const skillLoaded = event("skill.loaded", z.object({
  name: z.string().trim().min(1).max(64),
  source: z.string().trim().min(1).max(200),
}).strict());
const toolRequested = event("tool.requested", z.object({
  toolCallId: id,
  toolName: z.string().trim().min(1).max(100),
  arguments: safeUnknown,
}).strict());
const toolCompleted = event("tool.completed", z.object({
  toolCallId: id,
  toolName: z.string().trim().min(1).max(100),
  result: safeUnknown,
}).strict());
const toolFailed = event("tool.failed", z.object({
  toolCallId: id,
  toolName: z.string().trim().min(1).max(100),
  error: z.string().trim().min(1).max(2000),
}).strict());
const approvalRequired = event("approval.required", z.object({
  approvalId: id,
  toolCallId: id,
  toolName: z.string().trim().min(1).max(100),
  arguments: safeUnknown,
  riskLevel: z.enum(["low", "medium", "high"]),
  expiresAt: z.string().datetime({ offset: true }),
}).strict());
const approvalResolved = event("approval.resolved", z.object({
  approvalId: id,
  status: z.enum(["approved", "denied", "expired", "cancelled"]),
}).strict());
const attachmentUploaded = event("attachment.uploaded", z.object({
  attachmentId: id,
  name: z.string().trim().min(1).max(255),
  mediaType: z.enum(["image/jpeg", "image/png", "image/webp"]),
  byteSize: z.int().positive(),
  width: z.int().positive(),
  height: z.int().positive(),
}).strict());
const attachmentRejected = event("attachment.rejected", z.object({
  name: z.string().trim().min(1).max(255),
  reason: z.string().trim().min(1).max(200),
}).strict());
const modelConfigured = event("model.configured", z.object({
  providerId: z.string().trim().min(1).max(80),
  modelId: z.string().trim().min(1).max(200),
  supportsImages: z.boolean(),
  source: z.enum(["session", "environment"]),
}).strict());
const modelValidationFailed = event("model.validation_failed", z.object({
  code: z.string().trim().min(1).max(100),
  message: z.string().trim().min(1).max(2000),
}).strict());
const runStarted = event("run.started", z.object({
  modelId: z.string().trim().min(1).max(200),
}).strict());
const runCancelled = event("run.cancelled", z.object({
  reason: z.enum(["user", "timeout", "session_deleted", "shutdown"]),
}).strict());
const runCompleted = event("run.completed", z.object({
  messageId: id,
}).strict());
const runFailed = event("run.failed", z.object({
  code: z.string().trim().min(1).max(100),
  message: z.string().trim().min(1).max(2000),
}).strict());

export const HarnessEventSchema = z.discriminatedUnion("type", [
  sessionStarted,
  sessionEnded,
  userMessage,
  assistantDelta,
  assistantCompleted,
  skillRequested,
  skillLoaded,
  toolRequested,
  toolCompleted,
  toolFailed,
  approvalRequired,
  approvalResolved,
  attachmentUploaded,
  attachmentRejected,
  modelConfigured,
  modelValidationFailed,
  runStarted,
  runCancelled,
  runCompleted,
  runFailed,
]);

export type HarnessEvent = z.infer<typeof HarnessEventSchema>;
export type HarnessEventType = HarnessEvent["type"];
export type HarnessEventOf<T extends HarnessEventType> = Extract<HarnessEvent, { type: T }>;
export type EventInput<T extends HarnessEventType> = Omit<
  HarnessEventOf<T>,
  "eventId" | "sessionId" | "sequence" | "timestamp"
>;