import { z } from "zod";

import { Backend, Capability, RiskLevel } from "./enums.js";
import { hash } from "./hash.js";
import { DbIntVersion } from "./version.js";

export type JsonValue =
  | null
  | boolean
  | number
  | string
  | JsonValue[]
  | { [key: string]: JsonValue };

export type JsonObject = { [key: string]: JsonValue };

export const JsonObject: z.ZodType<JsonObject, z.ZodTypeDef, unknown> = z.lazy(() =>
  z
    .custom<Record<string, unknown>>(isPlainJsonRecordInput, {
      message: "expected plain JSON object",
    })
    .superRefine((value, context) => {
      for (const key of Reflect.ownKeys(value)) {
        if (typeof key !== "string" || !Object.prototype.propertyIsEnumerable.call(value, key)) {
          context.addIssue({
            code: z.ZodIssueCode.custom,
            message: "expected plain JSON object",
          });
          continue;
        }

        if (!isSafeJsonObjectKey(key)) {
          context.addIssue({
            code: z.ZodIssueCode.custom,
            message: "reserved object key",
            path: [key],
          });
        }
      }
    })
    .pipe(z.record(z.string(), JsonValue)),
);

export const JsonValue: z.ZodType<JsonValue, z.ZodTypeDef, unknown> = z.lazy(() =>
  z.union([z.null(), z.boolean(), z.number().finite(), z.string(), JsonArray, JsonObject]),
);

export const JsonArray: z.ZodType<JsonValue[], z.ZodTypeDef, unknown> = z.lazy(() =>
  z
    .custom<unknown[]>(Array.isArray, { message: "expected JSON array" })
    .superRefine((value, context) => {
      for (const key of Reflect.ownKeys(value)) {
        if (key === "length") {
          continue;
        }

        if (typeof key !== "string" || !Object.prototype.propertyIsEnumerable.call(value, key)) {
          context.addIssue({
            code: z.ZodIssueCode.custom,
            message: "expected JSON array",
          });
          continue;
        }

        const index = Number(key);
        if (
          !Number.isInteger(index) ||
          index < 0 ||
          index >= value.length ||
          String(index) !== key
        ) {
          context.addIssue({
            code: z.ZodIssueCode.custom,
            message: "expected JSON array",
          });
        }
      }
    })
    .pipe(z.array(JsonValue)),
);

export const Persona = z
  .object({
    name: z.string().min(1),
    version: DbIntVersion,
    backend: Backend,
    capabilities: z.array(Capability),
    maxRiskLevel: RiskLevel,
    allowedRoles: z.array(z.string().min(1)).optional(),
    promptConfig: z
      .object({
        systemPrompt: z.string().optional(),
        instructionsPrelude: z.string().optional(),
      })
      .strict()
      .default({})
      .transform((value) => {
        const promptConfig: { systemPrompt?: string; instructionsPrelude?: string } = {};
        if (value.systemPrompt !== undefined) {
          promptConfig.systemPrompt = value.systemPrompt;
        }
        if (value.instructionsPrelude !== undefined) {
          promptConfig.instructionsPrelude = value.instructionsPrelude;
        }

        return promptConfig;
      }),
    modelConfig: JsonObject.default({}),
  })
  .strict();

export type Persona = z.infer<typeof Persona>;

export function personaHash(persona: Persona | undefined): string {
  if (!persona) {
    throw new TypeError("persona is required");
  }

  const hashSubject = {
    name: persona.name,
    version: persona.version,
    capabilities: persona.capabilities,
    backend: persona.backend,
    maxRiskLevel: persona.maxRiskLevel,
    promptConfig: persona.promptConfig,
    modelConfig: persona.modelConfig,
  };

  return hash(
    persona.allowedRoles === undefined
      ? hashSubject
      : { ...hashSubject, allowedRoles: persona.allowedRoles },
  );
}

function isSafeJsonObjectKey(key: string) {
  return key !== "__proto__" && key !== "constructor" && key !== "prototype";
}

function isPlainJsonRecordInput(value: unknown) {
  if (value === null || typeof value !== "object" || Array.isArray(value)) {
    return false;
  }

  const prototype = Object.getPrototypeOf(value);
  return prototype === Object.prototype || prototype === null;
}