import z from "zod"
import fuzzysort from "fuzzysort"
import { Config } from "../config/config"
import { sortBy } from "remeda"
import { NoSuchModelError, type Provider as SDK } from "ai"
import { Log } from "../util/log"
import { NamedError } from "@opencode-ai/util/error"
import { Auth } from "../auth"
import { Env } from "../env"
import { Instance } from "../project/instance"
import { SecureClient } from "tinfoil"
import { createOpenAICompatible } from "@ai-sdk/openai-compatible"
import type { LanguageModelV2 } from "@ai-sdk/provider"

export namespace Provider {
  const log = Log.create({ service: "provider" })

  const TINFOIL_API_URL = "https://inference.tinfoil.sh/v1"

  interface TinfoilState {
    secureClient: SecureClient
    apiKey: string
    baseURL: string
  }

  let tinfoilState: TinfoilState | null = null

  async function initTinfoil(): Promise<TinfoilState | null> {
    if (tinfoilState) return tinfoilState

    const config = await Config.get()
    const providerConfig = config.provider?.["tinfoil"]

    const apiKey = await (async () => {
      const env = Env.all()
      const envKeys = providerConfig?.env ?? ["TINFOIL_API_KEY"]
      const envKey = envKeys.map((item: string) => env[item]).find(Boolean)
      if (envKey) return envKey
      const auth = await Auth.get("tinfoil")
      if (auth?.type === "api") return auth.key
      if (providerConfig?.options?.apiKey) return providerConfig.options.apiKey
      return undefined
    })()

    if (!apiKey) return null

    const baseURL = providerConfig?.options?.baseURL
    const enclaveURL = providerConfig?.options?.enclaveURL
    const configRepo = providerConfig?.options?.configRepo
    const transport = providerConfig?.options?.transport

    const secureClient = new SecureClient({
      baseURL,
      enclaveURL,
      configRepo,
      transport: transport ?? "tls",
    })

    await secureClient.ready()

    const finalBaseURL = baseURL || secureClient.getBaseURL() || TINFOIL_API_URL

    tinfoilState = {
      secureClient,
      apiKey,
      baseURL: finalBaseURL,
    }

    return tinfoilState
  }

  interface TinfoilModel {
    id: string
    object: string
    created: number
    owned_by: string
    type: string
    context_window?: number
    multimodal?: boolean
  }

  interface TinfoilModelsResponse {
    object: string
    data: TinfoilModel[]
  }

  async function fetchTinfoilModels(state: TinfoilState): Promise<Record<string, Model>> {
    try {
      const response = await state.secureClient.fetch(`${state.baseURL}/models`, {
        headers: {
          Authorization: `Bearer ${state.apiKey}`,
        },
      })

      if (!response.ok) {
        log.error("Failed to fetch tinfoil models", { status: response.status })
        return {}
      }

      const data = (await response.json()) as TinfoilModelsResponse
      const models: Record<string, Model> = {}

      // Only include chat and code models
      const supportedTypes = ["chat", "code"]

      for (const model of data.data) {
        if (!supportedTypes.includes(model.type)) {
          continue
        }

        const supportsImage = model.multimodal === true

        models[model.id] = {
          id: model.id,
          providerID: "tinfoil",
          name: model.id,
          api: {
            id: model.id,
            url: state.baseURL,
            npm: "tinfoil",
          },
          status: "active",
          headers: {},
          options: {},
          cost: {
            input: 0,
            output: 0,
            cache: { read: 0, write: 0 },
          },
          limit: {
            context: model.context_window ?? 128000,
            output: 4096,
          },
          capabilities: {
            temperature: true,
            reasoning: false,
            attachment: supportsImage,
            toolcall: true,
            input: { text: true, audio: false, image: supportsImage, video: false, pdf: false },
            output: { text: true, audio: false, image: false, video: false, pdf: false },
            interleaved: false,
          },
          release_date: new Date().toISOString().split("T")[0],
        }
      }

      return models
    } catch (e) {
      log.error("Error fetching tinfoil models", { error: e })
      return {}
    }
  }

  export const Model = z
    .object({
      id: z.string(),
      providerID: z.string(),
      api: z.object({
        id: z.string(),
        url: z.string(),
        npm: z.string(),
      }),
      name: z.string(),
      family: z.string().optional(),
      capabilities: z.object({
        temperature: z.boolean(),
        reasoning: z.boolean(),
        attachment: z.boolean(),
        toolcall: z.boolean(),
        input: z.object({
          text: z.boolean(),
          audio: z.boolean(),
          image: z.boolean(),
          video: z.boolean(),
          pdf: z.boolean(),
        }),
        output: z.object({
          text: z.boolean(),
          audio: z.boolean(),
          image: z.boolean(),
          video: z.boolean(),
          pdf: z.boolean(),
        }),
        interleaved: z.union([
          z.boolean(),
          z.object({
            field: z.enum(["reasoning_content", "reasoning_details"]),
          }),
        ]),
      }),
      cost: z.object({
        input: z.number(),
        output: z.number(),
        cache: z.object({
          read: z.number(),
          write: z.number(),
        }),
        experimentalOver200K: z
          .object({
            input: z.number(),
            output: z.number(),
            cache: z.object({
              read: z.number(),
              write: z.number(),
            }),
          })
          .optional(),
      }),
      limit: z.object({
        context: z.number(),
        output: z.number(),
      }),
      status: z.enum(["alpha", "beta", "deprecated", "active"]),
      options: z.record(z.string(), z.any()),
      headers: z.record(z.string(), z.string()),
      release_date: z.string(),
      variants: z.record(z.string(), z.record(z.string(), z.any())).optional(),
    })
    .meta({
      ref: "Model",
    })
  export type Model = z.infer<typeof Model>

  export const Info = z
    .object({
      id: z.string(),
      name: z.string(),
      source: z.enum(["env", "config", "custom", "api"]),
      env: z.string().array(),
      key: z.string().optional(),
      options: z.record(z.string(), z.any()),
      models: z.record(z.string(), Model),
    })
    .meta({
      ref: "Provider",
    })
  export type Info = z.infer<typeof Info>

  // Simplified state for tinfoil-only provider
  const state = Instance.state(async () => {
    using _ = log.time("state")

    const tinfoil = await initTinfoil()
    if (!tinfoil) {
      log.warn("Tinfoil not configured - no API key found")
      return {
        models: new Map<string, LanguageModelV2>(),
        providers: {} as { [providerID: string]: Info },
        sdk: new Map<number, SDK>(),
      }
    }

    const models = await fetchTinfoilModels(tinfoil)

    const provider: Info = {
      id: "tinfoil",
      name: "Tinfoil",
      source: "env",
      env: ["TINFOIL_API_KEY"],
      key: tinfoil.apiKey,
      options: {
        baseURL: tinfoil.baseURL,
        apiKey: tinfoil.apiKey,
        fetch: tinfoil.secureClient.fetch,
      },
      models,
    }

    log.info("found", { providerID: "tinfoil", modelCount: Object.keys(models).length })

    return {
      models: new Map<string, LanguageModelV2>(),
      providers: { tinfoil: provider },
      sdk: new Map<number, SDK>(),
    }
  })

  export async function list() {
    return state().then((s) => s.providers)
  }

  async function getSDK(model: Model): Promise<SDK> {
    const s = await state()
    const provider = s.providers[model.providerID]
    if (!provider) throw new InitError({ providerID: model.providerID })

    const options = { ...provider.options }
    options["includeUsage"] = true

    const key = Bun.hash.xxHash32(JSON.stringify({ options }))
    const existing = s.sdk.get(key)
    if (existing) return existing

    const customFetch = options["fetch"]

    const loaded = createOpenAICompatible({
      name: "tinfoil",
      baseURL: options["baseURL"] as string,
      apiKey: options["apiKey"] as string,
      fetch: customFetch ?? fetch,
    })
    s.sdk.set(key, loaded)
    return loaded as SDK
  }

  export async function getProvider(providerID: string) {
    return state().then((s) => s.providers[providerID])
  }

  export async function getModel(providerID: string, modelID: string) {
    const s = await state()
    const provider = s.providers[providerID]
    if (!provider) {
      const availableProviders = Object.keys(s.providers)
      const matches = fuzzysort.go(providerID, availableProviders, { limit: 3, threshold: -10000 })
      const suggestions = matches.map((m) => m.target)
      throw new ModelNotFoundError({ providerID, modelID, suggestions })
    }

    const info = provider.models[modelID]
    if (!info) {
      const availableModels = Object.keys(provider.models)
      const matches = fuzzysort.go(modelID, availableModels, { limit: 3, threshold: -10000 })
      const suggestions = matches.map((m) => m.target)
      throw new ModelNotFoundError({ providerID, modelID, suggestions })
    }
    return info
  }

  export async function getLanguage(model: Model): Promise<LanguageModelV2> {
    const s = await state()
    const key = `${model.providerID}/${model.id}`
    if (s.models.has(key)) return s.models.get(key)!

    const sdk = await getSDK(model)

    try {
      const language = sdk.languageModel(model.api.id) as LanguageModelV2
      s.models.set(key, language)
      return language
    } catch (e) {
      if (e instanceof NoSuchModelError)
        throw new ModelNotFoundError(
          {
            modelID: model.id,
            providerID: model.providerID,
          },
          { cause: e },
        )
      throw e
    }
  }

  export async function closest(providerID: string, query: string[]) {
    const s = await state()
    const provider = s.providers[providerID]
    if (!provider) return undefined
    for (const item of query) {
      for (const modelID of Object.keys(provider.models)) {
        if (modelID.includes(item))
          return {
            providerID,
            modelID,
          }
      }
    }
  }

  export async function getSmallModel(_providerID: string) {
    const cfg = await Config.get()

    if (cfg.small_model) {
      const parsed = parseModel(cfg.small_model)
      return getModel(parsed.providerID, parsed.modelID)
    }

    const provider = await state().then((s) => s.providers["tinfoil"])
    if (provider) {
      const models = Object.keys(provider.models)
      if (models.length > 0) {
        return getModel("tinfoil", models[0])
      }
    }

    return undefined
  }

  export function sort(models: Model[]) {
    return sortBy(models, [(model) => model.id, "asc"])
  }

  export async function defaultModel() {
    const cfg = await Config.get()
    if (cfg.model) return parseModel(cfg.model)

    const provider = await list().then((val) => val["tinfoil"])
    if (!provider) throw new Error("Tinfoil provider not configured. Set TINFOIL_API_KEY environment variable.")
    const [model] = sort(Object.values(provider.models))
    if (!model) throw new Error("No models available from Tinfoil")
    return {
      providerID: "tinfoil",
      modelID: model.id,
    }
  }

  export function parseModel(model: string) {
    const [providerID, ...rest] = model.split("/")
    return {
      providerID: providerID,
      modelID: rest.join("/"),
    }
  }

  export const ModelNotFoundError = NamedError.create(
    "ProviderModelNotFoundError",
    z.object({
      providerID: z.string(),
      modelID: z.string(),
      suggestions: z.array(z.string()).optional(),
    }),
  )

  export const InitError = NamedError.create(
    "ProviderInitError",
    z.object({
      providerID: z.string(),
    }),
  )
}