import Foundation
import MarkBase
import Hummingbird

@main
struct EmbeddingServerMain {
    static func main() async throws {
        let args = CommandLine.arguments
        let modelName = args.count > 1 ? args[1] : "embeddinggemma-300m"
        let port = args.count > 2 ? Int(args[2]) ?? 8084 : 8084
        
        let modelPath = NSString(string: "~/MarkBaseEngine/models/\(modelName)").expandingTildeInPath

        print("MarkBaseEngine Embedding Server")
        print("  Model: \(modelName)")
        print("  Port: \(port)")
        print("  Path: \(modelPath)")

        let engine = try MarkBaseEngine(autoCompile: true)
        let embedModel = try EmbeddingGemmaModel(modelDir: modelPath, engine: engine)
        let layers = embedModel.config.numHiddenLayers
        let hiddenSize = embedModel.config.hiddenSize

        print("EmbeddingGemma loaded (\(layers) layers, hidden=\(hiddenSize))")

        // Use actor to serialize access to embedModel
        let embedder = Embedder(model: embedModel)

        let router = Router()

        router.get("/") { _, _ in
            return "{\"server\":\"MarkBaseEngine Embedding\",\"model\":\"embeddinggemma-300m\",\"layers\":\(layers),\"hidden_size\":\(hiddenSize)}"
        }

        router.get("/health") { _, _ in
            return "{\"status\":\"healthy\",\"model\":\"embeddinggemma-300m\",\"layers\":\(layers)}"
        }

        router.post("/v1/embeddings") { request, _ in
            let buffer = try await request.body.collect(upTo: .max)
            let data = Data(buffer: buffer)

            guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any],
                  let input = json["input"] else {
                return "{\"error\":\"missing 'input' field\"}"
            }

            let inputs: [String]
            if let str = input as? String { inputs = [str] }
            else if let arr = input as? [String] { inputs = arr }
            else { return "{\"error\":\"'input' must be string or array\"}" }

            var embeddings: [[String: Any]] = []
            for (i, text) in inputs.enumerated() {
                let result = try await embedder.embed(text: text)
                embeddings.append(["object": "embedding", "index": i, "embedding": result.embedding, "usage_ms": result.ms])
            }

            let id = UUID().uuidString
            let ts = Int(Date().timeIntervalSince1970)
            let response: [String: Any] = [
                "id": id, "object": "list", "created": ts, "model": "embeddinggemma-300m",
                "data": embeddings
            ]

            let jsonData = try JSONSerialization.data(withJSONObject: response)
            return String(data: jsonData, encoding: .utf8) ?? "{}"
        }

        let app = Application(
            router: router,
            configuration: .init(address: .hostname("0.0.0.0", port: port))
        )

        print("Server starting on port \(port)...")
        try await app.run()
    }
}

struct EmbeddingResult {
    let embedding: [Float]
    let ms: Int
}

actor Embedder {
    private let model: EmbeddingGemmaModel
    
    init(model: EmbeddingGemmaModel) {
        self.model = model
    }
    
    func embed(text: String) throws -> EmbeddingResult {
        let t0 = Date()
        let embedding = try model.embed(text: text)
        let ms = Int(Date().timeIntervalSince(t0) * 1000)
        return EmbeddingResult(embedding: embedding, ms: ms)
    }
}