package json

import (
	"bufio"
	"fmt"
	"main/walk"
)

type JSONWriterState int
const (
	JSONWriterStateBeforeValue JSONWriterState = iota
	JSONWriterStateAfterValue
	JSONWriterStateInArray
	JSONWriterStateInMap
)

func NewJSONWriter(writer *bufio.Writer) *JSONWriter {
	return &JSONWriter {
		path: nil,
		writer: writer,
		state: JSONWriterStateBeforeValue,
	}
}

type JSONWriter struct {
	path []walk.PathSegment
	writer *bufio.Writer
	state JSONWriterState
}

func (writer *JSONWriter) navigateTo(keepLen int, path []walk.PathSegment, state JSONWriterState) {
	for {
		if keepLen > len(writer.path) {
			panic("keepLen > len(writer.path)")
		} else if len(writer.path) == keepLen {
			if len(path) == 0 {
				switch writer.state {
				case JSONWriterStateBeforeValue:
					switch state {
					case JSONWriterStateBeforeValue:
						return
					case JSONWriterStateAfterValue:
						panic("Cannot go from BeforeValue to AfterValue in navigateTo")
					case JSONWriterStateInArray:
						writer.writer.WriteRune('[')
						writer.state = JSONWriterStateInArray
					case JSONWriterStateInMap:
						writer.writer.WriteRune('{')
						writer.state = JSONWriterStateInMap
					}
				case JSONWriterStateAfterValue:
					if state == JSONWriterStateAfterValue {
						return
					} else {
						if keepLen == 0 {
							writer.writer.WriteRune('\n')
							writer.state = JSONWriterStateBeforeValue
						} else {
							writer.writer.WriteRune(',')
							path = writer.path[len(writer.path) - 1:]
							writer.path = writer.path[:len(writer.path) - 1]
							keepLen -= 1
							switch path[0].(type) {
							case string:
								writer.state = JSONWriterStateInMap
							case int:
								writer.state = JSONWriterStateInArray
							}
						}
					}
				case JSONWriterStateInArray:
					if state == JSONWriterStateInArray {
						return
					} else {
						writer.writer.WriteRune(']')
						writer.state = JSONWriterStateAfterValue
					}
				case JSONWriterStateInMap:
					if state == JSONWriterStateInMap {
						return
					} else {
						writer.writer.WriteRune('}')
						writer.state = JSONWriterStateAfterValue
					}
				}
			} else {
				// len(path) > 0
				switch writer.state {
				case JSONWriterStateBeforeValue:
					switch path[0].(type) {
					case string:
						writer.writer.WriteRune('{')
						writer.state = JSONWriterStateInMap
					case int:
						writer.writer.WriteRune('[')
						writer.state = JSONWriterStateInArray
					}
				case JSONWriterStateAfterValue:
					if keepLen == 0 {
						writer.writer.WriteRune('\n')
						writer.state = JSONWriterStateBeforeValue
					} else {
						writer.writer.WriteRune(',')
						path = append(writer.path[len(writer.path) - 1:], path...)
						writer.path = writer.path[:len(writer.path) - 1]
						keepLen -= 1
						switch path[0].(type) {
						case string:
							writer.state = JSONWriterStateInMap
						case int:
							writer.state = JSONWriterStateInArray
						}
					}
				case JSONWriterStateInArray:
					switch path[0].(type) {
					case string:
						writer.writer.WriteRune(']')
						writer.state = JSONWriterStateAfterValue
					case int:
						writer.path = append(writer.path, path[0])
						path = path[1:]
						keepLen += 1
						writer.state = JSONWriterStateBeforeValue
					}
				case JSONWriterStateInMap:
					switch p := path[0].(type) {
					case string:
						fmt.Fprintf(writer.writer, "%q:", p)
						writer.path = append(writer.path, p)
						path = path[1:]
						keepLen += 1
						writer.state = JSONWriterStateBeforeValue
					case int:
						writer.writer.WriteRune('}')
						writer.state = JSONWriterStateAfterValue
					}
				}
			}
		} else {
			switch writer.state {
			case JSONWriterStateBeforeValue:
				panic("Cannot close structures from BeforeValue in navigateTo")
			case JSONWriterStateAfterValue:
				if len(writer.path) == keepLen + 1 {
					if len(path) == 0 {
						switch writer.path[len(writer.path) - 1].(type) {
						case string:
							if state == JSONWriterStateInMap {
								writer.writer.WriteRune(',')
								writer.path = writer.path[:len(writer.path) - 1]
								writer.state = JSONWriterStateInMap
							} else {
								writer.writer.WriteRune('}')
								writer.path = writer.path[:len(writer.path) - 1]
							}
						case int:
							if state == JSONWriterStateInArray {
								writer.writer.WriteRune(',')
								writer.path = writer.path[:len(writer.path) - 1]
								writer.state = JSONWriterStateInArray
							} else {
								writer.writer.WriteRune(']')
								writer.path = writer.path[:len(writer.path) - 1]
							}
						}
					} else {
						switch writer.path[len(writer.path) - 1].(type) {
						case string:
							switch path[0].(type) {
							case string:
								writer.writer.WriteRune(',')
								writer.path = writer.path[:len(writer.path) - 1]
								writer.state = JSONWriterStateInMap
							case int:
								writer.writer.WriteRune('}')
								writer.path = writer.path[:len(writer.path) - 1]
							}
						case int:
							switch path[0].(type) {
							case string:
								writer.writer.WriteRune(']')
								writer.path = writer.path[:len(writer.path) - 1]
							case int:
								writer.writer.WriteRune(',')
								writer.path = writer.path[:len(writer.path) - 1]
								writer.state = JSONWriterStateInArray
							}
						}
					}
				} else {
					switch writer.path[len(writer.path) - 1].(type) {
					case string:
						writer.writer.WriteRune('}')
						writer.path = writer.path[:len(writer.path) - 1]
					case int:
						writer.writer.WriteRune(']')
						writer.path = writer.path[:len(writer.path) - 1]
					}
				}
			case JSONWriterStateInArray:
				writer.writer.WriteRune(']')
				writer.state = JSONWriterStateAfterValue
			case JSONWriterStateInMap:
				writer.writer.WriteRune('}')
				writer.state = JSONWriterStateAfterValue
			}
		}
	}
}

func (writer *JSONWriter) inMapAt(keepLen int, path []walk.PathSegment) bool {
	if len(path) != 0 {
		return false
	}

	if keepLen == len(writer.path) {
		return writer.state == JSONWriterStateInMap
	}

	_, isString := writer.path[0].(string)
	return isString
}

func (writer *JSONWriter) inArrayAt(keepLen int, path []walk.PathSegment) bool {
	if len(path) != 0 {
		return false
	}

	if keepLen == len(writer.path) {
		return writer.state == JSONWriterStateInArray
	}

	_, isInt := writer.path[0].(int)
	return isInt
}

func countPrefix(current []walk.PathSegment, target []walk.PathSegment) int {
	for i, c := range current {
		if i >= len(target) || c != target[i] {
			return i
		}
	}
	return len(current)
}

/*
{
	"a": {
		"b": {
			"c": null
		},
		"d": [],
		"e": {},
		"f": {"g": 5}
	}
}

write(true, true, [], {"a": ...})
	write(true, true, ["a"], {"b": ..., "d": [], "e": {}, "f": ...})
		write(true, false, ["a", "b"], {"c": null})
			write(true, false, ["a", "b", "c"], null)
				navigateTo(["a", "b", "c"], BeforeValue)
			> null
		> }
	> , "d":
		write(false, false, ["a", "d"], [])
		> []
	> , "e":
		write(false, false, ["a", "e"], {})
		> {}
	> , "f":
		write(false, true, ["a", "f"], {"g": 5})
		> {"g":
			write(false, true, ["a", "f", "g"], 5)
			> 5
			path = ["a", "f", "g"]
			state = AfterValue
*/

/*
{
	"a": {}
}

write(true, true, [], {"a": {}})
	write(true, true, ["a"], {})
		if not in map at ["a"]
			navigateTo(["a"], InMap)
*/

/*
{
	"a": {},
	"b": null,
	"c": {}
}

write(true, true, [], {...})
	write(true, false, ["a"], {})
		if in map at ["a"]
			navigateTo(["a"], AfterValue)
		else
			navigateTo(["a"], BeforeValue)
			> {}
> , "b":
	write(false, false, ["b"], null)
	> null
> , "c":
	write(false, true, ["c"], {})
	> {
	path = ["c"]
	state = InMap
*/

func (writer *JSONWriter) write(first bool, last bool, path []walk.PathSegment, value walk.Value) {
	switch value := value.(type) {
	case walk.NullValue:
		if first {
			keepLen := countPrefix(writer.path, path)
			writer.navigateTo(keepLen, path[keepLen:], JSONWriterStateBeforeValue)
		}
		writer.writer.WriteString("null")
		if last {
			writer.path = path
			writer.state = JSONWriterStateAfterValue
		}
	case walk.BoolValue:
		if first {
			keepLen := countPrefix(writer.path, path)
			writer.navigateTo(keepLen, path[keepLen:], JSONWriterStateBeforeValue)
		}
		if value {
			writer.writer.WriteString("true")
		} else {
			writer.writer.WriteString("false")
		}
		if last {
			writer.path = path
			writer.state = JSONWriterStateAfterValue
		}
	case walk.NumberValue:
		if first {
			keepLen := countPrefix(writer.path, path)
			writer.navigateTo(keepLen, path[keepLen:], JSONWriterStateBeforeValue)
		}
		fmt.Fprintf(writer.writer, "%v", value)
		if last {
			writer.path = path
			writer.state = JSONWriterStateAfterValue
		}
	case walk.StringValue:
		if first {
			keepLen := countPrefix(writer.path, path)
			writer.navigateTo(keepLen, path[keepLen:], JSONWriterStateBeforeValue)
		}
		fmt.Fprintf(writer.writer, "%q", value)
		if last {
			writer.path = path
			writer.state = JSONWriterStateAfterValue
		}
	case walk.ArrayValue:
		if len(value) == 0 {
			if first {
				if last {
					keepLen := countPrefix(writer.path, path)
					if !writer.inArrayAt(keepLen, path[keepLen:]) {
						writer.navigateTo(keepLen, path[keepLen:], JSONWriterStateInArray)
					}
				} else {
					keepLen := countPrefix(writer.path, path)
					if writer.inArrayAt(keepLen, path[keepLen:]) {
						writer.navigateTo(keepLen, path[keepLen:], JSONWriterStateAfterValue)
					} else {
						writer.navigateTo(keepLen, path[keepLen:], JSONWriterStateBeforeValue)
						writer.writer.WriteString("[]")
					}
				}
			} else {
				if last {
					writer.writer.WriteRune('[')
					writer.path = path
					writer.state = JSONWriterStateInArray
				} else {
					writer.writer.WriteString("[]")
				}
			}
		} else {
			if !first {
				writer.writer.WriteRune('[')
			}
			for i, el := range value {
				if i != 0 {
					writer.writer.WriteRune(',')
				}
				writer.write(first && i == 0, last && i == len(value) - 1, append(path, el.Index), el.Value)
			}
			if !last {
				writer.writer.WriteRune(']')
			}
		}
	case walk.MapValue:
		if len(value) == 0 {
			if first {
				if last {
					keepLen := countPrefix(writer.path, path)
					if !writer.inMapAt(keepLen, path[keepLen:]) {
						writer.navigateTo(keepLen, path[keepLen:], JSONWriterStateInMap)
					}
				} else {
					keepLen := countPrefix(writer.path, path)
					if writer.inMapAt(keepLen, path[keepLen:]) {
						writer.navigateTo(keepLen, path[keepLen:], JSONWriterStateAfterValue)
					} else {
						writer.navigateTo(keepLen, path[keepLen:], JSONWriterStateBeforeValue)
						writer.writer.WriteString("{}")
					}
				}
			} else {
				if last {
					writer.writer.WriteRune('{')
					writer.path = path
					writer.state = JSONWriterStateInMap
				} else {
					writer.writer.WriteString("{}")
				}
			}
		} else {
			if !first {
				writer.writer.WriteRune('{')
			}
			for i, el := range value {
				if i != 0 {
					writer.writer.WriteRune(',')
				}
				if i != 0 || !first {
					fmt.Fprintf(writer.writer, "%q:", el.Key)
				}
				writer.write(first && i == 0, last && i == len(value) - 1, append(path, el.Key), el.Value)
			}
			if !last {
				writer.writer.WriteRune('}')
			}
		}
	}
}

func (writer *JSONWriter) Write(value walk.Value) error {
	writer.write(true, true, nil, value)
	return nil
}

func (writer *JSONWriter) AssertDone() {
	switch writer.state {
	case JSONWriterStateInArray:
		writer.writer.WriteRune(']')
	case JSONWriterStateInMap:
		writer.writer.WriteRune('}')
	}
	for i := len(writer.path) - 1; i >= 0; i -= 1 {
		switch writer.path[i].(type) {
		case int:
			writer.writer.WriteRune(']')
		case string:
			writer.writer.WriteRune('}')
		default:
			panic("Invalid path segment type")
		}
	}
	writer.writer.Flush()
}