/** Model orchestrator — run multiple models simultaneously with GPU-aware scheduling */

import { Effect } from "effect"
import * as Log from "@saeeol/core/util/log"
import { Process } from "@/util/process"
import { Global } from "@saeeol/core/global"
import type { ModelArtifact, ModelInstance, BackendType, BackendStatus, GPUProfile } from "./types"
import * as GPU from "./gpu"
import * as Manager from "./model-manager"
import { which } from "@/util/which"
import { iife } from "@/util/iife"
import path from "path"
import { mkdir } from "fs/promises"

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

const running = new Map<string, ModelInstance>()
let nextPort = 11435 // Start after default Ollama port

/** Detect which local backends are installed */
export async function detectBackends(): Promise<BackendStatus[]> {
  const backends: Array<{ type: BackendType; cmds: string[]; defaultPort: number }> = [
    { type: "ollama", cmds: ["ollama"], defaultPort: 11434 },
    { type: "lmstudio", cmds: ["lms"], defaultPort: 1234 },
    { type: "vllm", cmds: ["vllm", "python -m vllm"], defaultPort: 8000 },
    { type: "llama.cpp", cmds: ["llama-server", "llama-cpp-server"], defaultPort: 8080 },
  ]

  const results: BackendStatus[] = []
  for (const b of backends) {
    let cmd = ""
    let available = false
    for (const c of b.cmds) {
      const found = await which(c.split(" ")[0])
      if (found) { cmd = c; available = true; break }
    }

    results.push({
      type: b.type,
      available,
      endpoint: `http://localhost:${b.defaultPort}`,
      version: undefined,
      loadedModels: [],
    })
  }

  return results
}

/** Check if a backend server is responding */
export async function isAlive(endpoint: string): Promise<boolean> {
  try {
    const res = await fetch(endpoint, { signal: AbortSignal.timeout(3000) })
    return res.ok || res.status === 404 // 404 is fine, means server is up
  } catch {
    return false
  }
}

/** Allocate a free port for a new model instance */
function allocPort(): number {
  return nextPort++
}

/** Load a model via Ollama */
async function loadViaOllama(artifact: ModelArtifact, modelPath: string): Promise<ModelInstance> {
  const port = 11434 // Ollama default

  // Register model with Ollama if not already present
  const modelfile = `FROM ${modelPath}`
  const name = artifact.repo.replace("/", "-").toLowerCase()

  const result = await Process.run(["ollama", "create", name, "-f", "-"])

  log.info("ollama create result", { code: result.code })

  return {
    id: name,
    artifact,
    status: "running",
    port,
    endpoint: `http://localhost:${port}`,
  }
}

/** Load a model via llama.cpp server */
async function loadViaLlamaCpp(artifact: ModelArtifact, modelPath: string, gpuIndex?: number): Promise<ModelInstance> {
  const port = allocPort()

  const args = ["-m", modelPath, "--port", String(port), "-c", "4096", "-ngl", "99"]
  if (gpuIndex !== undefined) args.push("--gpu", String(gpuIndex))

  const proc = Bun.spawn(["llama-server", ...args], {
    stdout: "pipe",
    stderr: "pipe",
    windowsHide: true,
    detached: true,
  })

  // Wait for server to start
  for (let i = 0; i < 30; i++) {
    await new Promise((r) => setTimeout(r, 1000))
    if (await isAlive(`http://localhost:${port}`)) break
  }

  const instance: ModelInstance = {
    id: `llamacpp-${port}`,
    artifact,
    status: "running",
    pid: proc.pid,
    port,
    gpuIndex,
    endpoint: `http://localhost:${port}/v1`,
  }

  running.set(instance.id, instance)
  return instance
}

/** Load a model via vLLM */
async function loadViaVLLM(artifact: ModelArtifact, modelPath: string, gpuIndex?: number): Promise<ModelInstance> {
  const port = allocPort()

  const args = [
    "serve", artifact.repo,
    "--port", String(port),
    "--served-model-name", artifact.repo,
    "--trust-remote-code",
  ]

  const proc = Bun.spawn(["vllm", ...args], {
    stdout: "pipe",
    stderr: "pipe",
    windowsHide: true,
    detached: true,
    env: {
      ...process.env,
      ...(gpuIndex !== undefined ? { CUDA_VISIBLE_DEVICES: String(gpuIndex) } : {}),
    },
  })

  // vLLM takes longer to start
  for (let i = 0; i < 120; i++) {
    await new Promise((r) => setTimeout(r, 2000))
    if (await isAlive(`http://localhost:${port}`)) break
  }

  const instance: ModelInstance = {
    id: `vllm-${port}`,
    artifact,
    status: "running",
    pid: proc.pid,
    port,
    gpuIndex,
    endpoint: `http://localhost:${port}/v1`,
  }

  running.set(instance.id, instance)
  return instance
}

/** Smart-load a model: pick the best backend and GPU automatically */
export async function load(
  artifact: ModelArtifact,
  opts?: { backend?: BackendType; gpuIndex?: number },
): Promise<ModelInstance> {
  const gpuProfile = await Effect.runPromise(GPU.profile)
  const vramNeeded = GPU.estimateVRAM(artifact.sizeBytes, artifact.quantization)

  // Pick GPU
  const gpuIndex = opts?.gpuIndex ?? (() => {
    const gpu = GPU.findBestGPU(gpuProfile, vramNeeded)
    return gpu?.index
  })()

  if (gpuIndex === undefined && gpuProfile.cudaAvailable) {
    log.warn("insufficient VRAM, model may run on CPU", { needed: vramNeeded, available: gpuProfile.availableVRAMMB })
  }

  // Install if needed
  const modelPath = await Manager.install(artifact)

  // Pick backend
  const backends = await detectBackends()
  const preferred = opts?.backend ?? iife(() => {
    if (artifact.format === "gguf") {
      if (backends.find((b) => b.type === "llama.cpp")?.available) return "llama.cpp" as BackendType
      if (backends.find((b) => b.type === "ollama")?.available) return "ollama" as BackendType
    }
    if (backends.find((b) => b.type === "vllm")?.available) return "vllm" as BackendType
    if (backends.find((b) => b.type === "ollama")?.available) return "ollama" as BackendType
    return "llama.cpp" as BackendType
  })

  log.info("loading model", { repo: artifact.repo, backend: preferred, gpuIndex })

  // Load via chosen backend
  const instance: ModelInstance = {
    id: `${preferred}-${artifact.repo}`,
    artifact,
    status: "starting",
    gpuIndex,
    vramUsageMB: vramNeeded,
  }

  running.set(instance.id, { ...instance })

  switch (preferred) {
    case "ollama":
      return loadViaOllama(artifact, modelPath)
    case "llama.cpp":
      return loadViaLlamaCpp(artifact, modelPath, gpuIndex)
    case "vllm":
      return loadViaVLLM(artifact, modelPath, gpuIndex)
    default:
      return loadViaLlamaCpp(artifact, modelPath, gpuIndex)
  }
}

/** Unload a running model */
export async function unload(instanceId: string): Promise<void> {
  const inst = running.get(instanceId)
  if (!inst) return

  if (inst.pid) {
    try { process.kill(inst.pid) } catch { /* already dead */ }
  }

  ;(inst as any).status = "stopped"
  ;(inst as any).pid = undefined
  running.delete(instanceId)
  log.info("unloaded model", { id: instanceId })
}

/** List all running model instances */
export function runningInstances(): ModelInstance[] {
  return [...running.values()]
}

/** Get a running instance by its endpoint */
export function findByEndpoint(url: string): ModelInstance | undefined {
  return [...running.values()].find((i) => i.endpoint === url)
}

/** Load multiple models with GPU-aware scheduling */
export async function loadBatch(
  artifacts: ModelArtifact[],
): Promise<ModelInstance[]> {
  const gpuProfile = await Effect.runPromise(GPU.profile)
  const results: ModelInstance[] = []

  // Sort by size descending — load largest first to ensure they get GPU space
  const sorted = [...artifacts].sort((a, b) => b.sizeBytes - a.sizeBytes)

  // Track per-GPU allocation
  const gpuAlloc = new Map<number, number>() // gpuIndex -> usedMB
  for (const g of gpuProfile.gpus) gpuAlloc.set(g.index, g.vramTotalMB - g.vramFreeMB)

  for (const artifact of sorted) {
    const vramNeeded = GPU.estimateVRAM(artifact.sizeBytes, artifact.quantization)

    // Find GPU with most remaining space
    let bestGpu: number | undefined
    let bestFree = 0
    for (const [idx, used] of gpuAlloc) {
      const gpu = gpuProfile.gpus.find((g) => g.index === idx)
      if (!gpu) continue
      const free = gpu.vramTotalMB - used
      if (free >= vramNeeded && free > bestFree) {
        bestGpu = idx
        bestFree = free
      }
    }

    if (bestGpu !== undefined) {
      gpuAlloc.set(bestGpu, (gpuAlloc.get(bestGpu) ?? 0) + vramNeeded)
    }

    try {
      const instance = await load(artifact, { gpuIndex: bestGpu })
      results.push(instance)
    } catch (e) {
      log.error("failed to load model in batch", { repo: artifact.repo, error: e })
      results.push({
        id: `error-${artifact.repo}`,
        artifact,
        status: "error",
        error: e instanceof Error ? e.message : String(e),
      })
    }
  }

  return results
}

/** Unload all running models */
export async function unloadAll(): Promise<void> {
  for (const id of [...running.keys()]) {
    await unload(id)
  }
}
