markbaseengine/Sources/MarkBaseServer/APIRouter.swift

import Foundation
import MarkBase

// ─────────────────────────────────────────────────────────────
// Complete API Router Implementation
// ─────────────────────────────────────────────────────────────

/// API endpoint handler
public final class APIRouter: @unchecked Sendable {
    private let modelManager: ModelManager
    private let metricsCollector: MetricsCollector
    private let concurrencyController: DynamicConcurrencyController
    private let requestQueue: RequestQueue
    
    public init(
        modelManager: ModelManager,
        metricsCollector: MetricsCollector = .shared,
        concurrencyController: DynamicConcurrencyController
    ) {
        self.modelManager = modelManager
        self.metricsCollector = metricsCollector
        self.concurrencyController = concurrencyController
        self.requestQueue = RequestQueue(controller: concurrencyController)
    }
    
    // ─────────────────────────────────────────────────────────────
    // Health & Info Endpoints
    // ─────────────────────────────────────────────────────────────
    
    /// GET /health
    public func handleHealth() async -> [String: Any] {
        let currentModel = await modelManager.getCurrentModel()
        
        return [
            "status": "healthy",
            "model": currentModel?.name ?? "none",
            "model_id": currentModel?.id ?? "",
            "loaded": currentModel?.loaded ?? false,
            "version": "1.0.0"
        ]
    }
    
    /// GET /v1/models
    public func handleListModels() async -> [String: Any] {
        let models = await modelManager.listModels()
        
        return [
            "object": "list",
            "data": models.map { model in
                [
                    "id": model.id,
                    "object": "model",
                    "created": Int(Date().timeIntervalSince1970),
                    "owned_by": "markbase",
                    "loaded": model.loaded
                ]
            }
        ]
    }
    
    // ─────────────────────────────────────────────────────────────
    // Model Management Endpoints
    // ─────────────────────────────────────────────────────────────
    
    /// POST /v1/models/load
    public func handleLoadModel(modelId: String) async throws -> [String: Any] {
        try await modelManager.loadModel(id: modelId)
        
        return [
            "status": "success",
            "message": "Model loaded: \(modelId)"
        ]
    }
    
    /// POST /v1/models/unload
    public func handleUnloadModel() async -> [String: Any] {
        await modelManager.unloadModel()
        
        return [
            "status": "success",
            "message": "Model unloaded"
        ]
    }
    
    /// POST /v1/models/switch
    public func handleSwitchModel(modelId: String) async throws -> [String: Any] {
        try await modelManager.switchModel(to: modelId)
        
        return [
            "status": "success",
            "message": "Model switched to: \(modelId)"
        ]
    }
    
    // ─────────────────────────────────────────────────────────────
    // Chat Completions Endpoint
    // ─────────────────────────────────────────────────────────────
    
    /// POST /v1/chat/completions
    public func handleChatCompletion(
        messages: [ChatMessage],
        config: GenerationConfig
    ) async throws -> ChatCompletionResponse {
        try await requestQueue.execute { [weak self] in
            guard let self = self else {
                throw ModelError.noModelLoaded
            }
            
            let generator = try await self.modelManager.getGenerator()
            let tokenizer = try await self.modelManager.getTokenizer()
            let _ = try await self.modelManager.getModel()
            
            // Build prompt
            let prompt = self.buildChatPrompt(messages: messages, tokenizer: tokenizer)
            
            // Generate
            let startTime = Date()
            let response = try generator.generateComplete(prompt: prompt, config: config)
            let duration = Date().timeIntervalSince(startTime)
            
            // Record metrics
            let tokens = tokenizer.encode(text: response).count
            let resolvedModelId = (await modelManager.getCurrentModel())?.id ?? "unknown"
            self.metricsCollector.recordRequest(duration: duration, tokens: tokens, model: resolvedModelId)
            
            return ChatCompletionResponse(
                id: self.generateId("chatcmpl"),
                object: "chat.completion",
                created: Int(Date().timeIntervalSince1970),
                model: resolvedModelId,
                choices: [
                    Choice(
                        index: 0,
                        message: ChatMessage(role: "assistant", content: response),
                        finish_reason: "stop"
                    )
                ],
                usage: Usage(
                    promptTokens: tokenizer.encode(text: prompt).count,
                    completionTokens: tokens,
                    totalTokens: tokenizer.encode(text: prompt + response).count
                )
            )
        }
    }
    
    // ─────────────────────────────────────────────────────────────
    // Embeddings Endpoint
    // ─────────────────────────────────────────────────────────────
    
    /// POST /v1/embeddings
    public func handleEmbeddings(
        input: EmbeddingsRequest.InputType
    ) async throws -> EmbeddingsResponse {
        try await requestQueue.execute { [weak self] in
            guard let self = self else {
                throw ModelError.noModelLoaded
            }
            
            let tokenizer = try await self.modelManager.getTokenizer()
            
            var embeddings: [EmbeddingData] = []
            var totalTokens = 0
            
            switch input {
            case .string(let text):
                let embedding = try await self.generateEmbedding(text: text)
                let tokens = tokenizer.encode(text: text).count
                totalTokens += tokens
                embeddings.append(EmbeddingData(index: 0, embedding: embedding))
                
            case .strings(let texts):
                for (index, text) in texts.enumerated() {
                    let embedding = try await self.generateEmbedding(text: text)
                    let tokens = tokenizer.encode(text: text).count
                    totalTokens += tokens
                    embeddings.append(EmbeddingData(index: index, embedding: embedding))
                }
                
            case .tokens(let tokens):
                let text = tokenizer.decode(tokens: tokens)
                let embedding = try await self.generateEmbedding(text: text)
                totalTokens += tokens.count
                embeddings.append(EmbeddingData(index: 0, embedding: embedding))
                
            case .tokensList(let tokensList):
                for (index, tokens) in tokensList.enumerated() {
                    let text = tokenizer.decode(tokens: tokens)
                    let embedding = try await self.generateEmbedding(text: text)
                    totalTokens += tokens.count
                    embeddings.append(EmbeddingData(index: index, embedding: embedding))
                }
            }
            
            return EmbeddingsResponse(
                data: embeddings,
                model: (await self.modelManager.getCurrentModel())?.id ?? "unknown",
                usage: EmbeddingUsage(
                    prompt_tokens: totalTokens,
                    total_tokens: totalTokens
                )
            )
        }
    }
    
    // ─────────────────────────────────────────────────────────────
    // Private Helpers
    // ─────────────────────────────────────────────────────────────
    
    private func buildChatPrompt(messages: [ChatMessage], tokenizer: Tokenizer) -> String {
        var prompt = ""
        
        for message in messages {
            let role = message.role == "assistant" ? "model" : message.role
            prompt += "<|turn>\(role)\n\(message.content ?? "")<turn|>\n"
        }
        
        prompt += "<|turn>model\n"
        return prompt
    }
    
    private func generateEmbedding(text: String) async throws -> [Float] {
        let tokenizer = try await modelManager.getTokenizer()
        let model = try await modelManager.getModel()
        
        let tokens = tokenizer.encode(text: text)
        
        var lastHidden: [Float] = []
        for (position, tokenId) in tokens.enumerated() {
            lastHidden = try model.forward(tokenId: tokenId, position: position)
        }
        
        return lastHidden
    }
    
    private func generateId(_ prefix: String) -> String {
        let uuid = UUID().uuidString.replacingOccurrences(of: "-", with: "")
        return "\(prefix)-\(uuid.prefix(29))"
    }
}