#include "json.h"

#include <stdlib.h>
#include <string.h>

typedef struct {
  const char *text;
  size_t length;
  size_t pos;
  bool error;
} Cursor;

static void skip_whitespace(Cursor *c) {
  while (c->pos < c->length) {
    char ch = c->text[c->pos];
    if (ch == ' ' || ch == '\t' || ch == '\n' || ch == '\r') c->pos++;
    else break;
  }
}

static char peek(Cursor *c) {
  return c->pos < c->length ? c->text[c->pos] : '\0';
}

static JsonValue *json_value_new(JsonType type) {
  JsonValue *v = calloc(1, sizeof(JsonValue));
  v->type = type;
  return v;
}

static JsonValue *parse_value(Cursor *c);

/* Dynamically-growable byte buffer, used to build a parsed string
 * (including \uXXXX escapes re-encoded as UTF-8). */
typedef struct {
  char *data;
  size_t len;
  size_t cap;
} StringBuf;

static void sbuf_init(StringBuf *b) {
  b->cap = 32;
  b->len = 0;
  b->data = malloc(b->cap);
  b->data[0] = '\0';
}

static void sbuf_push(StringBuf *b, char ch) {
  if (b->len + 2 > b->cap) {
    b->cap *= 2;
    b->data = realloc(b->data, b->cap);
  }
  b->data[b->len++] = ch;
  b->data[b->len] = '\0';
}

static void sbuf_push_utf8(StringBuf *b, unsigned int codepoint) {
  if (codepoint <= 0x7F) {
    sbuf_push(b, (char)codepoint);
  } else if (codepoint <= 0x7FF) {
    sbuf_push(b, (char)(0xC0 | (codepoint >> 6)));
    sbuf_push(b, (char)(0x80 | (codepoint & 0x3F)));
  } else {
    sbuf_push(b, (char)(0xE0 | (codepoint >> 12)));
    sbuf_push(b, (char)(0x80 | ((codepoint >> 6) & 0x3F)));
    sbuf_push(b, (char)(0x80 | (codepoint & 0x3F)));
  }
}

static int hex_digit(char ch) {
  if (ch >= '0' && ch <= '9') return ch - '0';
  if (ch >= 'a' && ch <= 'f') return ch - 'a' + 10;
  if (ch >= 'A' && ch <= 'F') return ch - 'A' + 10;
  return -1;
}

static char *parse_string_raw(Cursor *c) {
  if (peek(c) != '"') {
    c->error = true;
    return NULL;
  }
  c->pos++;

  StringBuf buf;
  sbuf_init(&buf);

  while (c->pos < c->length) {
    char ch = c->text[c->pos];
    if (ch == '"') {
      c->pos++;
      return buf.data;
    }
    if (ch == '\\') {
      c->pos++;
      if (c->pos >= c->length) break;
      char esc = c->text[c->pos];
      switch (esc) {
        case '"': sbuf_push(&buf, '"'); c->pos++; break;
        case '\\': sbuf_push(&buf, '\\'); c->pos++; break;
        case '/': sbuf_push(&buf, '/'); c->pos++; break;
        case 'b': sbuf_push(&buf, '\b'); c->pos++; break;
        case 'f': sbuf_push(&buf, '\f'); c->pos++; break;
        case 'n': sbuf_push(&buf, '\n'); c->pos++; break;
        case 'r': sbuf_push(&buf, '\r'); c->pos++; break;
        case 't': sbuf_push(&buf, '\t'); c->pos++; break;
        case 'u': {
          c->pos++;
          if (c->pos + 4 > c->length) {
            c->error = true;
            free(buf.data);
            return NULL;
          }
          unsigned int cp = 0;
          for (int i = 0; i < 4; i++) {
            int d = hex_digit(c->text[c->pos + (size_t)i]);
            if (d < 0) {
              c->error = true;
              free(buf.data);
              return NULL;
            }
            cp = (cp << 4) | (unsigned int)d;
          }
          c->pos += 4;
          sbuf_push_utf8(&buf, cp);
          break;
        }
        default:
          c->error = true;
          free(buf.data);
          return NULL;
      }
    } else {
      sbuf_push(&buf, ch);
      c->pos++;
    }
  }

  c->error = true;
  free(buf.data);
  return NULL;
}

static JsonValue *parse_string(Cursor *c) {
  char *s = parse_string_raw(c);
  if (!s) return NULL;
  JsonValue *v = json_value_new(JSON_STRING);
  v->as.string = s;
  return v;
}

static JsonValue *parse_number(Cursor *c) {
  const char *start = c->text + c->pos;
  char *end = NULL;
  double val = strtod(start, &end);
  if (end == start) {
    c->error = true;
    return NULL;
  }
  c->pos += (size_t)(end - start);
  JsonValue *v = json_value_new(JSON_NUMBER);
  v->as.number = val;
  return v;
}

static bool match_literal(Cursor *c, const char *literal) {
  size_t len = strlen(literal);
  if (c->pos + len > c->length) return false;
  if (strncmp(c->text + c->pos, literal, len) != 0) return false;
  c->pos += len;
  return true;
}

static JsonValue *parse_array(Cursor *c) {
  c->pos++; /* consume '[' */
  JsonValue *v = json_value_new(JSON_ARRAY);
  int cap = 0;

  skip_whitespace(c);
  if (peek(c) == ']') {
    c->pos++;
    return v;
  }

  while (true) {
    skip_whitespace(c);
    JsonValue *item = parse_value(c);
    if (!item) {
      json_free(v);
      return NULL;
    }

    if (v->as.array.count >= cap) {
      cap = cap == 0 ? 4 : cap * 2;
      v->as.array.items = realloc(v->as.array.items, (size_t)cap * sizeof(JsonValue *));
    }
    v->as.array.items[v->as.array.count++] = item;

    skip_whitespace(c);
    char ch = peek(c);
    if (ch == ',') {
      c->pos++;
      continue;
    }
    if (ch == ']') {
      c->pos++;
      break;
    }
    c->error = true;
    json_free(v);
    return NULL;
  }
  return v;
}

static JsonValue *parse_object(Cursor *c) {
  c->pos++; /* consume '{' */
  JsonValue *v = json_value_new(JSON_OBJECT);
  int cap = 0;

  skip_whitespace(c);
  if (peek(c) == '}') {
    c->pos++;
    return v;
  }

  while (true) {
    skip_whitespace(c);
    if (peek(c) != '"') {
      c->error = true;
      json_free(v);
      return NULL;
    }
    char *key = parse_string_raw(c);
    if (!key) {
      json_free(v);
      return NULL;
    }

    skip_whitespace(c);
    if (peek(c) != ':') {
      c->error = true;
      free(key);
      json_free(v);
      return NULL;
    }
    c->pos++;
    skip_whitespace(c);

    JsonValue *val = parse_value(c);
    if (!val) {
      free(key);
      json_free(v);
      return NULL;
    }

    if (v->as.object.count >= cap) {
      cap = cap == 0 ? 4 : cap * 2;
      v->as.object.keys = realloc(v->as.object.keys, (size_t)cap * sizeof(char *));
      v->as.object.values = realloc(v->as.object.values, (size_t)cap * sizeof(JsonValue *));
    }
    v->as.object.keys[v->as.object.count] = key;
    v->as.object.values[v->as.object.count] = val;
    v->as.object.count++;

    skip_whitespace(c);
    char ch = peek(c);
    if (ch == ',') {
      c->pos++;
      continue;
    }
    if (ch == '}') {
      c->pos++;
      break;
    }
    c->error = true;
    json_free(v);
    return NULL;
  }
  return v;
}

static JsonValue *parse_value(Cursor *c) {
  skip_whitespace(c);
  char ch = peek(c);
  if (ch == '{') return parse_object(c);
  if (ch == '[') return parse_array(c);
  if (ch == '"') return parse_string(c);
  if (ch == '-' || (ch >= '0' && ch <= '9')) return parse_number(c);
  if (match_literal(c, "true")) {
    JsonValue *v = json_value_new(JSON_BOOL);
    v->as.boolean = true;
    return v;
  }
  if (match_literal(c, "false")) {
    JsonValue *v = json_value_new(JSON_BOOL);
    v->as.boolean = false;
    return v;
  }
  if (match_literal(c, "null")) return json_value_new(JSON_NULL);
  c->error = true;
  return NULL;
}

JsonValue *json_parse(const char *text, size_t length) {
  Cursor c = {text, length, 0, false};
  JsonValue *root = parse_value(&c);
  if (!root) return NULL;

  skip_whitespace(&c);
  if (c.pos != c.length) {
    json_free(root);
    return NULL;
  }
  return root;
}

void json_free(JsonValue *value) {
  if (!value) return;
  switch (value->type) {
    case JSON_STRING:
      free(value->as.string);
      break;
    case JSON_ARRAY:
      for (int i = 0; i < value->as.array.count; i++) json_free(value->as.array.items[i]);
      free(value->as.array.items);
      break;
    case JSON_OBJECT:
      for (int i = 0; i < value->as.object.count; i++) {
        free(value->as.object.keys[i]);
        json_free(value->as.object.values[i]);
      }
      free(value->as.object.keys);
      free(value->as.object.values);
      break;
    default:
      break;
  }
  free(value);
}

const JsonValue *json_object_get(const JsonValue *value, const char *key) {
  if (!value || value->type != JSON_OBJECT) return NULL;
  for (int i = 0; i < value->as.object.count; i++) {
    if (strcmp(value->as.object.keys[i], key) == 0) return value->as.object.values[i];
  }
  return NULL;
}

int json_array_count(const JsonValue *value) {
  if (!value || value->type != JSON_ARRAY) return 0;
  return value->as.array.count;
}

const JsonValue *json_array_get(const JsonValue *value, int index) {
  if (!value || value->type != JSON_ARRAY) return NULL;
  if (index < 0 || index >= value->as.array.count) return NULL;
  return value->as.array.items[index];
}

double json_as_number(const JsonValue *value) {
  if (!value || value->type != JSON_NUMBER) return 0.0;
  return value->as.number;
}

const char *json_as_string(const JsonValue *value) {
  if (!value || value->type != JSON_STRING) return "";
  return value->as.string;
}
