import { z } from "zod";
import { TELEGRAM_IMAGE_MAX_BASE64_CHARS } from "../../connectors/telegram-image-contract.js";
import { capturedProposalsSchema, proposalCapabilitiesSchema } from "../proposals.js";
import { conversationMessagesSchema } from "../conversation-history.js";

export const SANDBOX_PROTOCOL_VERSION = 7;
export const MAX_RUN_PACKET_BYTES = 12 * 1024 * 1024;
export const MAX_RESULT_FRAME_BYTES = 2 * 1024 * 1024;
export const MAX_TRACE_ITEMS = 256;
export const MAX_TRACE_BYTES = 128 * 1024;

const imageSchema = z.object({
  type: z.literal("image"),
  data: z.string().max(TELEGRAM_IMAGE_MAX_BASE64_CHARS),
  mimeType: z.string().min(1).max(100),
}).strict();

export const sandboxRunPacketSchema = z.object({
  version: z.literal(SANDBOX_PROTOCOL_VERSION),
  runId: z.string().min(1).max(200),
  systemPrompt: z.string().max(256 * 1024),
  prompt: z.string().max(1_000_000),
  messages: conversationMessagesSchema.optional(),
  images: z.array(imageSchema).max(8),
  proposals: proposalCapabilitiesSchema.optional(),
  model: z.object({
    provider: z.enum(["tinker", "openai-compatible"]),
    id: z.string().min(1).max(500),
    reasoning: z.boolean(),
    thinkingLevel: z.enum(["off", "minimal", "low", "medium", "high", "xhigh", "max"]).default("off"),
    acceptsImages: z.boolean(),
    jsonObjectResponseFormat: z.boolean(),
    jsonSchemaResponseFormat: z.record(z.string(), z.unknown()).optional(),
    contextWindow: z.number().int().positive().max(2_000_000),
    maxTokens: z.number().int().positive().max(32_000),
  }).strict(),
  broker: z.object({
    socketPath: z.literal("/broker/provider.sock"),
    capability: z.string().regex(/^[a-f0-9]{64}$/),
    origin: z.literal("http://provider.invalid"),
    route: z.literal("/v1/chat/completions"),
  }).strict(),
}).strict();

const traceSchema = z.object({
  kind: z.string().min(1).max(200),
  data: z.unknown(),
}).strict();

export const sandboxResultSchema = z.discriminatedUnion("status", [
  z.object({
    version: z.literal(SANDBOX_PROTOCOL_VERSION),
    runId: z.string().min(1).max(200),
    status: z.literal("completed"),
    finalMessage: z.unknown(),
    proposals: capturedProposalsSchema,
    traces: z.array(traceSchema).max(MAX_TRACE_ITEMS),
    observedRevision: z.string().min(1).max(500).optional(),
  }).strict(),
  z.object({
    version: z.literal(SANDBOX_PROTOCOL_VERSION),
    runId: z.string().min(1).max(200),
    status: z.literal("failed"),
    code: z.enum([
      "invalid-packet",
      "provider-rejected",
      "provider-client-error",
      "provider-transient-error",
      "provider-timeout",
      "model-failed",
      "protocol-failed",
      "output-limit",
    ]),
    message: z.string().min(1).max(500),
    reason: z.enum([
      "assistant-missing",
      "assistant-error-stop",
      "proposal-tool-error",
      "proposal-result-invalid",
      "worker-exception",
    ]).optional(),
    traces: z.array(traceSchema).max(MAX_TRACE_ITEMS).optional(),
    proposalFailureCode: z.enum([
      "arguments-invalid",
      "evidence-outside-snapshot",
      "target-outside-snapshot",
      "proposal-call-limit",
      "duplicate-proposal-kind",
      "proposal-tool-unknown",
    ]).optional(),
  }).strict(),
]);

export type SandboxRunPacket = z.infer<typeof sandboxRunPacketSchema>;
export type SandboxResult = z.infer<typeof sandboxResultSchema>;

export interface BrokerRequest {
  version: 1;
  capability: string;
  runId: string;
  origin: string;
  path: string;
  method: "POST";
  model: string;
  headers: Record<string, string>;
  body: string;
}

export interface BrokerResponse {
  version: 1;
  status: "completed" | "rejected";
  httpStatus?: number;
  headers?: Record<string, string>;
  body?: string;
  code?: string;
  message?: string;
}

export function encodeFrame(value: unknown, maxBytes: number): Buffer {
  const payload = Buffer.from(JSON.stringify(value), "utf8");
  if (payload.length > maxBytes) throw new Error(`Frame exceeds ${maxBytes} bytes`);
  const header = Buffer.allocUnsafe(4);
  header.writeUInt32BE(payload.length, 0);
  return Buffer.concat([header, payload]);
}

export function decodeSingleFrame(buffer: Buffer, maxBytes: number): unknown {
  if (buffer.length < 4) throw new Error("Truncated frame header");
  const size = buffer.readUInt32BE(0);
  if (size > maxBytes) throw new Error(`Frame exceeds ${maxBytes} bytes`);
  if (buffer.length !== size + 4) {
    throw new Error(buffer.length < size + 4 ? "Truncated frame payload" : "Duplicate or trailing frame data");
  }
  return JSON.parse(buffer.subarray(4).toString("utf8")) as unknown;
}