import Foundation

enum InputKind: String, Codable, Sendable {
    case document
    case query
}

enum ModelSize: String, Codable, Sendable, CaseIterable {
    case s06 = "0.6B"
    case s4 = "4B"
    case s8 = "8B"

    /// Official Qwen org F16 weights. mlx-community publishes only quantized
    /// variants for these sizes; the F16 repos are the canonical correctness
    /// oracle. mlx-swift-lm's `Qwen3Model.sanitize` consumes them as-is.
    var repoId: String {
        switch self {
        case .s06: return "Qwen/Qwen3-Embedding-0.6B"
        case .s4:  return "Qwen/Qwen3-Embedding-4B"
        case .s8:  return "Qwen/Qwen3-Embedding-8B"
        }
    }

    var label: String {
        switch self {
        case .s06: return "qwen3-embedding-0.6B"
        case .s4:  return "qwen3-embedding-4B"
        case .s8:  return "qwen3-embedding-8B"
        }
    }
}

struct EmbedRequest: Codable, Sendable {
    let texts: [String]
    let input_type: InputKind
    let size: ModelSize?
}

struct EmbedResponse: Codable, Sendable {
    let vectors: [[Float]]
    let dim: Int
    let model: String
}