sandbox.ts

import { z } from "zod";

import {
  DEFAULT_MEMORY_MIB,
  DEFAULT_TEMPLATE_NAME,
  DEFAULT_TIMEOUT,
  DEFAULT_VCPU,
} from "@/config/constants.js";

/** Supported outbound network policy modes for a sandbox. */
export const NetworkPolicyMode = {
  ALLOW_ALL: "allow-all",
  DENY_ALL: "deny-all",
  CUSTOM: "custom",
} as const;

export const SandboxState = {
  STARTING: "starting",
  RUNNING: "running",
  SNAPSHOTTING: "snapshotting",
  PAUSED: "paused",
  UNPAUSING: "unpausing",
  DELETING: "deleting",
  DELETED: "deleted",
} as const;

export const networkPolicyModeSchema = z.enum([
  NetworkPolicyMode.ALLOW_ALL,
  NetworkPolicyMode.DENY_ALL,
  NetworkPolicyMode.CUSTOM,
]);
export type NetworkPolicyMode = z.infer<typeof networkPolicyModeSchema>;

function isValidDomainPattern(value: string): boolean {
  const trimmed = value.trim();
  if (!trimmed) {
    return false;
  }

  const host = trimmed.startsWith("*.") ? trimmed.slice(2) : trimmed;
  if (!host || host.startsWith(".") || host.endsWith(".")) {
    return false;
  }
  if (/^\d+\.\d+\.\d+\.\d+$/.test(host) || host.includes(":")) {
    return false;
  }

  const labels = host.split(".");
  if (labels.length < 2) {
    return false;
  }
  return labels.every(
    (label) =>
      label.length > 0 &&
      !label.startsWith("-") &&
      !label.endsWith("-") &&
      /^[A-Za-z0-9-]+$/.test(label),
  );
}

function isValidCidr(value: string): boolean {
  if (value.indexOf("/") !== value.lastIndexOf("/")) {
    return false;
  }
  const [address, prefix] = value.split("/");
  if (!address || prefix === undefined || !/^\d+$/.test(prefix)) {
    return false;
  }
  const octets = address.split(".");
  if (octets.length !== 4) {
    return false;
  }
  if (!octets.every((octet) => /^\d+$/.test(octet) && Number(octet) >= 0 && Number(octet) <= 255)) {
    return false;
  }
  const prefixNumber = Number(prefix);
  return prefixNumber >= 0 && prefixNumber <= 32;
}

const domainPatternSchema = z.string().refine(isValidDomainPattern, {
  message: "domain must be a valid domain pattern",
});

const cidrSchema = z.string().refine(isValidCidr, {
  message: "CIDR must be a valid IPv4 CIDR block",
});

export const networkPolicySchema = z.object({
  mode: networkPolicyModeSchema,
  allowedDomains: z.array(domainPatternSchema).max(50).optional(),
  allowedCidrs: z.array(cidrSchema).max(10).optional(),
  transforms: z
    .array(
      z.object({
        domain: domainPatternSchema,
        injectHeaders: z.record(z.string(), z.string()).optional(),
        stripHeaders: z.array(z.string()).optional(),
      }),
    )
    .max(20)
    .optional(),
});
/** Network policy configuration applied to sandbox egress traffic. */
export type NetworkPolicy = z.infer<typeof networkPolicySchema>;

const mountPathSchema = z
  .string()
  .startsWith("/", "mountPath must be an absolute path")
  .min(2, "mountPath must be an absolute path");

const mountPrefixSchema = z
  .string()
  .refine((value) => value.length === 0 || (!value.startsWith("/") && value.endsWith("/") && !value.includes("..")), {
    message: "prefix must be relative, must not contain '..', and must end with '/'",
  });

export const objectStorageMountSchema = z.object({
  type: z.literal("object-storage"),
  bucket: z.string().trim().min(1, "bucket must be a non-empty string"),
  mountPath: mountPathSchema,
  endpoint: z.string().url("endpoint must be a valid URL"),
  prefix: mountPrefixSchema.optional(),
  readOnly: z.boolean().optional(),
  accessKeyId: z.string().optional(),
  secretAccessKey: z.string().optional(),
});
export type ObjectStorageMount = z.infer<typeof objectStorageMountSchema>;

export const objectStorageMountUpdateSchema = z
  .object({
    bucket: z.string().trim().min(1, "bucket must be a non-empty string").optional(),
    mountPath: mountPathSchema.optional(),
    endpoint: z.string().url("endpoint must be a valid URL").optional(),
    prefix: mountPrefixSchema.optional(),
    readOnly: z.boolean().optional(),
    accessKeyId: z.string().optional(),
    secretAccessKey: z.string().optional(),
  })
  .refine((value) => Object.keys(value).length > 0, {
    message: "mount update must include at least one field",
  });
export type ObjectStorageMountUpdate = z.infer<typeof objectStorageMountUpdateSchema>;

export const objectStorageMountSummarySchema = z.object({
  id: z.string(),
  type: z.literal("object-storage"),
  bucket: z.string(),
  mountPath: z.string(),
  prefix: z.string().optional(),
  readOnly: z.boolean().optional(),
});
export type ObjectStorageMountSummary = z.infer<typeof objectStorageMountSummarySchema>;

export const sandboxStateSchema = z.enum([
  SandboxState.STARTING,
  SandboxState.RUNNING,
  SandboxState.SNAPSHOTTING,
  SandboxState.PAUSED,
  SandboxState.UNPAUSING,
  SandboxState.DELETING,
  SandboxState.DELETED,
]);
/** Lifecycle states reported for a sandbox. */
export type SandboxState = z.infer<typeof sandboxStateSchema>;

export const sandboxDataSchema = z
  .object({
    id: z.string(),
    templateId: z.string(),
    templateName: z.string().optional(),
    state: sandboxStateSchema,
    vcpu: z.number(),
    memory: z.number(),
    disk: z.number(),
    timeout: z.number().optional(),
    autoPause: z.boolean().optional(),
    envVars: z.record(z.string(), z.string()).optional(),
    networkPolicy: networkPolicySchema.optional(),
    mounts: z.array(objectStorageMountSummarySchema).optional(),
    createdAt: z.string(),
    updatedAt: z.string().optional(),
  })
  .catchall(z.unknown());
/** Sandbox resource returned by the control plane API. */
export type SandboxData = z.infer<typeof sandboxDataSchema>;

export const createSnapshotParamsSchema = z.object({
  name: z.string().trim().min(1).max(64).optional(),
  killSandboxAfter: z.boolean().optional(),
});
/** Parameters accepted when creating a snapshot from a sandbox. */
export type CreateSnapshotParams = z.infer<typeof createSnapshotParamsSchema>;

export const sandboxListItemSchema = z
  .object({
    id: z.string(),
    templateId: z.string(),
    state: sandboxStateSchema,
    launchTime: z.string().optional(),
    stateChangeTime: z.string().optional(),
    timeoutAt: z.number().optional(),
    createdAt: z.string(),
  })
  .catchall(z.unknown());
/** Sandbox summary returned by the list sandboxes API. */
export type SandboxListItem = z.infer<typeof sandboxListItemSchema>;

export const listSandboxesResponseSchema = z
  .object({
    items: z.array(sandboxListItemSchema),
    totalItems: z.number().int().nonnegative(),
  })
  .catchall(z.unknown());
/** Paginated sandbox list response. */
export type ListSandboxesResponse = z.infer<typeof listSandboxesResponseSchema>;

export const createPresignedUrlParamsSchema = z.object({
  port: z.number().int().min(1).max(65535),
  expiresIn: z.number().int().min(1).optional(),
});
export type CreatePresignedUrlParams = z.infer<typeof createPresignedUrlParamsSchema>;

export const presignedUrlSchema = z
  .object({
    id: z.string(),
    token: z.string(),
    url: z.string().url(),
    sandboxId: z.string(),
    port: z.number().int().min(1).max(65535),
    expiresAt: z.string(),
    createdAt: z.string(),
  })
  .catchall(z.unknown());
export type PresignedUrl = z.infer<typeof presignedUrlSchema>;

export const listSandboxesParamsSchema = z
  .object({
    state: z
      .enum([
        SandboxState.STARTING,
        SandboxState.SNAPSHOTTING,
        SandboxState.RUNNING,
        SandboxState.PAUSED,
        SandboxState.UNPAUSING,
        SandboxState.DELETING,
      ])
      .optional(),
    sort: z.enum(["created_at", "state"]).optional(),
    orderBy: z.enum(["asc", "desc"]).optional(),
    page: z.number().int().min(1).optional(),
    pageSize: z.number().int().min(1).max(100).optional(),
  })
  .passthrough();
/** Parameters accepted when listing sandboxes. */
export type ListSandboxesParams = z.infer<typeof listSandboxesParamsSchema>;

export const createSandboxParamsSchema = z
  .object({
    templateName: z.string().optional(),
    vcpu: z.number().int().positive().optional(),
    memory: z.number().int().positive().optional(),
    timeout: z.number().int().positive().optional(),
    autoPause: z.boolean().optional(),
    otelExport: z.boolean().optional(),
    telemetry: z.boolean().optional(),
    envVars: z.record(z.string(), z.string()).optional(),
    networkPolicy: networkPolicySchema.optional(),
    mounts: z.array(objectStorageMountSchema).max(8).optional(),
  })
  .refine((value) => hasUniqueMountPaths(value.mounts), {
    path: ["mounts"],
    message: "mounts must use unique mountPath values",
  });
/** Parameters accepted when creating a sandbox. */
export type CreateSandboxParams = z.infer<typeof createSandboxParamsSchema>;

export const createSandboxRuntimeParamsSchema = z
  .object(
    {
      templateName: z.preprocess(
        (value) => value ?? DEFAULT_TEMPLATE_NAME,
        z
          .string({ invalid_type_error: "templateName must be a string" })
          .trim()
          .min(1, "templateName must be 1-64 characters")
          .max(64, "templateName must be 1-64 characters"),
      ),
      vcpu: z.preprocess(
        (value) => value ?? DEFAULT_VCPU,
        z
          .number({ invalid_type_error: "vcpu must be between 1 and 8" })
          .int("vcpu must be between 1 and 8")
          .min(1, "vcpu must be between 1 and 8")
          .max(8, "vcpu must be between 1 and 8"),
      ),
      memory: z.preprocess(
        (value) => value ?? DEFAULT_MEMORY_MIB,
        z
          .number({ invalid_type_error: "memory must be even and between 512 and 8192" })
          .int("memory must be even and between 512 and 8192")
          .min(512, "memory must be even and between 512 and 8192")
          .max(8192, "memory must be even and between 512 and 8192")
          .refine((value) => value % 2 === 0, {
            message: "memory must be even and between 512 and 8192",
          }),
      ),
      timeout: z.preprocess(
        (value) => value ?? DEFAULT_TIMEOUT,
        z
          .number({ invalid_type_error: "timeout must be between 1 and 28800" })
          .int("timeout must be between 1 and 28800")
          .min(1, "timeout must be between 1 and 28800")
          .max(28800, "timeout must be between 1 and 28800"),
      ),
      autoPause: z.boolean().optional(),
      otelExport: z.boolean().optional(),
      telemetry: z.boolean().optional(),
      envVars: z.record(z.string(), z.string()).optional(),
      networkPolicy: networkPolicySchema.optional(),
      mounts: z.array(objectStorageMountSchema).max(8, "mounts must contain at most 8 entries").optional(),
    },
    { invalid_type_error: "params must be an object" },
  )
  .passthrough()
  .refine((value) => hasUniqueMountPaths(value.mounts), {
    path: ["mounts"],
    message: "mounts must use unique mountPath values",
  });

function hasUniqueMountPaths(mounts: ObjectStorageMount[] | undefined): boolean {
  if (mounts == null) {
    return true;
  }
  return new Set(mounts.map((mount) => mount.mountPath)).size === mounts.length;
}

type NetworkPolicyWire = {
  mode: NetworkPolicyMode;
  allow_domains?: string[];
  allow_cidrs?: string[];
  transforms?: Array<{
    domain: string;
    inject_headers?: Record<string, string>;
    strip_headers?: string[];
  }>;
};

type ObjectStorageMountWire = {
  type: "object-storage";
  bucket: string;
  mount_path: string;
  endpoint: string;
  prefix?: string;
  read_only?: boolean;
  access_key_id?: string;
  secret_access_key?: string;
};

/**
 * Converts SDK network policy input into the wire format expected by the API.
 *
 * @param policy SDK network policy configuration.
 * @returns The API wire-format network policy.
 */
export function toNetworkPolicyWire(
  policy: NetworkPolicy | undefined,
): NetworkPolicyWire | undefined {
  if (policy == null) return undefined;

  return {
    mode: policy.mode,
    allow_domains: policy.allowedDomains,
    allow_cidrs: policy.allowedCidrs,
    transforms: policy.transforms?.map((transform) => ({
      domain: transform.domain,
      inject_headers: transform.injectHeaders,
      strip_headers: transform.stripHeaders,
    })),
  };
}

export function toObjectStorageMountsWire(
  mounts: ObjectStorageMount[] | undefined,
): ObjectStorageMountWire[] | undefined {
  if (mounts == null) return undefined;

  return mounts.map((mount) => ({
    type: mount.type,
    bucket: mount.bucket,
    mount_path: mount.mountPath,
    endpoint: mount.endpoint,
    prefix: mount.prefix,
    read_only: mount.readOnly,
    access_key_id: mount.accessKeyId,
    secret_access_key: mount.secretAccessKey,
  }));
}

type ObjectStorageMountUpdateWire = {
  bucket?: string;
  mount_path?: string;
  endpoint?: string;
  prefix?: string;
  read_only?: boolean;
  access_key_id?: string;
  secret_access_key?: string;
};

export function toObjectStorageMountUpdateWire(
  mount: ObjectStorageMountUpdate,
): ObjectStorageMountUpdateWire {
  return {
    bucket: mount.bucket,
    mount_path: mount.mountPath,
    endpoint: mount.endpoint,
    prefix: mount.prefix,
    read_only: mount.readOnly,
    access_key_id: mount.accessKeyId,
    secret_access_key: mount.secretAccessKey,
  };
}