import { z } from "zod";
import { StreamingChunkSchema } from "../proxy/schemas";

export const ProviderTestErrorTypeSchema = z.enum([
  "timeout",
  "network_unreachable",
  "connection_refused",
  "auth_failed",
  "rate_limited",
  "server_error",
  "invalid_response",
  "unknown",
]);

const ProviderTestErrorSchema = z.object({
  message: z.string(),
  type: ProviderTestErrorTypeSchema,
  details: z.string().optional(),
  hint: z.string().optional(),
});

const ProviderTestContentSchema = z.object({
  type: z.enum(["text", "thinking"]),
  text: z.string().optional(),
  thinking: z.string().optional(),
});

export const ProviderTestResultSchema = z.object({
  success: z.boolean(),
  error: ProviderTestErrorSchema.optional(),
  model: z.string().optional(),
  inputTokens: z.number().optional(),
  outputTokens: z.number().optional(),
  cacheCreationInputTokens: z.number().optional(),
  cacheReadInputTokens: z.number().optional(),
  latencyMs: z.number().optional(),
  content: z.array(ProviderTestContentSchema).optional(),
  rawResponse: z.string().optional(),
  streaming: z.boolean().optional(),
  streamingChunks: z
    .object({
      chunks: z.array(StreamingChunkSchema),
      truncated: z.boolean().default(false),
    })
    .optional(),
  requestHeaders: z.record(z.string(), z.string()).optional(),
});

export const ProviderTestStateSchema = z.union([
  ProviderTestResultSchema,
  z.object({ notConfigured: z.literal(true) }),
  z.object({ testing: z.literal(true) }),
]);

const ProviderFormatTestResultsSchema = z.object({
  nonStreaming: ProviderTestStateSchema,
  streaming: ProviderTestStateSchema,
});

const ModelTestResultsSchema = z.object({
  anthropic: ProviderFormatTestResultsSchema,
  openai: ProviderFormatTestResultsSchema,
});

export const ProviderTestResultsSchema = z.object({
  anthropic: ProviderFormatTestResultsSchema,
  openai: ProviderFormatTestResultsSchema,
  models: z.record(z.string(), ModelTestResultsSchema).optional(),
  testedAt: z.string().optional(),
});

export type ProviderTestErrorType = z.infer<typeof ProviderTestErrorTypeSchema>;
export type ProviderTestResult = z.infer<typeof ProviderTestResultSchema>;
export type ProviderTestState = z.infer<typeof ProviderTestStateSchema>;
export type ProviderTestResults = z.infer<typeof ProviderTestResultsSchema>;
export type ModelTestResults = z.infer<typeof ModelTestResultsSchema>;

export function createPendingProviderTestResults(): ProviderTestResults {
  return {
    anthropic: {
      nonStreaming: { testing: true },
      streaming: { testing: true },
    },
    openai: {
      nonStreaming: { testing: true },
      streaming: { testing: true },
    },
  };
}

export function createFailedProviderTestResults(
  message: string,
  type: ProviderTestErrorType,
): ProviderTestResults {
  const createFormatResult = () => ({
    nonStreaming: {
      success: false as const,
      error: { message, type },
    },
    streaming: { notConfigured: true as const },
  });

  return {
    anthropic: createFormatResult(),
    openai: createFormatResult(),
  };
}
