/** LTM — VRAM/task scheduler */

import * as Log from "@saeeol/core/util/log"
import type { HardwareProfile } from "./types"
import * as GPU from "@/provider/local/gpu"
import * as Embedder from "@/provider/local/embedder"
import { Effect } from "effect"

const log = Log.create({ service: "ltm/scheduler" })

/** Current allocation state */
export interface Allocation {
  llm: number
  embedding: number
  available: number
}

/** Scheduling strategy */
export type Strategy = "concurrent" | "alternating" | "cpu-fallback" | "no-gpu"

/** Query current VRAM allocation state */
export async function allocation(): Promise<Allocation> {
  const gpu = await Effect.runPromise(GPU.profile)
  const embVRAM = Embedder.vramUsage()
  return {
    llm: gpu.totalVRAMMB - gpu.availableVRAMMB - embVRAM,
    embedding: embVRAM,
    available: gpu.availableVRAMMB,
  }
}

/** Return the scheduling strategy */
export async function strategy(): Promise<Strategy> {
  const gpu = await Effect.runPromise(GPU.profile)
  if (!gpu.cudaAvailable) return "no-gpu"
  if (gpu.totalVRAMMB >= 16384) return "concurrent"
  if (gpu.totalVRAMMB >= 8192) return "alternating"
  return "cpu-fallback"
}

/** Whether embedding background work can run */
export async function canRunEmbedding(): Promise<boolean> {
  const strat = await strategy()
  // no-gpu or cpu-fallback: always allowed (CPU processing)
  if (strat === "no-gpu" || strat === "cpu-fallback") return true
  if (strat === "concurrent") return true
  // alternating: need at least 2GB VRAM
  const alloc = await allocation()
  return alloc.available >= 2048
}

/** Whether LLM and embedding can run concurrently */
export async function canRunConcurrent(hw: HardwareProfile): Promise<boolean> {
  if (hw.totalVRAMMB >= 16384) return true
  if (hw.totalVRAMMB >= 8192 && Embedder.vramUsage() <= 1024) return true
  return false
}

/** Called when LLM requests VRAM — pauses embedding if needed */
export async function requestLLM(vramNeededMB: number): Promise<boolean> {
  const strat = await strategy()
  if (strat === "concurrent") return true
  if (strat === "no-gpu") return true

  const alloc = await allocation()
  if (alloc.available >= vramNeededMB) {
    log.info("LLM request: enough VRAM", { available: alloc.available, needed: vramNeededMB })
    return true
  }

  // alternating strategy: pause embedding to free VRAM
  if (strat === "alternating" && alloc.embedding > 0) {
    log.info("LLM request: pausing embedding to free VRAM", {
      embedding: alloc.embedding,
      needed: vramNeededMB,
    })
    await Embedder.stop()
    return true
  }

  // cpu-fallback: embedding already runs on CPU, GPU is LLM-exclusive
  if (strat === "cpu-fallback") return true

  log.warn("LLM request: insufficient VRAM", { available: alloc.available, needed: vramNeededMB })
  return false
}

/** Resume embedding server — called when LLM goes idle */
export async function resumeEmbedding(bake: {
  embeddingModel: string
  embeddingDimensions: number
}): Promise<void> {
  const current = Embedder.status()
  if (current?.status === "running") return

  const canRun = await canRunEmbedding()
  if (!canRun) {
    log.info("embedder resume skipped — VRAM budget exceeded")
    return
  }

  log.info("resuming embedding server after LLM idle")
  try {
    await Embedder.start({
      embeddingModel: bake.embeddingModel,
      embeddingDimensions: bake.embeddingDimensions,
      embeddingVRAMMB: 0,
      contextLength: 4096,
      batchSize: 16,
      numThread: 4,
      numGPU: 0,
      hardwareHash: "",
    })
  } catch (e) {
    log.error("failed to resume embedding server", { error: e })
  }
}