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

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

import {
  createRawContentExpectationsSchema,
  createRawKeyedExpectationsSchema,
  createRawValueExpectationSchema,
  statusCodePatternSchema,
  stringMapSchema,
} from "../../schema/fragments";

function createLlmOptionsSchema() {
  return Type.Object(
    {
      frequencyPenalty: Type.Optional(Type.Number()),
      maxOutputTokens: Type.Optional(Type.Integer({ minimum: 1 })),
      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 = {
  config: Type.Object(
    {
      authToken: Type.Optional(Type.String()),
      headers: Type.Optional(stringMapSchema),
      ignoreSSL: Type.Optional(Type.Boolean()),
      key: Type.Optional(Type.String()),
      mode: Type.Optional(Type.Union([Type.Literal("http"), Type.Literal("stream")])),
      model: Type.String({ minLength: 1 }),
      options: Type.Optional(createLlmOptionsSchema()),
      prompt: Type.String({ minLength: 1 }),
      provider: Type.Union([Type.Literal("openai"), Type.Literal("openai-responses"), Type.Literal("anthropic")]),
      providerOptions: Type.Optional(Type.Record(Type.String(), Type.Object({}, { additionalProperties: true }))),
      url: Type.String({ minLength: 1 }),
    },
    { additionalProperties: false },
  ),
  defaults: Type.Object(
    {
      headers: Type.Optional(stringMapSchema),
      ignoreSSL: Type.Optional(Type.Boolean()),
      mode: Type.Optional(Type.Union([Type.Literal("http"), Type.Literal("stream")])),
      options: Type.Optional(createLlmOptionsSchema()),
      providerOptions: Type.Optional(Type.Record(Type.String(), Type.Object({}, { additionalProperties: true }))),
    },
    { additionalProperties: false },
  ),
  expect: Type.Object(
    {
      durationMs: Type.Optional(createRawValueExpectationSchema()),
      finishReason: Type.Optional(createRawValueExpectationSchema()),
      headers: Type.Optional(createRawKeyedExpectationsSchema()),
      output: Type.Optional(createRawContentExpectationsSchema()),
      rawFinishReason: Type.Optional(createRawValueExpectationSchema()),
      status: Type.Optional(Type.Array(statusCodePatternSchema)),
      stream: Type.Optional(
        Type.Object(
          {
            completed: Type.Optional(Type.Boolean()),
            firstTokenMs: Type.Optional(createRawValueExpectationSchema()),
          },
          { additionalProperties: false },
        ),
      ),
      usage: Type.Optional(
        Type.Object(
          {
            inputTokens: Type.Optional(createRawValueExpectationSchema()),
            outputTokens: Type.Optional(createRawValueExpectationSchema()),
            totalTokens: Type.Optional(createRawValueExpectationSchema()),
          },
          { additionalProperties: false },
        ),
      ),
    },
    { additionalProperties: false },
  ),
};