package bjc.utils.misc;

import java.util.ArrayList;
import java.util.List;

/**
 * SmartJsonFormatter
 *
 * Pretty-prints JSON while keeping short/simple objects and arrays on a single line.
 *
 * Heuristics:
 * - Empty containers print inline: {} and [].
 * - Containers whose elements/values are all primitives (string, number, true, false, null)
 *   and whose inline form fits within the remaining line budget are printed inline.
 * - Otherwise, containers are expanded. For arrays of primitives that don't fit on one line,
 *   elements are packed into multiple lines with wrapping up to maxWidth.
 *
 * Configuration knobs:
 * - maxWidth: target maximum characters per line (default 100).
 * - indentSize: spaces per indent level (default 2).
 * - inlineMaxElements: hard cap of elements/fields allowed for inline (default 8).
 *
 * Usage:
 *   String pretty = SmartJsonFormatter.format(jsonString);
 *   // or with custom settings:
 *   String pretty = new SmartJsonFormatter(100, 2, 8).formatSmart(jsonString);
 *
 * Notes:
 * - Strict JSON only (no comments, no trailing commas).
 * - Returns the original text on parse error when using format(); throws otherwise via formatSmart().
 */
public class SmartJSONFormatter {

	private final int maxWidth;
	private final int indentSize;
	private final int inlineMaxElements;

	public SmartJSONFormatter() {
		this(100, 2, 8);
	}

	public SmartJSONFormatter(int maxWidth, int indentSize, int inlineMaxElements) {
		this.maxWidth = Math.max(40, maxWidth);
		this.indentSize = Math.max(0, indentSize);
		this.inlineMaxElements = Math.max(1, inlineMaxElements);
	}

	/** Safe helper: returns original JSON if parsing/formatting fails. */
	public static String format(String json) {
		try {
			return new SmartJSONFormatter().formatSmart(json);
		} catch (Exception e) {
			return json; // fail-safe: don't destructively change input if something is off
		}
	}

	/** Strict formatter: throws on invalid JSON. */
	public String formatSmart(String json) {
		Tokenizer tz = new Tokenizer(json);
		Parser p = new Parser(tz);
		Node root = p.parse();
		StringBuilder out = new StringBuilder(json.length() + 64);
		render(root, 0, 0, out);
		return out.toString();
	}

	// ---------- Rendering ----------

	private void render(Node node, int level, int currentLinePos, StringBuilder out) {
		if (node instanceof PrimitiveNode) {
			String text = ((PrimitiveNode) node).raw;
			out.append(text);
			return;
		}
		if (node instanceof ObjectNode) {
			ObjectNode obj = (ObjectNode) node;
			if (obj.fields.isEmpty()) {
				out.append("{}");
				return;
			}
			String inline = tryInlineObject(obj, level, currentLinePos);
			if (inline != null) {
				out.append(inline);
				return;
			}
			out.append("{\n");
			for (int i = 0; i < obj.fields.size(); i++) {
				Pair field = obj.fields.get(i);
				indent(out, level + 1);
				out.append(field.key.raw);
				out.append(": ");
				render(field.value, level + 1, (level + 1) * indentSize + field.key.raw.length() + 2, out);
				if (i < obj.fields.size() - 1)
					out.append(',');
				out.append('\n');
			}
			indent(out, level);
			out.append('}');
			return;
		}
		if (node instanceof ArrayNode) {
			ArrayNode arr = (ArrayNode) node;
			if (arr.values.isEmpty()) {
				out.append("[]");
				return;
			}
			String inline = tryInlineArray(arr, level, currentLinePos);
			if (inline != null) {
				out.append(inline);
				return;
			}
			// NEW: if array is all primitives, pack multiple per line up to maxWidth.
			if (isAllPrimitives(arr)) {
				renderArrayWrappedPrimitives(arr, level, out);
				return;
			}
			// Fallback: one element per line (recursively formatted)
			out.append("[\n");
			for (int i = 0; i < arr.values.size(); i++) {
				Node v = arr.values.get(i);
				indent(out, level + 1);
				render(v, level + 1, (level + 1) * indentSize, out);
				if (i < arr.values.size() - 1)
					out.append(',');
				out.append('\n');
			}
			indent(out, level);
			out.append(']');
		}
	}

	/**
	 * Pack primitive array elements into rows, wrapping before exceeding maxWidth.
	 */
	private void renderArrayWrappedPrimitives(ArrayNode arr, int level, StringBuilder out) {
		out.append("[\n");
		int indentLen = (level + 1) * indentSize;
		int linePos = indentLen;
		indent(out, level + 1);
		for (int i = 0; i < arr.values.size(); i++) {
			String s = ((PrimitiveNode) arr.values.get(i)).raw;
			String sep = (i < arr.values.size() - 1) ? ", " : "";
			int tokenLen = s.length() + sep.length();
			// If this token doesn't fit and we already have something on the line, wrap.
			if (linePos > indentLen && linePos + tokenLen > maxWidth) {
				out.append('\n');
				indent(out, level + 1);
				linePos = indentLen;
			}
			out.append(s);
			if (!sep.isEmpty())
				out.append(sep);
			linePos += tokenLen;
		}
		out.append('\n');
		indent(out, level);
		out.append(']');
	}

	private boolean isAllPrimitives(ArrayNode arr) {
		for (Node v : arr.values) {
			if (!(v instanceof PrimitiveNode))
				return false;
		}
		return true;
	}

	private String tryInlineObject(ObjectNode obj, int level, int currentLinePos) {
		if (obj.fields.size() > inlineMaxElements)
			return null;
		for (Pair f : obj.fields) {
			if (!(f.value instanceof PrimitiveNode))
				return null;
		}
		String inline = buildInlineObject(obj);
		int budget = maxWidth - currentLinePos;
		return inline.length() <= Math.max(0, budget) ? inline : null;
	}

	private String buildInlineObject(ObjectNode obj) {
		StringBuilder sb = new StringBuilder(16 * obj.fields.size() + 2);
		sb.append('{');
		for (int i = 0; i < obj.fields.size(); i++) {
			Pair f = obj.fields.get(i);
			sb.append(f.key.raw).append(": ").append(((PrimitiveNode) f.value).raw);
			if (i < obj.fields.size() - 1)
				sb.append(", ");
		}
		sb.append('}');
		return sb.toString();
	}

	private String tryInlineArray(ArrayNode arr, int level, int currentLinePos) {
		if (arr.values.size() > inlineMaxElements)
			return null;
		for (Node v : arr.values) {
			if (!(v instanceof PrimitiveNode))
				return null;
		}
		String inline = buildInlineArray(arr);
		int budget = maxWidth - currentLinePos;
		return inline.length() <= Math.max(0, budget) ? inline : null;
	}

	private String buildInlineArray(ArrayNode arr) {
		StringBuilder sb = new StringBuilder(16 * arr.values.size() + 2);
		sb.append('[');
		for (int i = 0; i < arr.values.size(); i++) {
			sb.append(((PrimitiveNode) arr.values.get(i)).raw);
			if (i < arr.values.size() - 1)
				sb.append(", ");
		}
		sb.append(']');
		return sb.toString();
	}

	private void indent(StringBuilder out, int level) {
		int n = level * indentSize;
		for (int i = 0; i < n; i++)
			out.append(' ');
	}

	// ---------- AST Nodes ----------

	private interface Node {
	}

	private static class PrimitiveNode implements Node {
		final String raw; // includes quotes for strings

		PrimitiveNode(String raw) {
			this.raw = raw;
		}
	}

	private static class ObjectNode implements Node {
		final List<Pair> fields = new ArrayList<>();
	}

	private static class ArrayNode implements Node {
		final List<Node> values = new ArrayList<>();
	}

	private static class Pair {
		final PrimitiveNode key; // must be a string token (with quotes)
		final Node value;

		Pair(PrimitiveNode key, Node value) {
			this.key = key;
			this.value = value;
		}
	}

	// ---------- Parser ----------

	private static class Parser {
		private final Tokenizer tz;
		private Token look;

		Parser(Tokenizer tz) {
			this.tz = tz;
			this.look = tz.next();
		}

		Node parse() {
			Node v = parseValue();
			expect(TokenType.EOF);
			return v;
		}

		private Node parseValue() {
			switch (look.type) {
			case LBRACE:
				return parseObject();
			case LBRACKET:
				return parseArray();
			case STRING:
			case NUMBER:
			case TRUE:
			case FALSE:
			case NULL:
				String raw = look.text;
				consume();
				return new PrimitiveNode(raw);
			default:
				throw error("Expected value");
			}
		}

		private ObjectNode parseObject() {
			expect(TokenType.LBRACE);
			ObjectNode obj = new ObjectNode();
			if (look.type == TokenType.RBRACE) {
				consume(); // empty {}
				return obj;
			}
			while (true) {
				if (look.type != TokenType.STRING)
					throw error("Object key must be a string");
				PrimitiveNode key = new PrimitiveNode(look.text);
				consume();
				expect(TokenType.COLON);
				Node value = parseValue();
				obj.fields.add(new Pair(key, value));
				if (look.type == TokenType.COMMA) {
					consume();
				} else if (look.type == TokenType.RBRACE) {
					consume();
					break;
				} else {
					throw error("Expected ',' or '}' in object");
				}
			}
			return obj;
		}

		private ArrayNode parseArray() {
			expect(TokenType.LBRACKET);
			ArrayNode arr = new ArrayNode();
			if (look.type == TokenType.RBRACKET) {
				consume(); // empty []
				return arr;
			}
			while (true) {
				arr.values.add(parseValue());
				if (look.type == TokenType.COMMA) {
					consume();
				} else if (look.type == TokenType.RBRACKET) {
					consume();
					break;
				} else {
					throw error("Expected ',' or ']' in array");
				}
			}
			return arr;
		}

		private void expect(TokenType t) {
			if (look.type != t)
				throw error("Expected " + t + " but found " + look.type);
			consume();
		}

		private void consume() {
			look = tz.next();
		}

		private IllegalArgumentException error(String msg) {
			return new IllegalArgumentException(msg + " at position " + tz.lastPosition());
		}
	}

	// ---------- Tokenizer ----------

	private enum TokenType {
		LBRACE, RBRACE, LBRACKET, RBRACKET, COLON, COMMA, STRING, NUMBER, TRUE, FALSE, NULL, EOF
	}

	private static class Token {
		final TokenType type;
		final String text; // raw lexeme for STRING/NUMBER/TRUE/FALSE/NULL
		final int pos;

		Token(TokenType type, String text, int pos) {
			this.type = type;
			this.text = text;
			this.pos = pos;
		}
	}

	private static class Tokenizer {
		private final String s;
		private int i = 0;
		private int lastPos = 0;

		Tokenizer(String s) {
			this.s = s;
		}

		int lastPosition() {
			return lastPos;
		}

		Token next() {
			skipWs();
			lastPos = i;
			if (i >= s.length())
				return new Token(TokenType.EOF, "", i);
			char c = s.charAt(i);
			switch (c) {
			case '{':
				i++;
				return new Token(TokenType.LBRACE, "{", lastPos);
			case '}':
				i++;
				return new Token(TokenType.RBRACE, "}", lastPos);
			case '[':
				i++;
				return new Token(TokenType.LBRACKET, "[", lastPos);
			case ']':
				i++;
				return new Token(TokenType.RBRACKET, "]", lastPos);
			case ':':
				i++;
				return new Token(TokenType.COLON, ":", lastPos);
			case ',':
				i++;
				return new Token(TokenType.COMMA, ",", lastPos);
			case '"':
				return readString();
			default:
				if (c == '-' || (c >= '0' && c <= '9'))
					return readNumber();
				if (isAlphaStart(c))
					return readKeyword();
				throw new IllegalArgumentException("Unexpected character '" + c + "' at position " + i);
			}
		}

		private void skipWs() {
			while (i < s.length()) {
				char c = s.charAt(i);
				if (c == ' ' || c == '\t' || c == '\n' || c == '\r') {
					i++;
					continue;
				}
				break;
			}
		}

		private Token readString() {
			int start = i;
			StringBuilder sb = new StringBuilder();
			sb.append('"');
			i++; // skip opening quote
			boolean escaped = false;
			while (i < s.length()) {
				char c = s.charAt(i++);
				sb.append(c);
				if (escaped) {
					escaped = false;
				} else if (c == '\\') {
					escaped = true;
				} else if (c == '"') {
					// end of string
					return new Token(TokenType.STRING, sb.toString(), start);
				}
			}
			throw new IllegalArgumentException("Unterminated string starting at " + start);
		}

		private Token readNumber() {
			int start = i;
			if (s.charAt(i) == '-')
				i++;
			if (i >= s.length())
				throw new IllegalArgumentException("Invalid number at end");
			if (s.charAt(i) == '0') {
				i++;
			} else if (isDigit(s.charAt(i))) {
				while (i < s.length() && isDigit(s.charAt(i)))
					i++;
			} else {
				throw new IllegalArgumentException("Invalid number at " + start);
			}
			if (i < s.length() && s.charAt(i) == '.') {
				i++;
				if (i >= s.length() || !isDigit(s.charAt(i)))
					throw new IllegalArgumentException("Invalid fraction at " + start);
				while (i < s.length() && isDigit(s.charAt(i)))
					i++;
			}
			if (i < s.length() && (s.charAt(i) == 'e' || s.charAt(i) == 'E')) {
				i++;
				if (i < s.length() && (s.charAt(i) == '+' || s.charAt(i) == '-'))
					i++;
				if (i >= s.length() || !isDigit(s.charAt(i)))
					throw new IllegalArgumentException("Invalid exponent at " + start);
				while (i < s.length() && isDigit(s.charAt(i)))
					i++;
			}
			return new Token(TokenType.NUMBER, s.substring(start, i), start);
		}

		private Token readKeyword() {
			int start = i;
			while (i < s.length() && isAlpha(s.charAt(i)))
				i++;
			String kw = s.substring(start, i);
			switch (kw) {
			case "true":
				return new Token(TokenType.TRUE, "true", start);
			case "false":
				return new Token(TokenType.FALSE, "false", start);
			case "null":
				return new Token(TokenType.NULL, "null", start);
			default:
				throw new IllegalArgumentException("Unexpected token '" + kw + "' at " + start);
			}
		}

		private boolean isAlphaStart(char c) {
			return (c >= 'a' && c <= 'z'); // JSON literals are lowercase
		}

		private boolean isAlpha(char c) {
			return (c >= 'a' && c <= 'z');
		}

		private boolean isDigit(char c) {
			return c >= '0' && c <= '9';
		}
	}

	// ---------- Simple demo (optional) ----------
	public static void main(String[] args) {
		String sample = "{\"a\":1,\"b\":[1,2,3,4,5,6,7,8,9,10],\"c\":[\"a very long string that will likely force wrapping\",2,3],\"d\":[{\"x\":1},{\"y\":2}],\"e\":[[1,2],[3,4]]}";
		System.out.println(SmartJSONFormatter.format(sample));
	}
}