import { z } from "zod"
import { eq, and } from "drizzle-orm"
import { Database } from "./drizzle"
import { ModelTable } from "./schema/model.sql"
import { ModelProviderPreferenceTable } from "./schema/model-provider-preference.sql"
import { Identifier } from "./identifier"
import { fn } from "./util/fn"
import { Actor } from "./actor"
import { Resource } from "@opencode-ai/console-resource"

export namespace ZenData {
  const FormatSchema = z.enum(["anthropic", "google", "openai", "oa-compat"])
  const TrialSchema = z.object({
    provider: z.string(),
    limits: z.array(
      z.object({
        limit: z.number(),
        client: z.enum(["cli", "desktop"]).optional(),
      }),
    ),
  })
  const RateLimitSchema = z.object({
    period: z.enum(["day", "rolling"]),
    value: z.number().int(),
    checkHeader: z.string().optional(),
    fallbackValue: z.number().int().optional(),
  })
  export type Format = z.infer<typeof FormatSchema>
  export type Trial = z.infer<typeof TrialSchema>
  export type RateLimit = z.infer<typeof RateLimitSchema>

  const ModelCostSchema = z.object({
    input: z.number(),
    output: z.number(),
    cacheRead: z.number().optional(),
    cacheWrite5m: z.number().optional(),
    cacheWrite1h: z.number().optional(),
  })

  const ModelSchema = z.object({
    name: z.string(),
    cost: ModelCostSchema,
    cost200K: ModelCostSchema.optional(),
    allowAnonymous: z.boolean().optional(),
    byokProvider: z.enum(["openai", "anthropic", "google"]).optional(),
    stickyProvider: z.enum(["strict", "prefer"]).optional(),
    trial: TrialSchema.optional(),
    rateLimit: RateLimitSchema.optional(),
    fallbackProvider: z.string().optional(),
    providers: z.array(
      z.object({
        id: z.string(),
        model: z.string(),
        weight: z.number().optional(),
        disabled: z.boolean().optional(),
        storeModel: z.string().optional(),
        headerMappings: z.record(z.string(), z.string()).optional(),
      }),
    ),
  })

  const ProviderSchema = z.object({
    api: z.string(),
    apiKey: z.string(),
    format: FormatSchema,
    headerMappings: z.record(z.string(), z.string()).optional(),
  })

  const ModelsSchema = z.object({
    models: z.record(z.string(), z.union([ModelSchema, z.array(ModelSchema.extend({ formatFilter: FormatSchema }))])),
    providers: z.record(z.string(), ProviderSchema),
  })

  export const validate = fn(ModelsSchema, (input) => {
    return input
  })

  export const list = fn(z.void(), () => {
    const json = JSON.parse(
      Resource.ZEN_MODELS1.value +
        Resource.ZEN_MODELS2.value +
        Resource.ZEN_MODELS3.value +
        Resource.ZEN_MODELS4.value +
        Resource.ZEN_MODELS5.value +
        Resource.ZEN_MODELS6.value +
        Resource.ZEN_MODELS7.value +
        Resource.ZEN_MODELS8.value +
        Resource.ZEN_MODELS9.value +
        Resource.ZEN_MODELS10.value,
    )
    return ModelsSchema.parse(json)
  })
}

export const ProviderPreferenceSchema = z.object({
  order: z.array(z.string()).optional(),
  only: z.array(z.string()).optional(),
  ignore: z.array(z.string()).optional(),
  allow_fallbacks: z.boolean().optional(),
})
export type ProviderPreference = z.infer<typeof ProviderPreferenceSchema>

export namespace Model {
  export const enable = fn(z.object({ model: z.string() }), ({ model }) => {
    Actor.assertAdmin()
    return Database.use((db) =>
      db.delete(ModelTable).where(and(eq(ModelTable.workspaceID, Actor.workspace()), eq(ModelTable.model, model))),
    )
  })

  export const disable = fn(z.object({ model: z.string() }), ({ model }) => {
    Actor.assertAdmin()
    return Database.use((db) =>
      db
        .insert(ModelTable)
        .values({
          id: Identifier.create("model"),
          workspaceID: Actor.workspace(),
          model: model,
        })
        .onDuplicateKeyUpdate({
          set: {
            timeDeleted: null,
          },
        }),
    )
  })

  export const listDisabled = fn(z.void(), () => {
    return Database.use((db) =>
      db
        .select({ model: ModelTable.model })
        .from(ModelTable)
        .where(eq(ModelTable.workspaceID, Actor.workspace()))
        .then((rows) => rows.map((row) => row.model)),
    )
  })

  export const isDisabled = fn(
    z.object({
      model: z.string(),
    }),
    ({ model }) => {
      return Database.use(async (db) => {
        const result = await db
          .select()
          .from(ModelTable)
          .where(and(eq(ModelTable.workspaceID, Actor.workspace()), eq(ModelTable.model, model)))
          .limit(1)

        return result.length > 0
      })
    },
  )

  export const setProviderPreference = fn(
    z.object({
      model: z.string(),
      preference: ProviderPreferenceSchema,
    }),
    ({ model, preference }) => {
      Actor.assertAdmin()
      return Database.use((db) =>
        db
          .insert(ModelProviderPreferenceTable)
          .values({
            id: Identifier.create("model_provider_pref"),
            workspaceID: Actor.workspace(),
            model: model,
            provider_preference: JSON.stringify(preference),
          })
          .onDuplicateKeyUpdate({
            set: {
              provider_preference: JSON.stringify(preference),
              timeDeleted: null,
            },
          }),
      )
    },
  )

  export const getProviderPreference = fn(
    z.object({ model: z.string() }),
    ({ model }) => {
      return Database.use(async (db) => {
        const result = await db
          .select({ preference: ModelProviderPreferenceTable.provider_preference })
          .from(ModelProviderPreferenceTable)
          .where(and(eq(ModelProviderPreferenceTable.workspaceID, Actor.workspace()), eq(ModelProviderPreferenceTable.model, model)))
          .limit(1)

        if (result.length === 0) return undefined
        try {
          return ProviderPreferenceSchema.parse(JSON.parse(result[0].preference))
        } catch {
          return undefined
        }
      })
    },
  )

  export const deleteProviderPreference = fn(
    z.object({ model: z.string() }),
    ({ model }) => {
      Actor.assertAdmin()
      return Database.use((db) =>
        db.delete(ModelProviderPreferenceTable).where(and(eq(ModelProviderPreferenceTable.workspaceID, Actor.workspace()), eq(ModelProviderPreferenceTable.model, model))),
      )
    },
  )

  export const listProviderPreferences = fn(z.void(), () => {
    return Database.use(async (db) => {
      const results = await db
        .select({ model: ModelProviderPreferenceTable.model, preference: ModelProviderPreferenceTable.provider_preference })
        .from(ModelProviderPreferenceTable)
        .where(eq(ModelProviderPreferenceTable.workspaceID, Actor.workspace()))

      const parsed: Record<string, ProviderPreference> = {}
      for (const row of results) {
        try {
          parsed[row.model] = ProviderPreferenceSchema.parse(JSON.parse(row.preference))
        } catch {}
      }
      return parsed
    })
  })
}
