/** RAG support — embedding models, rerankers, vector DB management */

import { Effect } from "effect"
import * as Log from "@saeeol/core/util/log"
import type { RAGAsset, RAGAssetType } from "./types"
import * as Hub from "./hub"
import * as Manager from "./model-manager"

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

/** Well-known embedding models with typical dimensions */
export const EMBEDDING_MODELS: Array<{ id: string; repo: string; name: string; dimensions: number; sizeBytes: number }> = [
  { id: "bge-small-en", repo: "BAAI/bge-small-en-v1.5", name: "BGE Small English", dimensions: 384, sizeBytes: 130_000_000 },
  { id: "bge-base-en", repo: "BAAI/bge-base-en-v1.5", name: "BGE Base English", dimensions: 768, sizeBytes: 420_000_000 },
  { id: "bge-large-en", repo: "BAAI/bge-large-en-v1.5", name: "BGE Large English", dimensions: 1024, sizeBytes: 1_300_000_000 },
  { id: "bge-m3", repo: "BAAI/bge-m3", name: "BGE M3 Multilingual", dimensions: 1024, sizeBytes: 2_200_000_000 },
  { id: "nomic-embed", repo: "nomic-ai/nomic-embed-text-v1.5", name: "Nomic Embed Text", dimensions: 768, sizeBytes: 550_000_000 },
  { id: "all-minilm-l6", repo: "sentence-transformers/all-MiniLM-L6-v2", name: "MiniLM L6", dimensions: 384, sizeBytes: 80_000_000 },
  { id: "gte-small", repo: "Alibaba-NLP/gte-base-en-v1.5", name: "GTE Base", dimensions: 768, sizeBytes: 430_000_000 },
]

/** Well-known reranker models */
export const RERANKER_MODELS: Array<{ id: string; repo: string; name: string; sizeBytes: number }> = [
  { id: "bge-reranker-base", repo: "BAAI/bge-reranker-base", name: "BGE Reranker Base", sizeBytes: 420_000_000 },
  { id: "bge-reranker-large", repo: "BAAI/bge-reranker-large", name: "BGE Reranker Large", sizeBytes: 1_300_000_000 },
  { id: "ms-marco-minilm", repo: "cross-encoder/ms-marco-MiniLM-L-6-v2", name: "MS MARCO MiniLM", sizeBytes: 80_000_000 },
]

/** Build a RAGAsset from a known embedding model */
export function embeddingAsset(model: typeof EMBEDDING_MODELS[number]): RAGAsset {
  return {
    id: model.id,
    name: model.name,
    type: "embedding",
    repo: model.repo,
    format: "safetensors",
    sizeBytes: model.sizeBytes,
    dimensions: model.dimensions,
  }
}

/** Build a RAGAsset from a known reranker model */
export function rerankerAsset(model: typeof RERANKER_MODELS[number]): RAGAsset {
  return {
    id: model.id,
    name: model.name,
    type: "reranker",
    repo: model.repo,
    format: "safetensors",
    sizeBytes: model.sizeBytes,
  }
}

/** Install an embedding model for RAG */
export async function installEmbedding(
  model: typeof EMBEDDING_MODELS[number],
  opts?: { onProgress?: (downloaded: number, total: number) => void; signal?: AbortSignal },
): Promise<string> {
  const asset = embeddingAsset(model)
  log.info("installing embedding model", { repo: asset.repo })

  // Download model config and weights
  const files = ["config.json", "tokenizer.json", "tokenizer_config.json"]
  let lastPath = ""
  for (const file of files) {
    try {
      lastPath = await Manager.installRAG(
        { repo: asset.repo, filename: file, sizeBytes: 0 },
        opts,
      )
    } catch { /* non-critical */ }
  }

  // Download safetensors weights
  const weightFile = "model.safetensors"
  lastPath = await Manager.installRAG(
    { repo: asset.repo, filename: weightFile, sizeBytes: asset.sizeBytes },
    opts,
  )

  return lastPath
}

/** Install a reranker model */
export async function installReranker(
  model: typeof RERANKER_MODELS[number],
  opts?: { onProgress?: (downloaded: number, total: number) => void; signal?: AbortSignal },
): Promise<string> {
  const asset = rerankerAsset(model)
  log.info("installing reranker model", { repo: asset.repo })

  const weightFile = "model.safetensors"
  return Manager.installRAG(
    { repo: asset.repo, filename: weightFile, sizeBytes: asset.sizeBytes },
    opts,
  )
}

/** List recommended RAG assets that fit in available VRAM */
export async function recommendForVRAM(vramMB: number): Promise<{
  embeddings: typeof EMBEDDING_MODELS
  rerankers: typeof RERANKER_MODELS
}> {
  const embeddings = EMBEDDING_MODELS.filter((m) => m.sizeBytes * 1.2 <= vramMB * 1024 * 1024)
  const rerankers = RERANKER_MODELS.filter((m) => m.sizeBytes * 1.2 <= vramMB * 1024 * 1024)
  return { embeddings, rerankers }
}

/** Search for RAG-compatible models on HuggingFace */
export async function search(query?: string, type?: RAGAssetType) {
  return Hub.searchRAG(type ?? "embedding", query, 20)
}
