import { z } from "zod";
import { approvalRequestSchema, imageAttachmentSchema, llmApiSchema, reportArtifactSchema } from "./schemas.js";

const baseShape = {
  eventId: z.string(),
  sessionId: z.string(),
  runId: z.string().optional(),
  timestamp: z.string().datetime(),
  sequence: z.number().int().positive(),
};

const event = <T extends string, S extends z.ZodType>(type: T, payload: S) =>
  z.object({ ...baseShape, type: z.literal(type), payload }).strict();

const empty = z.object({}).strict();
const errorPayload = z.object({ code: z.string(), message: z.string() }).strict();
const toolBase = z.object({ toolCallId: z.string(), toolName: z.string(), arguments: z.unknown().optional() }).strict();

export const harnessEventSchema = z.discriminatedUnion("type", [
  event("session.started", z.object({ cwd: z.string() }).strict()),
  event("session.ended", z.object({ reason: z.enum(["deleted", "shutdown", "error"]) }).strict()),
  event("user.message", z.object({ messageId: z.string(), text: z.string(), attachmentIds: z.array(z.string()) }).strict()),
  event("assistant.delta", z.object({ messageId: z.string(), delta: z.string() }).strict()),
  event("assistant.completed", z.object({ messageId: z.string() }).strict()),
  event("skill.requested", z.object({ name: z.string(), filePath: z.string(), toolCallId: z.string().optional(), source: z.enum(["read", "command"]) }).strict()),
  event("skill.loaded", z.object({ name: z.string(), filePath: z.string(), toolCallId: z.string().optional(), source: z.enum(["read", "command"]) }).strict()),
  event("report.generated", z.object({ assistantMessageId: z.string(), artifact: reportArtifactSchema }).strict()),
  event("tool.requested", toolBase),
  event("tool.completed", z.object({ toolCallId: z.string(), toolName: z.string(), result: z.unknown().optional() }).strict()),
  event("tool.failed", z.object({ toolCallId: z.string(), toolName: z.string(), error: z.string() }).strict()),
  event("approval.required", approvalRequestSchema),
  event("approval.resolved", z.object({ approvalId: z.string(), toolCallId: z.string(), status: z.enum(["approved", "denied", "expired", "cancelled"]), resolvedAt: z.string().datetime() }).strict()),
  event("attachment.uploaded", imageAttachmentSchema.omit({ sessionId: true })),
  event("attachment.rejected", z.object({ name: z.string(), code: z.string(), message: z.string() }).strict()),
  event("model.configured", z.object({ providerId: z.string(), modelId: z.string(), api: llmApiSchema, supportsImages: z.boolean() }).strict()),
  event("model.validation_failed", errorPayload),
  event("run.started", z.object({ modelId: z.string() }).strict()),
  event("run.cancelled", z.object({ reason: z.enum(["user", "timeout", "session_deleted", "shutdown"]) }).strict()),
  event("run.completed", empty),
  event("run.failed", errorPayload),
]);

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

export const sseControlEventSchema = z.object({
  type: z.literal("stream.gap"),
  oldestAvailableSequence: z.number().int().positive(),
  requestedSequence: z.number().int().nonnegative(),
}).strict();