package org.luaj.vm2.libs;

import java.util.HashSet;
import java.util.Set;

import org.luaj.vm2.LuaDouble;
import org.luaj.vm2.LuaInteger;
import org.luaj.vm2.LuaString;
import org.luaj.vm2.LuaTable;
import org.luaj.vm2.LuaUserdata;
import org.luaj.vm2.LuaValue;
import org.luaj.vm2.Varargs;

/**
 * Minimal built-in cjson-compatible module implemented in Java.
 * It supports require("cjson"), cjson.encode(...), cjson.decode(...), and cjson.null.
 */
public class CjsonLib extends TwoArgFunction {

	private static final Object NULL_SENTINEL = new Object();

	public LuaValue call(LuaValue modname, LuaValue env) {
		LuaValue pkg = env.get("package");
		if (!pkg.istable()) {
			return env;
		}
		LuaValue preload = pkg.get("preload");
		if (!preload.istable()) {
			return env;
		}
		preload.set("cjson", new loader());
		return env;
	}

	static LuaValue nullValue() {
		return new LuaUserdata(NULL_SENTINEL);
	}

	static boolean isNullSentinel(LuaValue value) {
		return value.isuserdata() && value.touserdata() == NULL_SENTINEL;
	}

	static final class loader extends TwoArgFunction {
		public LuaValue call(LuaValue modname, LuaValue env) {
			LuaTable module = new LuaTable(0, 3);
			module.set("encode", new encode());
			module.set("decode", new decode());
			module.set("null", nullValue());
			return module;
		}
	}

	static final class encode extends OneArgFunction {
		public LuaValue call(LuaValue value) {
			StringBuilder out = new StringBuilder();
			new Encoder().append(value, out, new HashSet());
			return LuaValue.valueOf(out.toString());
		}
	}

	static final class decode extends OneArgFunction {
		public LuaValue call(LuaValue value) {
			Parser parser = new Parser(value.checkjstring());
			LuaValue result = parser.parseValue();
			parser.skipWhitespace();
			if (!parser.eof()) {
				error("trailing garbage");
			}
			return result;
		}
	}

	private static final class Encoder {
		void append(LuaValue value, StringBuilder out, Set seen) {
			if (value.isnil() || isNullSentinel(value)) {
				out.append("null");
			} else if (value.isboolean()) {
				out.append(value.toboolean() ? "true" : "false");
			} else if (value.isnumber()) {
				double d = value.todouble();
				if (Double.isNaN(d) || Double.isInfinite(d)) {
					error("cannot encode NaN or infinite numbers");
				}
				out.append(value.tojstring());
			} else if (value.isstring()) {
				appendString(value.checkjstring(), out);
			} else if (value.istable()) {
				appendTable(value.checktable(), out, seen);
			} else {
				error("unsupported type for json encoding: " + value.typename());
			}
		}

		private void appendTable(LuaTable table, StringBuilder out, Set seen) {
			if (!seen.add(table)) {
				error("circular reference");
			}
			try {
				if (isArray(table)) {
					out.append('[');
					int n = table.length();
					for (int i = 1; i <= n; i++) {
						if (i > 1) {
							out.append(',');
						}
						append(table.get(i), out, seen);
					}
					out.append(']');
				} else {
					out.append('{');
					LuaValue key = LuaValue.NIL;
					boolean first = true;
					while (true) {
						Varargs next = table.next(key);
						key = next.arg1();
						if (key.isnil()) {
							break;
						}
						if (!first) {
							out.append(',');
						}
						first = false;
						appendString(key.checkjstring(), out);
						out.append(':');
						append(next.arg(2), out, seen);
					}
					out.append('}');
				}
			} finally {
				seen.remove(table);
			}
		}

		private boolean isArray(LuaTable table) {
			int n = table.length();
			int count = 0;
			LuaValue key = LuaValue.NIL;
			while (true) {
				Varargs next = table.next(key);
				key = next.arg1();
				if (key.isnil()) {
					break;
				}
				if (!key.isinttype()) {
					return false;
				}
				int index = key.toint();
				if (index < 1 || index > n) {
					return false;
				}
				count++;
			}
			return count == n;
		}

		private void appendString(String value, StringBuilder out) {
			out.append('"');
			for (int i = 0; i < value.length(); i++) {
				char c = value.charAt(i);
				switch (c) {
				case '"':
				case '\\':
					out.append('\\').append(c);
					break;
				case '\b':
					out.append("\\b");
					break;
				case '\f':
					out.append("\\f");
					break;
				case '\n':
					out.append("\\n");
					break;
				case '\r':
					out.append("\\r");
					break;
				case '\t':
					out.append("\\t");
					break;
				default:
					if (c < 0x20) {
						String hex = Integer.toHexString(c);
						out.append("\\u");
						for (int j = hex.length(); j < 4; j++) {
							out.append('0');
						}
						out.append(hex);
					} else {
						out.append(c);
					}
				}
			}
			out.append('"');
		}
	}

	private static final class Parser {
		private final String input;
		private int pos;

		Parser(String input) {
			this.input = input;
		}

		LuaValue parseValue() {
			skipWhitespace();
			if (eof()) {
				error("unexpected end of json");
			}
			char c = input.charAt(pos);
			switch (c) {
			case '{':
				return parseObject();
			case '[':
				return parseArray();
			case '"':
				return LuaValue.valueOf(parseString());
			case 't':
				expect("true");
				return LuaValue.TRUE;
			case 'f':
				expect("false");
				return LuaValue.FALSE;
			case 'n':
				expect("null");
				return nullValue();
			default:
				if (c == '-' || (c >= '0' && c <= '9')) {
					return parseNumber();
				}
				error("invalid json value");
				return LuaValue.NIL;
			}
		}

		private LuaValue parseObject() {
			LuaTable table = new LuaTable();
			pos++;
			skipWhitespace();
			if (consume('}')) {
				return table;
			}
			while (true) {
				skipWhitespace();
				if (eof() || input.charAt(pos) != '"') {
					error("expected string key");
				}
				String key = parseString();
				skipWhitespace();
				if (!consume(':')) {
					error("expected ':'");
				}
				table.set(key, parseValue());
				skipWhitespace();
				if (consume('}')) {
					return table;
				}
				if (!consume(',')) {
					error("expected ',' or '}'");
				}
			}
		}

		private LuaValue parseArray() {
			LuaTable table = new LuaTable();
			int index = 1;
			pos++;
			skipWhitespace();
			if (consume(']')) {
				return table;
			}
			while (true) {
				table.set(index++, parseValue());
				skipWhitespace();
				if (consume(']')) {
					return table;
				}
				if (!consume(',')) {
					error("expected ',' or ']'");
				}
			}
		}

		private LuaValue parseNumber() {
			int start = pos;
			if (input.charAt(pos) == '-') {
				pos++;
			}
			readDigits();
			boolean isFloat = false;
			if (!eof() && input.charAt(pos) == '.') {
				isFloat = true;
				pos++;
				readDigits();
			}
			if (!eof()) {
				char c = input.charAt(pos);
				if (c == 'e' || c == 'E') {
					isFloat = true;
					pos++;
					if (!eof()) {
						char sign = input.charAt(pos);
						if (sign == '+' || sign == '-') {
							pos++;
						}
					}
					readDigits();
				}
			}
			String number = input.substring(start, pos);
			try {
				return isFloat ? LuaValue.valueOf(Double.parseDouble(number)) : LuaValue.valueOf(Long.parseLong(number));
			} catch (NumberFormatException e) {
				error("invalid number");
				return LuaValue.NIL;
			}
		}

		private void readDigits() {
			int start = pos;
			while (!eof()) {
				char c = input.charAt(pos);
				if (c < '0' || c > '9') {
					break;
				}
				pos++;
			}
			if (start == pos) {
				error("invalid number");
			}
		}

		private String parseString() {
			StringBuffer out = new StringBuffer();
			pos++;
			while (!eof()) {
				char c = input.charAt(pos++);
				if (c == '"') {
					return out.toString();
				}
				if (c != '\\') {
					out.append(c);
					continue;
				}
				if (eof()) {
					error("unterminated escape");
				}
				char esc = input.charAt(pos++);
				switch (esc) {
				case '"':
				case '\\':
				case '/':
					out.append(esc);
					break;
				case 'b':
					out.append('\b');
					break;
				case 'f':
					out.append('\f');
					break;
				case 'n':
					out.append('\n');
					break;
				case 'r':
					out.append('\r');
					break;
				case 't':
					out.append('\t');
					break;
				case 'u':
					if (pos + 4 > input.length()) {
						error("invalid unicode escape");
					}
					out.append((char) Integer.parseInt(input.substring(pos, pos + 4), 16));
					pos += 4;
					break;
				default:
					error("invalid escape");
				}
			}
			error("unterminated string");
			return null;
		}

		void skipWhitespace() {
			while (!eof()) {
				char c = input.charAt(pos);
				if (c != ' ' && c != '\n' && c != '\r' && c != '\t') {
					break;
				}
				pos++;
			}
		}

		boolean eof() {
			return pos >= input.length();
		}

		private void expect(String token) {
			if (!input.regionMatches(pos, token, 0, token.length())) {
				error("expected '" + token + "'");
			}
			pos += token.length();
		}

		private boolean consume(char c) {
			if (!eof() && input.charAt(pos) == c) {
				pos++;
				return true;
			}
			return false;
		}
	}
}
