import { z } from "zod";
import { ArtifactDescriptor } from "./artifacts.js";
import { RunEvent } from "./events.js";
import { RunManifest, StageKey } from "./run.js";

export const WORKER_PROTOCOL_VERSION = 1;

export const WorkerStatus = z.enum(["online", "offline", "draining", "quarantined", "disabled"]);
export type WorkerStatus = z.infer<typeof WorkerStatus>;

export const NodeMetadata = z.object({
  machineId: z.string(),
  hostname: z.string(),
  platform: z.string(),
  arch: z.string(),
  cpuCores: z.number().int(),
  memoryBytes: z.number().int(),
  osVersion: z.string(),
  workerVersion: z.string(),
  protocolVersion: z.number().int(),
  sdkVersion: z.string().optional(),
  tools: z.record(z.string(), z.string()).default({}),
  profiles: z.array(z.string()).default(["standard"]),
  user: z.string(),
  sandboxAvailable: z.boolean(),
});
export type NodeMetadata = z.infer<typeof NodeMetadata>;

export const EnrollRequest = z.object({
  enrollmentSecret: z.string().min(16),
  name: z.string().min(2).max(64),
  lanIp: z.string().optional(),
  node: NodeMetadata,
});
export const EnrollResponse = z.object({
  workerId: z.string().uuid(),
  token: z.string(),
  protocolVersion: z.number(),
});

export const HeartbeatRequest = z.object({
  status: WorkerStatus.default("online"),
  capacity: z.number().int().min(0).max(32),
  running: z.array(z.string().uuid()),
  load: z
    .object({ load1: z.number(), memFreeBytes: z.number(), diskFreeBytes: z.number() })
    .partial()
    .default({}),
  node: NodeMetadata.partial().optional(),
  leases: z.array(z.object({ runId: z.string().uuid(), fencingToken: z.number().int() })).default([]),
});
export const HeartbeatResponse = z.object({
  ok: z.literal(true),
  drain: z.boolean(),
  cancel: z.array(z.string().uuid()),
  revokedLeases: z.array(z.string().uuid()),
  leaseTtlSeconds: z.number().int(),
});

export const LeaseRequest = z.object({ profiles: z.array(z.string()).default(["standard"]) });
export const LeasedJob = z.object({
  runId: z.string().uuid(),
  attemptId: z.string().uuid(),
  attemptNo: z.number().int(),
  fencingToken: z.number().int(),
  leaseExpiresAt: z.string().datetime(),
  manifest: RunManifest,
  instructionText: z.string(),
  resumeFromStage: StageKey.nullable(),
  checkpoint: z.record(z.string(), z.unknown()).nullable(),
  cancelRequested: z.boolean(),
});
export const LeaseResponse = z.object({ job: LeasedJob.nullable() });

export const LeaseRenewRequest = z.object({ fencingToken: z.number().int() });
export const LeaseRenewResponse = z.object({
  ok: z.boolean(),
  cancelRequested: z.boolean(),
  leaseExpiresAt: z.string().datetime().optional(),
});

export const EventsPostRequest = z.object({
  fencingToken: z.number().int(),
  events: z.array(RunEvent).min(1).max(500),
});

export const StageReport = z.object({
  fencingToken: z.number().int(),
  stage: StageKey,
  status: z.enum(["running", "completed", "failed", "skipped"]),
  costUsd: z.number().nonnegative().default(0),
  turns: z.number().int().nonnegative().default(0),
  usage: z.record(z.string(), z.unknown()).optional(),
  error: z.string().max(4000).optional(),
  checkpoint: z.record(z.string(), z.unknown()).optional(),
});

export const ArtifactUploadInit = z.object({ fencingToken: z.number().int(), artifact: ArtifactDescriptor });

export const CompleteAttemptRequest = z.object({
  fencingToken: z.number().int(),
  outcome: z.enum(["completed", "failed", "cancelled", "budget_exceeded", "timeout"]),
  error: z.string().max(4000).optional(),
  summary: z.record(z.string(), z.unknown()).optional(),
  costUsd: z.number().nonnegative().default(0),
  resourceUsage: z
    .object({
      cpuSeconds: z.number().optional(),
      wallSeconds: z.number().optional(),
      diskBytes: z.number().optional(),
      peakRssBytes: z.number().optional(),
    })
    .optional(),
});