import { test, expect, describe } from "bun:test"
import { Schema } from "effect"
import {
  GPUInfo, GPUProfile, Quantization, ModelFormat,
  ModelArtifact, ModelStatus, ModelInstance,
  BackendType, BackendStatus,
  RAGAssetType, RAGAsset,
  DownloadState, HFModelSearch,
} from "@/provider/local/types"
import * as GPU from "@/provider/local/gpu"
import * as Hub from "@/provider/local/hub"
import * as RAG from "@/provider/local/rag"

// ── Types ──

describe("local types", () => {
  test("Quantization accepts known values", () => {
    const vals = ["q4_k_m", "q5_k_m", "q8_0", "fp16", "bf16", "fp32"] as const
    for (const v of vals) {
      expect(Schema.decodeSync(Quantization)(v)).toBe(v)
    }
  })

  test("Quantization rejects invalid", () => {
    expect(() => Schema.decodeSync(Quantization)("invalid" as any)).toThrow()
  })

  test("ModelFormat accepts known values", () => {
    for (const v of ["gguf", "safetensors", "pytorch", "onnx", "awq", "gptq"] as const) {
      expect(Schema.decodeSync(ModelFormat)(v)).toBe(v)
    }
  })

  test("ModelStatus accepts all states", () => {
    for (const v of ["stopped", "starting", "running", "error", "downloading"] as const) {
      expect(Schema.decodeSync(ModelStatus)(v)).toBe(v)
    }
  })

  test("BackendType accepts known backends", () => {
    for (const v of ["ollama", "lmstudio", "vllm", "llama.cpp", "text-generation-webui"] as const) {
      expect(Schema.decodeSync(BackendType)(v)).toBe(v)
    }
  })

  test("RAGAssetType accepts known types", () => {
    for (const v of ["embedding", "reranker", "vectordb"] as const) {
      expect(Schema.decodeSync(RAGAssetType)(v)).toBe(v)
    }
  })
})

// ── GPU ──

describe("GPU utilities", () => {
  const mockProfile: GPUProfile = {
    gpus: [
      { index: 0, name: "RTX 4090", vramTotalMB: 24564, vramUsedMB: 2000, vramFreeMB: 22564, computeCapability: "8.9" },
      { index: 1, name: "RTX 3090", vramTotalMB: 24564, vramUsedMB: 8000, vramFreeMB: 16564, computeCapability: "8.6" },
    ],
    totalVRAMMB: 49128,
    availableVRAMMB: 39128,
    cudaAvailable: true,
  }

  test("estimateVRAM returns MB with 20% overhead", () => {
    // 1GB model = ~1220MB VRAM needed
    const vram = GPU.estimateVRAM(1024 * 1024 * 1024, "q4_k_m")
    expect(vram).toBeGreaterThan(1000)
    expect(vram).toBeLessThan(1500)
  })

  test("findBestGPU picks GPU with most free VRAM", () => {
    const gpu = GPU.findBestGPU(mockProfile, 10000)
    expect(gpu).toBeDefined()
    expect(gpu!.index).toBe(0) // 22564MB free > 16564MB free
  })

  test("findBestGPU returns undefined if nothing fits", () => {
    const gpu = GPU.findBestGPU(mockProfile, 50000)
    expect(gpu).toBeUndefined()
  })

  test("canFit returns true when enough VRAM", () => {
    expect(GPU.canFit(mockProfile, 30000)).toBe(true)
  })

  test("canFit returns false when not enough VRAM", () => {
    expect(GPU.canFit(mockProfile, 50000)).toBe(false)
  })

  test("suggestQuantization picks highest quality that fits", () => {
    const sizes: Record<string, number> = {
      fp16: 14_000_000_000,
      q8_0: 8_000_000_000,
      q5_k_m: 5_000_000_000,
      q4_k_m: 4_000_000_000,
      q2_k: 2_500_000_000,
    }
    const result = GPU.suggestQuantization(10000, sizes)
    expect(result).toBeDefined()
    if (result) expect(["q4_k_m", "q5_k_m", "q8_0", "fp16"]).toContain(result)
  })

  test("suggestQuantization returns undefined if nothing fits", () => {
    const sizes: Record<string, number> = {
      fp16: 100_000_000_000,
      q2_k: 50_000_000_000,
    }
    const result = GPU.suggestQuantization(100, sizes)
    expect(result).toBeUndefined()
  })
})

// ── Hub ──

describe("HuggingFace Hub utilities", () => {
  test("parseQuant extracts quantization from GGUF filename", () => {
    expect(Hub.parseQuant("model-q4_k_m.gguf")).toBe("q4_k_m")
    expect(Hub.parseQuant("Llama-3.1-8B-Q5_K_M.gguf")).toBe("q5_k_m")
    expect(Hub.parseQuant("model-q8_0.gguf")).toBe("q8_0")
    expect(Hub.parseQuant("model-fp16.gguf")).toBe("fp16")
  })

  test("parseQuant returns undefined for unknown", () => {
    expect(Hub.parseQuant("model.bin")).toBeUndefined()
    expect(Hub.parseQuant("config.json")).toBeUndefined()
  })

  test("listGGUF filters GGUF files from siblings", () => {
    const siblings = [
      { rfilename: "model-q4_k_m.gguf" },
      { rfilename: "config.json" },
      { rfilename: "tokenizer.json" },
      { rfilename: "model-q8_0.gguf" },
    ]
    const ggufs = Hub.listGGUF(siblings)
    expect(ggufs).toHaveLength(2)
    expect(ggufs).toContain("model-q4_k_m.gguf")
    expect(ggufs).toContain("model-q8_0.gguf")
  })

  test("downloadURL builds correct URL", () => {
    const url = Hub.downloadURL("meta-llama/Llama-3.1-8B", "model.gguf")
    expect(url).toBe("https://huggingface.co/meta-llama/Llama-3.1-8B/resolve/main/model.gguf")
  })

  test("buildArtifact creates valid artifact", () => {
    const artifact = Hub.buildArtifact("meta-llama/Llama-3.1-8B", "model-q4_k_m.gguf", 5_000_000_000, "abc123")
    expect(artifact.repo).toBe("meta-llama/Llama-3.1-8B")
    expect(artifact.filename).toBe("model-q4_k_m.gguf")
    expect(artifact.format).toBe("gguf")
    expect(artifact.quantization).toBe("q4_k_m")
    expect(artifact.sizeBytes).toBe(5_000_000_000)
    expect(artifact.sha256).toBe("abc123")
  })

  test("buildArtifact detects safetensors format", () => {
    const artifact = Hub.buildArtifact("BAAI/bge-small", "model.safetensors", 130_000_000)
    expect(artifact.format).toBe("safetensors")
  })
})

// ── RAG ──

describe("RAG assets", () => {
  test("EMBEDDING_MODELS has well-known entries", () => {
    expect(RAG.EMBEDDING_MODELS.length).toBeGreaterThan(0)
    const ids = RAG.EMBEDDING_MODELS.map(m => m.id)
    expect(ids).toContain("bge-small-en")
    expect(ids).toContain("nomic-embed")
    expect(ids).toContain("all-minilm-l6")
  })

  test("RERANKER_MODELS has well-known entries", () => {
    expect(RAG.RERANKER_MODELS.length).toBeGreaterThan(0)
    const ids = RAG.RERANKER_MODELS.map(m => m.id)
    expect(ids).toContain("bge-reranker-base")
  })

  test("embeddingAsset builds correct RAGAsset", () => {
    const model = RAG.EMBEDDING_MODELS[0]
    const asset = RAG.embeddingAsset(model)
    expect(asset.type).toBe("embedding")
    expect(asset.repo).toBe(model.repo)
    expect(asset.dimensions).toBe(model.dimensions)
  })

  test("rerankerAsset builds correct RAGAsset", () => {
    const model = RAG.RERANKER_MODELS[0]
    const asset = RAG.rerankerAsset(model)
    expect(asset.type).toBe("reranker")
    expect(asset.repo).toBe(model.repo)
  })

  test("recommendForVRAM filters by available memory", async () => {
    const { embeddings, rerankers } = await RAG.recommendForVRAM(500) // 500MB
    // Should only include small models
    for (const e of embeddings) {
      expect(e.sizeBytes).toBeLessThan(500 * 1024 * 1024 * 1.2)
    }
  })

  test("recommendForVRAM returns all when plenty of VRAM", async () => {
    const { embeddings } = await RAG.recommendForVRAM(100000) // 100GB
    expect(embeddings.length).toBe(RAG.EMBEDDING_MODELS.length)
  })
})
