import Foundation
import MarkBase
import Hummingbird

struct SimpleServerApp {
    static func main() async throws {
        // 支持命令行參數選擇模型
        let args = CommandLine.arguments
        let modelName = args.count > 1 ? args[1] : "E4B-MarkBase"
        let port = args.count > 2 ? Int(args[2]) ?? 8080 : 8080
        
        let modelPath = NSString(string: "~/MarkBaseEngine/models/\(modelName)").expandingTildeInPath

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

        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelPath)
        let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)

        print("✓ E4B loaded (\(model.numHiddenLayers) layers)")

        let router = Router()

        let layers = model.numHiddenLayers

        @Sendable func helpJSON() -> String {
            return """
            {
              "server": {
                "name": "MarkBaseEngine",
                "version": "1.0.0",
                "model": "E4B (\(layers) layers)",
                "framework": "Hummingbird 2.x + Metal GPU",
                "platform": "Apple Silicon",
                "base_url": "http://localhost:8080"
              },
              "endpoints": [
                {
                  "method": "GET",
                  "path": "/",
                  "summary": "API help and documentation",
                  "content_type": "application/json",
                  "curl": "curl http://localhost:8080/"
                },
                {
                  "method": "GET",
                  "path": "/help",
                  "summary": "API help and documentation (alias)",
                  "content_type": "application/json",
                  "curl": "curl http://localhost:8080/help"
                },
                {
                  "method": "GET",
                  "path": "/health",
                  "summary": "Health check - returns server status and model information",
                  "responses": {
                    "200": {
                      "description": "Server is healthy",
                      "body": "OK"
                    }
                  },
                  "curl": "curl http://localhost:8080/health"
                },
                {
                  "method": "GET",
                  "path": "/v1/models",
                  "summary": "List available models (OpenAI-compatible)",
                  "responses": {
                    "200": {
                      "description": "Model list",
                      "body": {
                        "id": "e4b",
                        "object": "model",
                        "owned_by": "markbase"
                      }
                    }
                  },
                  "curl": "curl http://localhost:8080/v1/models"
                },
                {
                  "method": "POST",
                  "path": "/v1/chat/completions",
                  "summary": "Chat completion (OpenAI-compatible API)",
                  "description": "Generate chat responses using the E4B model. Supports text-only input via messages array.",
                  "request": {
                    "body": {
                      "messages": [
                        {"role": "user", "content": "Hello"}
                      ],
                      "max_tokens": 100,
                      "temperature": 0.7,
                      "top_p": 0.95,
                      "top_k": 40,
                      "stream": false
                    }
                  },
                  "responses": {
                    "200": {
                      "description": "Successful completion",
                      "body": {
                        "id": "chatcmpl-xxx",
                        "object": "chat.completion",
                        "created": 1234567890,
                        "model": "e4b",
                        "choices": [
                          {
                            "index": 0,
                            "message": {
                              "role": "assistant",
                              "content": "Hello! How can I help you?"
                            },
                            "finish_reason": "stop"
                          }
                        ]
                      }
                    },
                    "400": {
                      "description": "Invalid request - missing or malformed messages field"
                    }
                  },
                  "parameters": [
                    {"name": "messages", "type": "array", "required": true, "description": "Array of message objects with 'role' (system/user/assistant/tool) and 'content' (string). Supports 'tool_calls' in assistant messages and 'tool_call_id'/'name' in tool messages."},
                    {"name": "max_tokens", "type": "integer", "required": false, "default": 100, "description": "Maximum number of tokens to generate (1-4096)"},
                    {"name": "temperature", "type": "float", "required": false, "default": 0.7, "description": "Sampling temperature (0.0-2.0)"},
                    {"name": "top_p", "type": "float", "required": false, "description": "Nucleus sampling threshold (0.0-1.0)"},
                    {"name": "top_k", "type": "integer", "required": false, "description": "Top-k sampling count"},
                    {"name": "tools", "type": "array", "required": false, "description": "Array of tool definitions (OpenAI format) for function calling. Each tool has 'type' (function) and 'function' with 'name', 'description', 'parameters'."},
                    {"name": "stream", "type": "boolean", "required": false, "default": false, "description": "Enable streaming response (not yet implemented)"}
                  ],
                  "curl": "curl -X POST http://localhost:8080/v1/chat/completions -H 'Content-Type: application/json' -d '{\"messages\":[{\"role\":\"user\",\"content\":\"Hello\"}],\"max_tokens\":100}'",
                  "examples": [
                    {
                      "title": "Text completion",
                      "curl": "curl -X POST http://localhost:8080/v1/chat/completions -H 'Content-Type: application/json' -d '{\"messages\":[{\"role\":\"user\",\"content\":\"Hello\"}],\"max_tokens\":100}'"
                    },
                    {
                      "title": "Function calling",
                      "curl": "curl -X POST http://localhost:8080/v1/chat/completions -H 'Content-Type: application/json' -d '{\"messages\":[{\"role\":\"user\",\"content\":\"Search for cats\"}],\"tools\":[{\"type\":\"function\",\"function\":{\"name\":\"find_file\",\"description\":\"Search for files\",\"parameters\":{\"type\":\"object\",\"properties\":{\"query\":{\"type\":\"string\"}},\"required\":[\"query\"]}}}],\"max_tokens\":200}'"
                    }
                  ]
                }
              ],
              "schema": {
                "content_type": "application/json",
                "error_format": {
                  "error": {
                    "message": "Error description",
                    "type": "error_type",
                    "code": 400
                  }
                },
                "error_codes": [
                  {"code": 400, "type": "invalid_request_error", "description": "Invalid request parameters"},
                  {"code": 404, "type": "not_found_error", "description": "Resource not found"},
                  {"code": 500, "type": "server_error", "description": "Internal server error"}
                ]
              },
              "documentation": {
                "api_spec": "docs/API_SPEC.md",
                "api_reference": "docs/API.md",
                "deployment": "docs/DEPLOYMENT.md",
                "performance": "docs/PERFORMANCE.md"
              },
              "notes": [
                "All responses are in JSON format",
                "Text generation only (multimodal not yet supported via API)",
                "E4B model with 42 layers, ~4B parameters",
                "For multimodal (vision/audio) support, use the MarkBase Swift library directly",
                "Streaming support is planned but not yet implemented",
                "Function calling uses native Gemma 4 special tokens",
                "Messages can include tool_calls (assistant) and tool responses (tool role) for multi-turn function calling"
              ]
            }
            """
        }

        @Sendable func healthResponse() -> String {
            return "{\"status\":\"healthy\",\"model\":\"e4b\",\"layers\":\(layers)}"
        }

        router.get("/") { _, _ in
            return helpJSON()
        }

        router.get("/help") { _, _ in
            return helpJSON()
        }

        router.get("/health") { _, _ in
            return healthResponse()
        }

        router.get("/v1/models") { _, _ in
            return "{\"id\":\"e4b\",\"object\":\"model\",\"owned_by\":\"markbase\"}"
        }

        router.post("/v1/chat/completions") { 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 messages = json["messages"] as? [[String: Any]] else {
                return "{\"error\":\"invalid request\",\"type\":\"invalid_request_error\",\"code\":400}"
            }

            let maxTokens = (json["max_tokens"] as? Int) ?? 100
            let temperature = Float(json["temperature"] as? Double ?? 0.7)
            let topK = json["top_k"] as? Int
            let topP = (json["top_p"] as? Double).map { Float($0) }
            let tools = json["tools"] as? [[String: Any]]

            let prompt = Gemma4Format.buildChatPrompt(messages: messages, tools: tools)

            let promptTokens = tokenizer.encode(text: prompt)

            let config = GenerationConfig(
                maxTokens: maxTokens,
                temperature: temperature,
                topK: topK,
                topP: topP
            )

            let generatedTokens = try generator.generateTokens(promptTokens: promptTokens, config: config)

            let id = UUID().uuidString
            let ts = Int(Date().timeIntervalSince1970)

            // Check for tool calls (token ID 48 = <|tool_call>)
            if generatedTokens.contains(48) {
                var results: [ToolCallResult] = []
                var i = 0
                while i < generatedTokens.count {
                    if generatedTokens[i] == 48 {
                        i += 1
                        var callTokens: [Int] = []
                        while i < generatedTokens.count && generatedTokens[i] != 49 {
                            callTokens.append(generatedTokens[i])
                            i += 1
                        }
                        if i < generatedTokens.count && generatedTokens[i] == 49 {
                            i += 1
                        }
                        let callText = tokenizer.decode(tokens: callTokens)
                        if let colonIdx = callText.firstIndex(of: ":"),
                           let braceIdx = callText.firstIndex(of: "{"),
                           let endBrace = callText.lastIndex(of: "}") {
                            let name = String(callText[callText.index(after: colonIdx)..<braceIdx]).trimmingCharacters(in: .whitespaces)
                            let rawArgs = String(callText[callText.index(after: braceIdx)..<endBrace])
                            let jsonArgs = Gemma4Format.gemma4ArgsToJSON(rawArgs)
                            results.append(ToolCallResult(
                                id: "call_\(UUID().uuidString.replacingOccurrences(of: "-", with: "").prefix(16))",
                                function: ToolCallFunction(name: name, arguments: jsonArgs)
                            ))
                        }
                    } else {
                        i += 1
                    }
                }

                let encoder = JSONEncoder()
                let callsData = try encoder.encode(results)
                let callsStr = String(data: callsData, encoding: .utf8) ?? "[]"

                return """
                {"id":"chatcmpl-\(id)","object":"chat.completion","created":\(ts),"model":"e4b","choices":[{"index":0,"message":{"role":"assistant","content":null,"tool_calls":\(callsStr)},"finish_reason":"tool_calls"}]}
                """
            } else {
                let response = tokenizer.decode(tokens: generatedTokens)
                let trimmed = response.trimmingCharacters(in: .whitespacesAndNewlines)
                let escaped = trimmed
                    .replacingOccurrences(of: "\\", with: "\\\\")
                    .replacingOccurrences(of: "\"", with: "\\\"")
                    .replacingOccurrences(of: "\n", with: "\\n")
                    .replacingOccurrences(of: "\r", with: "\\r")
                    .replacingOccurrences(of: "\t", with: "\\t")

                return """
                {"id":"chatcmpl-\(id)","object":"chat.completion","created":\(ts),"model":"e4b","choices":[{"index":0,"message":{"role":"assistant","content":"\(escaped)"},"finish_reason":"stop"}]}
                """
            }
        }

        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  /         - API help")
        print("  GET  /help     - API help")
        print("  GET  /health   - Health check")
        print("  GET  /v1/models - Model list")
        print("  POST /v1/chat/completions - Chat completion")
        print("")
        print("Model: \(modelName)")
        if modelName.contains("E4B") {
            print("  ⚠️ E4B is multimodal (Vision + Audio + Text)")
            print("  For text-only, use: 12B-it-MLX-8bit or 31B-it-8bit")
        } else {
            print("  ✓ LLM for text generation")
        }
        print("")

        try await app.run()
    }
}