#include "seobeo/seobeo.h"
#include <string.h>
#include <time.h>

static _Atomic uint32_t counter = 0;

static void build_response(char *out, size_t size, const char *id, const char *result) {
  snprintf(out, size, "{\"jsonrpc\":\"2.0\",\"id\":%s,\"result\":%s}", id, result);
}

static void build_error(char *out, size_t size, const char *id, int code, const char *msg) {
  snprintf(out, size,
    "{\"jsonrpc\":\"2.0\",\"id\":%s,\"error\":{\"code\":%d,\"message\":\"%s\"}}",
    id, code, msg);
}

static void handle_initialize(char *out, size_t size, const char *id) {
  const char *result = "{"
    "\"protocolVersion\":\"2024-11-05\","
    "\"capabilities\":{\"tools\":{}},"
    "\"serverInfo\":{\"name\":\"simple-mcp\",\"version\":\"1.0.0\"}"
  "}";
  build_response(out, size, id, result);
}

static void handle_tools_list(char *out, size_t size, const char *id) {
  const char *result = "{"
    "\"tools\":["
      "{"
        "\"name\":\"get_time\","
        "\"description\":\"Get current Unix timestamp\","
        "\"inputSchema\":{\"type\":\"object\",\"properties\":{}}"
      "},"
      "{"
        "\"name\":\"get_counter\","
        "\"description\":\"Get and increment a counter\","
        "\"inputSchema\":{\"type\":\"object\",\"properties\":{}}"
      "}"
    "]"
  "}";
  build_response(out, size, id, result);
}

static void handle_tools_call(char *out, size_t size, const char *id, Dowa_JSON_Entry *json) {
  Dowa_JSON_Value *params_val = Dowa_JSON_Get(json, "params");
  if (!params_val || params_val->type != DOWA_JSON_OBJECT)
  {
    build_error(out, size, id, -32602, "Missing params");
    return;
  }

  Dowa_JSON_Entry *params = (Dowa_JSON_Entry *)params_val->object_val;
  char *tool_name = Dowa_JSON_Get_String(params, "name");

  if (!tool_name) {
    build_error(out, size, id, -32602, "Missing tool name");
    return;
  }

  char result[256];
  if (strcmp(tool_name, "get_time") == 0) {
    time_t now = time(NULL);
    snprintf(result, sizeof(result),
      "{\"content\":[{\"type\":\"text\",\"text\":\"Current timestamp: %ld\"}]}",
      (long)now);
  } else if (strcmp(tool_name, "get_counter") == 0) {
    uint32_t val = ++counter;
    snprintf(result, sizeof(result),
      "{\"content\":[{\"type\":\"text\",\"text\":\"Counter value: %u\"}]}",
      val);
  } else {
    build_error(out, size, id, -32601, "Unknown tool");
    return;
  }

  build_response(out, size, id, result);
}

Seobeo_Request_Entry* HandleMCP(Seobeo_Request_Entry *req, Dowa_Arena *arena)
{
  Seobeo_Request_Entry *resp = NULL;

  void *body_kv = Dowa_HashMap_Get_Ptr(req, "Body");
  if (!body_kv)
  {
    char *err = "{\"jsonrpc\":\"2.0\",\"id\":null,\"error\":{\"code\":-32700,\"message\":\"No body\"}}";
    Dowa_HashMap_Push_Arena(resp, "status", "400", arena);
    Dowa_HashMap_Push_Arena(resp, "content-type", "application/json", arena);
    Dowa_HashMap_Push_Arena(resp, "body", err, arena);
    return resp;
  }

  const char *body = ((Seobeo_Request_Entry*)body_kv)->value;
  int32 body_len = strlen(body);

  Dowa_JSON_Value parsed = Dowa_JSON_Parse(body, body_len, arena);
  if (parsed.type != DOWA_JSON_OBJECT || !parsed.object_val)
  {
    char *err = "{\"jsonrpc\":\"2.0\",\"id\":null,\"error\":{\"code\":-32700,\"message\":\"Parse error\"}}";
    Dowa_HashMap_Push_Arena(resp, "status", "400", arena);
    Dowa_HashMap_Push_Arena(resp, "content-type", "application/json", arena);
    Dowa_HashMap_Push_Arena(resp, "body", err, arena);
    return resp;
  }

  Dowa_JSON_Entry *json = (Dowa_JSON_Entry *)parsed.object_val;
  char *method = Dowa_JSON_Get_String(json, "method");
  Dowa_JSON_Value *id_val = Dowa_JSON_Get(json, "id");
  char id_buf[32] = "null";
  if (id_val)
  {
    if (id_val->type == DOWA_JSON_NUMBER)
      snprintf(id_buf, sizeof(id_buf), "%.0f", id_val->num_val);
    else if (id_val->type == DOWA_JSON_STRING)
      snprintf(id_buf, sizeof(id_buf), "\"%s\"", id_val->str_val);
  }

  char *response = Dowa_Arena_Allocate(arena, 2048);

  if (!method)
    build_error(response, 2048, id_buf, -32600, "Missing method");
  else if (strcmp(method, "initialize") == 0)
    handle_initialize(response, 2048, id_buf);
  else if (strcmp(method, "tools/list") == 0)
    handle_tools_list(response, 2048, id_buf);
  else if (strcmp(method, "tools/call") == 0)
    handle_tools_call(response, 2048, id_buf, json);
  else if (strcmp(method, "notifications/initialized") == 0)
    snprintf(response, 2048, "{\"jsonrpc\":\"2.0\",\"id\":%s,\"result\":{}}", id_buf);
  else
    build_error(response, 2048, id_buf, -32601, "Method not found");

  Dowa_HashMap_Push_Arena(resp, "status", "200", arena);
  Dowa_HashMap_Push_Arena(resp, "content-type", "application/json", arena);
  Dowa_HashMap_Push_Arena(resp, "body", response, arena);
  return resp;
}

int main(void) {
  Seobeo_Router_Init();
  Seobeo_Router_Register("POST", "/mcp", HandleMCP);

  Seobeo_Log(SEOBEO_INFO, "MCP server running on http://localhost:8080/mcp\n");
  Seobeo_Web_Server_Start(NULL, "8080", SEOBEO_MODE_FORK, 0);
  return 0;
}