/** Local embedding server — Ollama-based embedding model lifecycle management */

import { Effect } from "effect"
import * as Log from "@saeeol/core/util/log"
import { Process } from "@/util/process"
import type { EmbeddingServer, EmbedderStatus, HardwareProfile, LLMBakeParams } from "@/ltm/types"
import type { BackendType } from "@/provider/local/types"
import * as GPU from "@/provider/local/gpu"
import * as RAG from "@/provider/local/rag"
import * as Hub from "@/provider/local/hub"
import * as Manager from "@/provider/local/model-manager"
import { iife } from "@/util/iife"
import * as Bus from "@/bus"
import { LTMEvent } from "@/ltm/events"

const log = Log.create({ service: "local/embedder" })

let server: EmbeddingServer | undefined

async function isOllamaRunning(endpoint: string): Promise<boolean> {
  try {
    const res = await fetch(`${endpoint}/api/tags`, { signal: AbortSignal.timeout(3000) })
    return res.ok
  } catch {
    return false
  }
}

async function isModelLoaded(endpoint: string, model: string): Promise<boolean> {
  try {
    const res = await fetch(`${endpoint}/api/tags`, { signal: AbortSignal.timeout(3000) })
    if (!res.ok) return false
    const data = (await res.json()) as { models?: Array<{ name: string }> }
    return (data.models ?? []).some((m) => m.name === model || m.name === `${model}:latest`)
  } catch {
    return false
  }
}

/** Check if an embedding model is installed in Ollama */
async function isModelInstalled(endpoint: string, model: string): Promise<boolean> {
  try {
    const res = await fetch(`${endpoint}/api/tags`, { signal: AbortSignal.timeout(5000) })
    if (!res.ok) return false
    const data = (await res.json()) as { models?: Array<{ name: string }> }
    return (data.models ?? []).some((m) =>
      m.name === model || m.name === `${model}:latest` || m.name.replace(":latest", "") === model,
    )
  } catch {
    return false
  }
}

async function pullModel(endpoint: string, model: string): Promise<void> {
  log.info("pulling embedding model", { model })
  const ollamaModel = getOllamaModelName(model)

  try {
    const res = await fetch(`${endpoint}/api/pull`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({ name: ollamaModel, stream: false }),
      signal: AbortSignal.timeout(300_000), // 5min timeout
    })
    if (!res.ok) {
      const err = await res.text().catch(() => "unknown error")
      throw new Error(`Ollama pull failed: ${res.status} ${err}`)
    }
    log.info("model pulled", { model: ollamaModel })
  } catch (e) {
    log.error("failed to pull model", { model: ollamaModel, error: e })
    throw e
  }
}

/** Map internal model ID to Ollama model name */
function getOllamaModelName(modelId: string): string {
  const map: Record<string, string> = {
    "bge-small-en": "nomic-embed-text", // lightest embedding model in Ollama
    "bge-base-en": "nomic-embed-text",
    "bge-large-en": "nomic-embed-text",
    "bge-m3": "bge-m3",
    "nomic-embed": "nomic-embed-text",
    "all-minilm-l6": "all-minilm",
    "gte-small": "nomic-embed-text",
  }
  return map[modelId] ?? modelId
}

async function embedViaOllama(endpoint: string, model: string, texts: string[]): Promise<number[][]> {
  const ollamaModel = getOllamaModelName(model)
  const res = await fetch(`${endpoint}/api/embed`, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({ model: ollamaModel, input: texts }),
    signal: AbortSignal.timeout(30_000),
  })

  if (!res.ok) {
    const err = await res.text().catch(() => "unknown error")
    throw new Error(`Embedding failed: ${res.status} ${err}`)
  }

  const data = (await res.json()) as { embeddings?: number[][] }
  if (!data.embeddings || !Array.isArray(data.embeddings)) {
    throw new Error("Invalid embedding response")
  }
  return data.embeddings
}

function estimateEmbeddingVRAM(modelId: string): number {
  const model = RAG.EMBEDDING_MODELS.find((m) => m.id === modelId)
  if (!model) return 200
  return Math.ceil(model.sizeBytes * 1.2 / (1024 * 1024))
}

/** Start the embedding server */
export async function start(bake: LLMBakeParams): Promise<EmbeddingServer> {
  if (server && server.status === "running") {
    log.info("embedder already running", { model: server.model })
    return server
  }

  const endpoint = iife(() => {
    const env = process.env.OLLAMA_BASE_URL ?? process.env.OLLAMA_HOST
    return env ?? "http://localhost:11434"
  })

  server = {
    id: "local-embedding",
    model: bake.embeddingModel,
    status: "starting",
    endpoint,
    dimensions: bake.embeddingDimensions,
    vramMB: estimateEmbeddingVRAM(bake.embeddingModel),
  }

  void Bus.publish(LTMEvent.EmbedderStatusChanged, { status: "starting", model: bake.embeddingModel })
  const running = await isOllamaRunning(endpoint)
  if (!running) {
    log.warn("Ollama not running, embedding server unavailable", { endpoint })
    server = { ...server, status: "error" }
    void Bus.publish(LTMEvent.EmbedderStatusChanged, { status: "error", model: bake.embeddingModel })
    return server
  }
  const ollamaModel = getOllamaModelName(bake.embeddingModel)
  const installed = await isModelInstalled(endpoint, ollamaModel)
  if (!installed) {
    log.info("model not installed, pulling", { model: ollamaModel })
    await pullModel(endpoint, ollamaModel)
  }
  const loaded = await isModelLoaded(endpoint, ollamaModel)
  if (!loaded) {
    // Warmup call — load model into memory
    try {
      await embedViaOllama(endpoint, bake.embeddingModel, ["warmup"])
      log.info("model warmed up", { model: ollamaModel })
    } catch (e) {
      log.error("warmup failed", { model: ollamaModel, error: e })
      server = { ...server, status: "error" }
      void Bus.publish(LTMEvent.EmbedderStatusChanged, { status: "error", model: bake.embeddingModel })
      return server
    }
  }

    server = { ...server, status: "running" }
  void Bus.publish(LTMEvent.EmbedderStatusChanged, { status: "running", model: bake.embeddingModel })
  log.info("embedding server started", { model: bake.embeddingModel, dimensions: bake.embeddingDimensions })

  return server
}

/** Stop the embedding server */
export async function stop(): Promise<void> {
  if (!server) return
  server = { ...server!, status: "stopped" }
  void Bus.publish(LTMEvent.EmbedderStatusChanged, { status: "stopped" })
  log.info("embedding server stopped")
  server = undefined
}

/** Current server status */
export function status(): EmbeddingServer | undefined {
  return server
}

/** Embed an array of texts → embedding vectors */
export async function embed(texts: string[]): Promise<number[][]> {
  if (!server || server.status !== "running") {
    throw new Error("Embedding server not running")
  }
  return embedViaOllama(server.endpoint, server.model, texts)
}

/** Embed a single text → embedding vector */
export async function embedOne(text: string): Promise<number[]> {
  const vectors = await embed([text])
  return vectors[0]!
}

/** VRAM usage in MB */
export function vramUsage(): number {
  return server?.vramMB ?? 0
}