import { Effect } from "effect"
import { mapValues, mergeDeep } from "remeda"
import * as Log from "@saeeol/core/util/log"
import { EffectBridge } from "@/effect/bridge"
import { InstanceState } from "@/effect/instance-state"
import { applyConfigModels, cleanupProviders, fromModelsDevProvider } from "./provider-conversion"
import { custom, LOCAL_PROVIDERS, CUSTOM_PROVIDERS, CLOUD_PROVIDERS } from "./custom-loaders"
import { ModelID, ProviderID } from "./schema"
import type { State, CustomModelLoader, CustomVarsLoader, CustomDiscoverModels } from "./provider-types"
import type { Info } from "./provider-schemas"
import type { ProviderCategory } from "./provider-schemas"
import {
  saeeolCustomLoaders,
  patchCustomLoaderResult,
} from "@/saeeol/provider/provider"
import type { Config } from "@/config/config"
import type { Auth } from "../auth"
import type { Plugin } from "../plugin"
import type { Hooks } from "@saeeol/plugin"
import type { Env } from "../env"
import type { AppFileSystem } from "@saeeol/core/filesystem"
import * as ModelsDev from "./models"

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

export function initState(deps: {
  fs: AppFileSystem.Interface
  config: { get: () => Effect.Effect<Config.Info> }
  auth: { get: (id: string) => Effect.Effect<Auth.Info | undefined>; all: () => Effect.Effect<Record<string, Auth.Info>> }
  env: { all: () => Effect.Effect<Record<string, string | undefined>>; get: (key: string) => Effect.Effect<string | undefined> }
  plugin: { list: () => Effect.Effect<Hooks[]> }
  modelsDevSvc: { get: () => Effect.Effect<Record<string, ModelsDev.Provider>> }
}) {
  return InstanceState.make<State>(() =>
    Effect.gen(function* () {
      using _ = log.time("state")
      const bridge = yield* EffectBridge.make()
      const cfg = yield* deps.config.get()
      const modelsDev = yield* deps.modelsDevSvc.get()
      const database = mapValues(modelsDev, fromModelsDevProvider)

      const providers: Record<ProviderID, Info> = {} as Record<ProviderID, Info>
      const languages = new Map<string, any>()
      const modelLoaders: Record<string, CustomModelLoader> = {}
      const varsLoaders: Record<string, CustomVarsLoader> = {}
      const sdk = new Map<string, any>()
      const discoveryLoaders: Record<string, CustomDiscoverModels> = {}
      const dep = {
        auth: (id: string) => deps.auth.get(id).pipe(Effect.orDie),
        config: () => deps.config.get(),
        env: () => deps.env.all(),
        get: (key: string) => deps.env.get(key),
      }

      log.info("init")

      function getCategory(id: string): ProviderCategory {
        if (LOCAL_PROVIDERS.includes(id as any)) return "local"
        if (CUSTOM_PROVIDERS.includes(id as any)) return "custom"
        return "cloud"
      }

      function mergeProvider(providerID: ProviderID, provider: Partial<Info>) {
        const existing = providers[providerID]
        if (!provider.category) provider.category = getCategory(providerID)
        if (existing) {
          // @ts-expect-error
          providers[providerID] = mergeDeep(existing, provider)
          return
        }
        const match = database[providerID]
        if (!match) return
        // @ts-expect-error
        providers[providerID] = mergeDeep(match, provider)
      }

      const plugins = yield* deps.plugin.list()

      const configProviders = Object.entries(cfg.provider ?? {})
      const disabled = new Set(cfg.disabled_providers ?? [])
      const enabled = cfg.enabled_providers ? new Set(cfg.enabled_providers) : null

      function isProviderAllowed(providerID: ProviderID): boolean {
        if (enabled && !enabled.has(providerID)) return false
        if (disabled.has(providerID)) return false
        return true
      }

      for (const hook of plugins) {
        const p = hook.provider
        const models = p?.models
        if (!p || !models) continue

        const providerID = ProviderID.make(p.id)
        if (disabled.has(providerID)) continue

        const provider = database[providerID]
        if (!provider) continue
        const pluginAuth = yield* deps.auth.get(providerID).pipe(Effect.orDie)

        provider.models = yield* Effect.promise(async () => {
          const next = await models(provider, { auth: pluginAuth })
          return Object.fromEntries(
            Object.entries(next).map(([id, model]: [string, any]) => [
              id,
              {
                ...(model as Record<string, unknown>),
                id: ModelID.make(id),
                providerID,
              },
            ]),
          ) as any
        })
      }

      applyConfigModels(database, configProviders, modelsDev)

      const auths = yield* deps.auth.all().pipe(Effect.orDie)
      const envs = yield* deps.env.all()
      for (const [id, provider] of Object.entries(database)) {
        const providerID = ProviderID.make(id)
        if (disabled.has(providerID)) continue
        if (
          auths[providerID]?.type === "oauth" &&
          plugins.some((x) => x.auth?.provider === providerID && x.auth.loader)
        ) {
          continue
        }
        const apiKey = provider.env.map((item) => envs[item]).find(Boolean)
        if (!apiKey) continue
        mergeProvider(providerID, {
          source: "env",
          key: provider.env.length === 1 ? apiKey : undefined,
        })
      }

      for (const [id, provider] of Object.entries(auths)) {
        const providerID = ProviderID.make(id)
        if (disabled.has(providerID)) continue
        if (provider.type === "api") {
          mergeProvider(providerID, {
            source: "api",
            key: provider.key,
          })
        }
      }

      for (const plugin of plugins) {
        if (!plugin.auth) continue
        const providerID = ProviderID.make(plugin.auth.provider)
        if (disabled.has(providerID)) continue

        const stored = yield* deps.auth.get(providerID).pipe(Effect.orDie)
        if (!stored) continue
        if (!plugin.auth.loader) continue

        const options = yield* Effect.promise(() =>
          plugin.auth!.loader!(
            () => bridge.promise(deps.auth.get(providerID).pipe(Effect.orDie)) as any,
            database[plugin.auth!.provider],
          ),
        )
        const opts = options ?? {}
        const patch: Partial<Info> = providers[providerID] ? { options: opts } : { source: "custom", options: opts }
        mergeProvider(providerID, patch)
      }

      const saeeolEnv = yield* deps.env.all()
      for (const [id, fn] of Object.entries({ ...custom(dep), ...saeeolCustomLoaders(dep) })) {
        const providerID = ProviderID.make(id)
        if (disabled.has(providerID)) continue
        const data = database[providerID]
        if (!data) {
          log.error("Provider does not exist in model list " + providerID)
          continue
        }
        const result = yield* fn(data)
        if (result) patchCustomLoaderResult(id, result, saeeolEnv) 
        if (result && (result.autoload || providers[providerID])) {
          if (result.getModel) modelLoaders[providerID] = result.getModel
          if (result.vars) varsLoaders[providerID] = result.vars
          if (result.discoverModels) discoveryLoaders[providerID] = result.discoverModels
          const opts = result.options ?? {}
          const patch: Partial<Info> = providers[providerID] ? { options: opts } : { source: "custom", options: opts }
          mergeProvider(providerID, patch)
        }
      }

      for (const [id, provider] of configProviders) {
        if (!provider) continue 
        const providerID = ProviderID.make(id)
        const oauth =
          auths[providerID]?.type === "oauth" && plugins.some((x) => x.auth?.provider === providerID && x.auth.loader)
        const partial: Partial<Info> = oauth ? {} : { source: "config" }
        if (provider.env) partial.env = provider.env
        if (provider.name) partial.name = provider.name
        if (provider.options) partial.options = provider.options
        mergeProvider(providerID, partial)
      }

      const gitlab = ProviderID.make("gitlab")
      if (discoveryLoaders[gitlab] && providers[gitlab] && isProviderAllowed(gitlab)) {
        yield* Effect.promise(async () => {
          try {
            const discovered = await discoveryLoaders[gitlab]()
            for (const [modelID, model] of Object.entries(discovered)) {
              if (!providers[gitlab].models[modelID]) {
                providers[gitlab].models[modelID] = model
              }
            }
          } catch (e) {
            log.warn("state discovery error", { id: "gitlab", error: e })
          }
        })
      }

      cleanupProviders(providers, cfg, isProviderAllowed)

      return {
        models: languages,
        providers,
        sdk,
        modelLoaders,
        varsLoaders,
      }
    }),
  )
}
