import { z } from "zod";

export const providerTypeSchema = z.enum([
  "anthropic",
  "openai",
  "google",
  "openai_compatible",
  "azure_openai",
]);

export type ProviderType = z.infer<typeof providerTypeSchema>;

export const modelTierSchema = z.enum(["premium", "fast", "local", "auto"]);

export type ModelTier = z.infer<typeof modelTierSchema>;

export const modelProviderSchema = z.object({
  id: z.string().uuid(),
  workspaceId: z.string().uuid(),
  alias: z.string().min(1),
  displayName: z.string().min(1),
  providerType: providerTypeSchema,
  modelId: z.string().min(1),
  baseUrl: z.string().url().nullable().optional(),
  apiKeyEnvVar: z.string().min(1).nullable().optional(),
  extraHeaders: z.record(z.string()).nullable().optional(),
  defaultParams: z
    .object({
      temperature: z.number().optional(),
      maxTokens: z.number().optional(),
    })
    .nullable()
    .optional(),
  metadata: z.record(z.unknown()).nullable().optional(),
  enabled: z.boolean().default(true),
});

export type ModelProvider = z.infer<typeof modelProviderSchema>;

export const roleModelRouteSchema = z.object({
  id: z.string().uuid(),
  workspaceId: z.string().uuid(),
  roleId: z.string().min(1),
  providerAliasId: z.string().uuid(),
  enabled: z.boolean().default(true),
  sortOrder: z.number().int().optional(),
});

export type RoleModelRoute = z.infer<typeof roleModelRouteSchema>;

export const resolvedModelSchema = z.object({
  alias: z.string().min(1),
  providerType: providerTypeSchema,
  modelId: z.string().min(1),
  baseUrl: z.string().url().optional(),
  apiKey: z.string(),
  extraHeaders: z.record(z.string()).optional(),
  defaultParams: z
    .object({
      temperature: z.number().optional(),
      maxTokens: z.number().optional(),
    })
    .optional(),
  tier: modelTierSchema.optional(),
});

export type ResolvedModel = z.infer<typeof resolvedModelSchema>;

export const convoyModelPolicySchema = z.object({
  roles: z.record(z.string()).optional(),
  escalate_to: z.string().min(1).optional(),
});

export type ConvoyModelPolicy = z.infer<typeof convoyModelPolicySchema>;