#include <cstring>
#include <exception>
#include <fstream>
#include <iostream>
#include <sstream>
#include <stdexcept>

#include "json.hpp"

#include <stdarg.h>

using namespace std;

char buffer[512];
char *format(char const *base, ...) {
  va_list lst;
  va_start(lst, base);

  vsnprintf(buffer, sizeof buffer, base, lst);

  va_end(lst);
  return buffer;
}

enum TokenType {
  TOK_OPEN_OBJ = 0,
  TOK_CLOSE_OBJ,
  TOK_OPEN_ARR,
  TOK_CLOSE_ARR,

  TOK_ASSIGNMENT,
  TOK_COMMA,

  TOK_STRING,
  TOK_NUMBER,

  TOK_EOF
};

char const *TName(TokenType tok) {
  switch (tok) {
  case TOK_OPEN_OBJ:
    return "{";
  case TOK_CLOSE_OBJ:
    return "}";
  case TOK_OPEN_ARR:
    return "[";
  case TOK_CLOSE_ARR:
    return "]";
  case TOK_ASSIGNMENT:
    return ":";
  case TOK_COMMA:
    return ",";
  case TOK_STRING:
    return "STRING!!!";
  case TOK_NUMBER:
    return "NUMBER FUUAAAAA";
  case TOK_EOF:
    return "EOF";
  }

  return "How the fuck did we even get here";
}

union TokenValue {
  long double d;
  char *s;
};

struct Token {
  TokenType type;
  TokenValue v;

  char const *Name(void) { return TName(type); }

  void PrettyPrint(void) {
    cout << Name() << ' ';
    if (type == TOK_STRING)
      printf("%s", v.s);
    if (type == TOK_NUMBER)
      cout << v.d;
    cout << endl;
  }
};

struct Tokeniser {
  istream *f = nullptr;

  void Init(istream *stream) { f = stream; }
  Token Next(void);

  Token ParseIdent(string &str, bool &in_string);
} TS;

const std::string WHITESPACE = " \n\r\t\f\v";

std::string ltrim(const std::string &s) {
  size_t start = s.find_first_not_of(WHITESPACE);
  return (start == std::string::npos) ? "" : s.substr(start);
}

std::string rtrim(const std::string &s) {
  size_t end = s.find_last_not_of(WHITESPACE);
  return (end == std::string::npos) ? "" : s.substr(0, end + 1);
}

std::string trim(const std::string &s) { return rtrim(ltrim(s)); }

Token Tokeniser::ParseIdent(string &str, bool &in_string) {
  Token res{.type = TOK_NUMBER};

  str = trim(str);

  for (size_t i = 0; i < str.length(); i++) {
    if (!isdigit(str[i]) && str[i] != '.' && str[i] != '\r' && str[i] != '\n') {
      res.type = TOK_STRING;
      break;
    }
  }

  if (res.type == TOK_NUMBER) {
    res.v.d = stold(str.c_str());
  } else {
    char *amogus = (char *)calloc(1, str.length() + 1);
    if ((str.front() == '"' || str.front() == '\'') &&
        (str.back() == '"' || str.back() == '\'')) {
      memcpy(amogus, str.c_str() + 1, str.length() - 2);
      amogus[str.length() - 1] = 0;
    } else {
      memcpy(amogus, str.c_str(), str.length());
    }
    res.v.s = amogus;
  }

  return res;
}

Token Tokeniser::Next(void) {
  char ch;
  string ident = "";
  bool string_mode = false;
  bool in_string = false;

  while (f->get(ch)) {
    static char const *chars = "{}[]:,";
    static unsigned char chars_len = strlen(chars);

    if (ch == '\r' || ch == '\n') {
      if (string_mode == false)
        continue;
    }

    if (ch == '"' || ch == '\'') {
      in_string = !in_string;
      string_mode = !string_mode;
    }

    if (!in_string)
      for (unsigned char i = 0; i < chars_len; i++) {
        if (ch == chars[i] || ch == ' ') {
          if (string_mode && ident.length() > 0) {
            auto a = ParseIdent(ident, in_string);
            string_mode = false;
            ident = "";
            f->unget();
            return a;
          }

          if (ch == ' ')
            continue;
          return Token{.type = (TokenType)i};
        } else {
          string_mode = true;
        }
      }

    if (string_mode)
      ident.push_back(ch);
  }

  return Token{.type = TOK_EOF};
}

void JSONNode::Free(void) {
  switch (type) {
  case JSON_STRING:
    if (v.s)
      delete v.s;
    break;
  case JSON_LIST:
    if (v.list) {
      for (auto i : *v.list)
        i.Free();
      delete v.list;
    }
    break;
  case JSON_OBJECT:
    if (v.object) {
      for (auto i : *v.object) {
        if (i.name)
          delete i.name;
        i.node.Free();
      }
      delete v.object;
    }
    break;
  default:
    break;
  }
}

long double &JSONNode::Number(void) {
  if (type != JSON_NUMBER)
    throw runtime_error("JSON node is not a number.");
  return v.ld;
}

long double JSONNode::Number(long double def) {
  if (type != JSON_NUMBER)
    return def;
  return v.ld;
}

char *&JSONNode::String(void) {
  if (type != JSON_STRING)
    throw runtime_error("JSON node is not a string.");
  return v.s;
}

char *JSONNode::String(char *def) {
  if (type != JSON_STRING)
    return def;

  return v.s;
}

vector<JSONNode> *&JSONNode::List(void) {
  if (type != JSON_LIST)
    throw runtime_error("JSON node is not a list.");
  return v.list;
}

vector<JSONKeyValue> *&JSONNode::Object(void) {
  if (type != JSON_OBJECT)
    throw runtime_error("JSON node is not a object.");
  return v.object;
}

string const JSONNode::Name(void) {
  switch (type) {
  case JSON_UNKNOWN:
    return "Unknown";
  case JSON_OBJECT:
    return "Object";
  case JSON_LIST:
    return "List";
  case JSON_KEY:
    return "Key";
  case JSON_NUMBER:
    return "Number";
  case JSON_STRING:
    return "String";
  }
  return "Unknown node";
}

JSONNode UNKNOWN{.type = JSON_UNKNOWN};

void JSONKeyValue::PrettyPrint(string indent, bool last) {
  printf("%s", indent.c_str());
  if (last) {
    printf("\\-");
    indent += "  ";
  } else {
    printf("|-");
    indent += "| ";
  }

  printf(" %s\n", name);
  node.PrettyPrint(indent, false);
}

JSONNode &JSONNode::operator[](string key) {
  for (size_t i = 0; i < v.object->size(); i++) {
    JSONKeyValue &kv = v.object->at(i);

    if (strcmp(kv.name, key.c_str()) == 0) {
      return kv.node;
    }
  }

  return UNKNOWN;
}

JSONNode &JSONNode::operator[](size_t index) { return v.list->at(index); }

void JSONNode::PrettyPrint(string indent = "", bool last = true) {
  printf("%s", indent.c_str());
  if (last) {
    printf("\\-");
    indent += "  ";
  } else {
    printf("|-");
    indent += "| ";
  }
  printf(" {%s}", Name().c_str());
  printf("\r\e[40C");

  if (type == JSON_NUMBER)
    cout << v.ld;
  if (type == JSON_STRING)
    cout << v.s;

  printf("\n");

  if (type == JSON_OBJECT) {
    for (size_t i = 0; i < v.object->size(); i++)
      v.object->at(i).PrettyPrint(indent, i == v.list->size() - 1);
  }
  if (type == JSON_LIST) {
    for (size_t i = 0; i < v.list->size(); i++)
      v.list->at(i).PrettyPrint(indent, i == v.list->size() - 1);
  }
}

struct ParserState {
  Token tok;

  Token Next(void) {
    tok = TS.Next();
    return tok;
  }

  bool Accept(TokenType other) {
    if (tok.type == other) {
      Next();
      return true;
    }
    return false;
  }

  bool Expect(TokenType other) {
    bool flag = false;
    if (Accept(other))
      flag = true;
    else {
      throw runtime_error(
          format("Expected token: {}, got {}.", tok.Name(), TName(other)));
    }
    return flag;
  }

  JSONNode ParseObject(void);
  JSONNode ParseArray(void);
  JSONNode ParseValue(void);

  JSONNode Parse(void);
} PS;

JSONNode ParserState::ParseObject(void) {
  JSONNode node = {.type = JSON_OBJECT,
                   .v = {.object = new vector<JSONKeyValue>}};

  if (Accept(TOK_CLOSE_OBJ))
    return node;

  while (1) {
    JSONKeyValue pair;
    if (tok.type != TOK_STRING) {
      throw runtime_error("Expected key.");
    }
    pair.name = tok.v.s;
    Next();
    Expect(TOK_ASSIGNMENT);
    pair.node = ParseValue();
    node.v.object->push_back(pair);
    Accept(TOK_COMMA);
    if (tok.type == TOK_CLOSE_OBJ)
      break;
  }

  Expect(TOK_CLOSE_OBJ);

  return node;
}

JSONNode ParserState::ParseArray(void) {
  JSONNode node = {.type = JSON_LIST, .v = {.list = new vector<JSONNode>}};

  if (Accept(TOK_CLOSE_ARR))
    return node;

  while (1) {
    node.v.list->push_back(ParseValue());
    Accept(TOK_COMMA);
    if (tok.type == TOK_CLOSE_ARR)
      break;
  }

  Expect(TOK_CLOSE_ARR);

  return node;
}

JSONNode ParserState::ParseValue(void) {
  TokenValue value = tok.v;
  if (Accept(TOK_OPEN_OBJ)) {
    JSONNode node = ParseObject();
    return node;
  } else if (Accept(TOK_OPEN_ARR)) {
    JSONNode node = ParseArray();
    return node;
  } else if (Accept(TOK_NUMBER)) {
    JSONNode node = {
        .type = JSON_NUMBER,
        .v = {.ld = value.d},
    };
    return node;
  } else if (Accept(TOK_STRING)) {
    JSONNode node = {
        .type = JSON_STRING,
        .v = {.s = value.s},
    };
    return node;
  } else if (Accept(TOK_EOF)) {
    return JSONNode{
        .type = JSON_UNKNOWN,
    };
  } else {
    throw runtime_error(format("Invalid value: %s", tok.Name()));
  }
}

JSONNode ParserState::Parse(void) {
  JSONNode root{JSON_UNKNOWN, {0}};

  Next();
  if (tok.type == TOK_EOF) {
    throw runtime_error("Empty file");
  }
  root = ParseValue();

  return root;
}

JSONNode JSONParseFile(char *file) {
  ifstream f(file);
  TS.Init(&f);
  return PS.Parse();
}

JSONNode JSONParseString(std::string string) {
  istringstream f(string);
  TS.Init(&f);
  return PS.Parse();
}