import { Type } from "@sinclair/typebox";

import type { CheckerSchemas } from "../types";

import {
  createAuthoringContentExpectationsSchema,
  createAuthoringFieldSchema,
  createAuthoringKeyedExpectationsSchema,
  createAuthoringStringMapSchema,
  createAuthoringValueExpectationSchema,
  createNormalizedContentExpectationsSchema,
  createNormalizedKeyedExpectationsSchema,
  createNormalizedValueExpectationSchema,
  statusCodePatternSchema,
  stringMapSchema,
} from "../../schema/fragments";

function createLlmOptionsSchema(kind: "authoring" | "normalized") {
  const maxOutputTokens = Type.Integer({ minimum: 1 });
  return Type.Object(
    {
      frequencyPenalty: Type.Optional(Type.Number()),
      maxOutputTokens: Type.Optional(
        kind === "authoring" ? createAuthoringFieldSchema(maxOutputTokens) : maxOutputTokens,
      ),
      presencePenalty: Type.Optional(Type.Number()),
      seed: Type.Optional(Type.Number()),
      stopSequences: Type.Optional(Type.Array(Type.String())),
      temperature: Type.Optional(Type.Number()),
      topK: Type.Optional(Type.Number()),
      topP: Type.Optional(Type.Number()),
    },
    { additionalProperties: false },
  );
}

export const llmCheckerSchemas: CheckerSchemas = {
  authoring: {
    config: createLlmConfigSchema("authoring"),
    expect: createLlmExpectSchema("authoring"),
  },
  normalized: {
    config: createLlmConfigSchema("normalized"),
    expect: createLlmExpectSchema("normalized"),
  },
};

function createLlmConfigSchema(kind: "authoring" | "normalized") {
  const bool = Type.Boolean();
  const mode = createLlmModeSchema();
  const provider = createLlmProviderSchema();
  return Type.Object(
    {
      authToken: Type.Optional(Type.String()),
      headers: Type.Optional(kind === "authoring" ? createAuthoringStringMapSchema() : stringMapSchema),
      ignoreSSL: Type.Optional(kind === "authoring" ? createAuthoringFieldSchema(bool) : bool),
      key: Type.Optional(Type.String()),
      mode: Type.Optional(kind === "authoring" ? createAuthoringFieldSchema(mode) : mode),
      model: Type.String({ minLength: 1 }),
      options: Type.Optional(createLlmOptionsSchema(kind)),
      prompt: Type.String({ minLength: 1 }),
      provider: kind === "authoring" ? createAuthoringFieldSchema(provider) : provider,
      providerOptions: Type.Optional(Type.Record(Type.String(), Type.Object({}, { additionalProperties: true }))),
      url: Type.String({ minLength: 1 }),
    },
    { additionalProperties: false },
  );
}

function createLlmExpectSchema(kind: "authoring" | "normalized") {
  const bool = Type.Boolean();
  const valueExpectation =
    kind === "authoring" ? createAuthoringValueExpectationSchema() : createNormalizedValueExpectationSchema();
  return Type.Object(
    {
      durationMs: Type.Optional(valueExpectation),
      finishReason: Type.Optional(valueExpectation),
      headers: Type.Optional(
        kind === "authoring" ? createAuthoringKeyedExpectationsSchema() : createNormalizedKeyedExpectationsSchema(),
      ),
      output: Type.Optional(
        kind === "authoring" ? createAuthoringContentExpectationsSchema() : createNormalizedContentExpectationsSchema(),
      ),
      rawFinishReason: Type.Optional(valueExpectation),
      status: Type.Optional(Type.Array(statusCodePatternSchema)),
      stream: Type.Optional(
        Type.Object(
          {
            completed: Type.Optional(kind === "authoring" ? createAuthoringFieldSchema(bool) : bool),
            firstTokenMs: Type.Optional(valueExpectation),
          },
          { additionalProperties: false },
        ),
      ),
      usage: Type.Optional(
        Type.Object(
          {
            inputTokens: Type.Optional(valueExpectation),
            outputTokens: Type.Optional(valueExpectation),
            totalTokens: Type.Optional(valueExpectation),
          },
          { additionalProperties: false },
        ),
      ),
    },
    { additionalProperties: false },
  );
}

function createLlmModeSchema() {
  return Type.Union([Type.Literal("http"), Type.Literal("stream")]);
}

function createLlmProviderSchema() {
  return Type.Union([Type.Literal("openai"), Type.Literal("openai-responses"), Type.Literal("anthropic")]);
}
