/** Local model orchestration — shared types */

import { Schema } from "effect"
import { optionalOmitUndefined } from "@/util/schema"

// ── GPU ──

export const GPUInfo = Schema.Struct({
  index: Schema.Number,
  name: Schema.String,
  vramTotalMB: Schema.Number,
  vramUsedMB: Schema.Number,
  vramFreeMB: Schema.Number,
  computeCapability: optionalOmitUndefined(Schema.String),
  driverVersion: optionalOmitUndefined(Schema.String),
  cudaVersion: optionalOmitUndefined(Schema.String),
})
export type GPUInfo = Schema.Schema.Type<typeof GPUInfo>

export const GPUProfile = Schema.Struct({
  gpus: Schema.Array(GPUInfo),
  totalVRAMMB: Schema.Number,
  availableVRAMMB: Schema.Number,
  cudaAvailable: Schema.Boolean,
})
export type GPUProfile = Schema.Schema.Type<typeof GPUProfile>

// ── Model Artifact ──

export const Quantization = Schema.Literals([
  "q2_k", "q3_k_s", "q3_k_m", "q3_k_l",
  "q4_0", "q4_1", "q4_k_s", "q4_k_m",
  "q5_0", "q5_1", "q5_k_s", "q5_k_m",
  "q6_k", "q8_0", "fp16", "bf16", "fp32",
])
export type Quantization = Schema.Schema.Type<typeof Quantization>

export const ModelFormat = Schema.Literals(["gguf", "safetensors", "pytorch", "onnx", "awq", "gptq"])
export type ModelFormat = Schema.Schema.Type<typeof ModelFormat>

export const ModelArtifact = Schema.Struct({
  id: Schema.String,
  repo: Schema.String,                     // e.g. "meta-llama/Llama-3.1-8B-Instruct"
  filename: Schema.String,                 // e.g. "model-q4_k_m.gguf"
  format: ModelFormat,
  quantization: Quantization,
  sizeBytes: Schema.Number,
  sha256: optionalOmitUndefined(Schema.String),
})
export type ModelArtifact = Schema.Schema.Type<typeof ModelArtifact>

// ── Model Instance (running process) ──

export const ModelStatus = Schema.Literals(["stopped", "starting", "running", "error", "downloading"])
export type ModelStatus = Schema.Schema.Type<typeof ModelStatus>

export const ModelInstance = Schema.Struct({
  id: Schema.String,
  artifact: ModelArtifact,
  status: ModelStatus,
  pid: optionalOmitUndefined(Schema.Number),
  port: optionalOmitUndefined(Schema.Number),
  gpuIndex: optionalOmitUndefined(Schema.Number),
  vramUsageMB: optionalOmitUndefined(Schema.Number),
  endpoint: optionalOmitUndefined(Schema.String),
  error: optionalOmitUndefined(Schema.String),
})
export type ModelInstance = Schema.Schema.Type<typeof ModelInstance>

// ── Backend (Ollama, LM Studio, vLLM, etc.) ──

export const BackendType = Schema.Literals(["ollama", "lmstudio", "vllm", "llama.cpp", "text-generation-webui"])
export type BackendType = Schema.Schema.Type<typeof BackendType>

export const BackendStatus = Schema.Struct({
  type: BackendType,
  available: Schema.Boolean,
  endpoint: Schema.String,
  version: optionalOmitUndefined(Schema.String),
  loadedModels: Schema.Array(Schema.String),
})
export type BackendStatus = Schema.Schema.Type<typeof BackendStatus>

// ── RAG / Vector DB ──

export const RAGAssetType = Schema.Literals(["embedding", "reranker", "vectordb"])
export type RAGAssetType = Schema.Schema.Type<typeof RAGAssetType>

export const RAGAsset = Schema.Struct({
  id: Schema.String,
  name: Schema.String,
  type: RAGAssetType,
  repo: Schema.String,
  format: ModelFormat,
  sizeBytes: Schema.Number,
  dimensions: optionalOmitUndefined(Schema.Number),
})
export type RAGAsset = Schema.Schema.Type<typeof RAGAsset>

// ── Download Progress ──

export const DownloadState = Schema.Struct({
  id: Schema.String,
  bytesDownloaded: Schema.Number,
  bytesTotal: Schema.Number,
  speedMBps: Schema.Finite,
  eta: optionalOmitUndefined(Schema.Number),
  done: Schema.Boolean,
  error: optionalOmitUndefined(Schema.String),
})
export type DownloadState = Schema.Schema.Type<typeof DownloadState>

// ── HuggingFace Hub ──

export const HFModelSearch = Schema.Struct({
  id: Schema.String,
  name: Schema.String,
  author: Schema.String,
  downloads: Schema.Number,
  likes: Schema.Number,
  tags: Schema.Array(Schema.String),
  pipelineTag: optionalOmitUndefined(Schema.String),
  libraryName: optionalOmitUndefined(Schema.String),
})
export type HFModelSearch = Schema.Schema.Type<typeof HFModelSearch>

export const HFSibling = Schema.Struct({
  rfilename: Schema.String,
})
export type HFSibling = Schema.Schema.Type<typeof HFSibling>

export const HFModelInfo = Schema.Struct({
  id: Schema.String,
  modelId: Schema.String,
  sha: Schema.String,
  siblings: Schema.Array(HFSibling),
  tags: Schema.Array(Schema.String),
  downloads: Schema.Number,
  likes: Schema.Number,
  private: Schema.Boolean,
})
export type HFModelInfo = Schema.Schema.Type<typeof HFModelInfo>
