import Foundation
import MarkBase
import Hummingbird

struct EmbeddingServerApp {
    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("═══════════════════════════════════════════════════════════════════")
        print("  MarkBaseEngine Embedding Server")
        print("═══════════════════════════════════════════════════════════════════")
        print("  Model: \(modelName)")
        print("  Port: \(port)")
        print("  Path: \(modelPath)")
        print("")

        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))")

        let router = Router()

        router.get("/") { _, _ in
            return """
            {
              "server": {
                "name": "MarkBaseEngine Embedding",
                "model": "embeddinggemma-300m",
                "layers": \(layers),
                "hidden_size": \(hiddenSize),
                "output_dim": \(hiddenSize),
                "framework": "Hummingbird 2.x + Metal GPU"
              },
              "endpoints": [
                {"method": "POST", "path": "/v1/embeddings", "summary": "Generate text embeddings"}
              ]
            }
            """
        }

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

        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\":\"invalid request\",\"type\":\"invalid_request_error\",\"code\":400,\"message\":\"missing 'input' field\"}"
            }

            let modelId = (json["model"] as? String) ?? "embeddinggemma-300m"
            let encodingFormat = (json["encoding_format"] as? String) ?? "float"

            let inputs: [String]
            if let str = input as? String { inputs = [str] }
            else if let arr = input as? [String] { inputs = arr }
            else {
                return "{\"error\":\"invalid request\",\"type\":\"invalid_request_error\",\"code\":400,\"message\":\"'input' must be string or array of strings\"}"
            }

            var embeddings: [[String: Any]] = []
            for (i, text) in inputs.enumerated() {
                let t0 = Date()
                let embedding = try embedModel.embed(text: text)
                let duration = Date().timeIntervalSince(t0)

                let embeddingData: [String: Any]
                if encodingFormat == "base64" {
                    let base64 = embedding.withUnsafeBytes { Data($0).base64EncodedString() }
                    embeddingData = ["object": "embedding", "index": i, "embedding": base64, "usage_ms": Int(duration * 1000)]
                } else {
                    embeddingData = ["object": "embedding", "index": i, "embedding": embedding, "usage_ms": Int(duration * 1000)]
                }
                embeddings.append(embeddingData)
            }

            let id = UUID().uuidString
            let ts = Int(Date().timeIntervalSince1970)
            let response: [String: Any] = [
                "id": id, "object": "list", "created": ts, "model": modelId,
                "data": embeddings,
                "usage": ["prompt_tokens": inputs.reduce(0) { $0 + $1.components(separatedBy: .whitespaces).count }, "total_tokens": inputs.reduce(0) { $0 + $1.components(separatedBy: .whitespaces).count }]
            ]

            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)...")
        print("Endpoints:")
        print("  GET  /             - Info")
        print("  GET  /health       - Health check")
        print("  POST /v1/embeddings - Text embeddings")
        print("")

        try await app.run()
    }
}