package walk

import (
	"bufio"
	"math"
	"strings"
	"strconv"
)

type JSONInStructure int
const (
	JSONInRoot JSONInStructure = iota
	JSONInMap
	JSONInArray
	JSONInValueEnd
)

type JSONIn struct {
	path []Atom
	reader *bufio.Reader
	structure []JSONInStructure
}

func NewJSONIn(reader *bufio.Reader) JSONIn {
	return JSONIn {
		path: make([]Atom, 0, 256),
		reader: reader,
		structure: []JSONInStructure{JSONInRoot},
	}
}

func isWhitespace(r rune) bool {
	for _, ws := range " \t\r\n" {
		if r == ws {
			return true
		}
	}
	return false
}

func isNumberRune(r rune) bool {
	return '0' <= r && r <= '9' || r == '.'
}

func (in *JSONIn) popPath() {
	if len(in.path) == 0 {
		panic("Tried to pop from empty path")
	}
	finalAtom := in.path[len(in.path) - 1]
	if finalAtom.Typ != AtomStringTerminal {
		in.path = in.path[:len(in.path) - 1]
		return
	}
	i := len(in.path) - 2
	for {
		if i < 0 {
			panic("Missing string begin in path")
		}
		if in.path[i].Typ == AtomStringTerminal {
			break
		}
		i--
	}
	in.path = in.path[:i]
}

func (in *JSONIn) nextNonWsRune() (rune, error) {
	for {
		r, _, err := in.reader.ReadRune()
		if err != nil {
			return 0, err
		}
		if !isWhitespace(r) {
			return r, nil
		}
	}
}

func (in *JSONIn) requireString(criteria string) {
	for _, r := range criteria {
		in.require(r)
	}
}

func (in *JSONIn) require(criterion rune) {
	r, _, err := in.reader.ReadRune()
	if err != nil {
		panic("Error while reading required rune: " + err.Error())
	}
	if r != criterion {
		panic("Required rune not read")
	}
}

func (in *JSONIn) readString(out []Atom) []Atom {
	// TODO: improve
	out = append(out, NewAtomStringTerminal())
	for {
		r, _, err := in.reader.ReadRune()
		if err != nil {
			panic("Missing closing terminal in string input: " + err.Error())
		}
		if r == '"' {
			break
		}
		if r == '\\' {
			r, _, err = in.reader.ReadRune()
			if err != nil {
				panic("Missing rune after \\")
			}
			if len(out) == cap(out) {
				newOut := make([]Atom, len(out), cap(out) * 2)
				copy(newOut, out)
				out = newOut
			}
			out = append(out, NewAtomStringRune(r))
			continue
		}
		if len(out) == cap(out) {
			newOut := make([]Atom, len(out), cap(out) * 2)
			copy(newOut, out)
			out = newOut
		}
		out = append(out, NewAtomStringRune(r))
	}
	out = append(out, NewAtomStringTerminal())
	return out
}

func (in *JSONIn) Read() (WalkItem, error) {
	restart:
	// TODO: Escaping
	// TODO: Don't allow trailing commas
	// TODO: Proper float parsing with e and stuff
	r, err := in.nextNonWsRune()
	if err != nil {
		return WalkItem {}, err
	}
	state := in.structure[len(in.structure) - 1]
	switch state {
		case JSONInMap:
			in.popPath()
			if r == '}' {
				in.structure[len(in.structure) - 1] = JSONInValueEnd
				return WalkItem {
					Value: []Atom{NewAtomTerminal(MapEnd)},
					Path: in.path,
				}, nil
			}
			if r != '"' {
				panic("Expected key, found something else")
			}
			in.path = in.readString(in.path)
			r, err = in.nextNonWsRune()
			if err != nil {
				panic("Expected : got: " + err.Error())
			}
			if r != ':' {
				panic("Expected : after key")
			}
			r, err = in.nextNonWsRune()
			if err != nil {
				panic("Missing map value after key: " + err.Error())
			}
		case JSONInArray:
			if r == ']' {
				in.structure[len(in.structure) - 1] = JSONInValueEnd
				in.popPath()
				return WalkItem {
					Value: []Atom{NewAtomTerminal(ArrayEnd)},
					Path: in.path,
				}, nil
			}
			prevIndex := in.path[len(in.path) - 1]
			if prevIndex.Typ == AtomNull {
				prevIndex.Typ = AtomNumber
				prevIndex.data = math.Float64bits(0)
			} else if prevIndex.Typ == AtomNumber {
				prevIndex.data = math.Float64bits(math.Float64frombits(prevIndex.data) + 1)
			} else {
				panic("Invalid index in array input")
			}
			in.path[len(in.path) - 1] = prevIndex
		case JSONInRoot:
		case JSONInValueEnd:
			in.structure = in.structure[:len(in.structure) - 1]
			underState := in.structure[len(in.structure) - 1]
			if underState == JSONInRoot {
				panic("More input after root JSON object ends")
			} else if underState == JSONInMap && r == '}' {
				in.structure[len(in.structure) - 1] = JSONInValueEnd
				in.popPath()
				return WalkItem {
					Value: []Atom{NewAtomTerminal(MapEnd)},
					Path: in.path,
				}, nil
			} else if underState == JSONInArray && r == ']' {
				in.structure[len(in.structure) - 1] = JSONInValueEnd
				in.popPath()
				return WalkItem {
					Value: []Atom{NewAtomTerminal(ArrayEnd)},
					Path: in.path,
				}, nil
			}
			if r != ',' {
				panic("Expected , after JSON value, found: \"" + string(r) + "\"")
			}
			goto restart
		default:
			panic("Invalid JSONIn state")
	}
	switch r {
		case 'n':
			in.requireString("ull")
			in.structure = append(in.structure, JSONInValueEnd)
			return WalkItem {
				Value: []Atom{NewAtomNull()},
				Path: in.path,
			}, nil
		case 'f':
			in.requireString("alse")
			in.structure = append(in.structure, JSONInValueEnd)
			return WalkItem {
				Value: []Atom{NewAtomBool(false)},
				Path: in.path,
			}, nil
		case 't':
			in.requireString("rue")
			in.structure = append(in.structure, JSONInValueEnd)
			return WalkItem {
				Value: []Atom{NewAtomBool(true)},
				Path: in.path,
			}, nil
		case '"':
			value := make([]Atom, 0, 64)
			value = in.readString(value)
			in.structure = append(in.structure, JSONInValueEnd)
			return WalkItem {
				Value: value,
				Path: in.path,
			}, nil
		case '{':
			in.structure = append(in.structure, JSONInMap)
			in.path = append(in.path, NewAtomNull())
			return WalkItem {
				Value: []Atom{NewAtomTerminal(MapBegin)},
				Path: in.path[:len(in.path) - 1],
			}, nil
		case '[':
			in.structure = append(in.structure, JSONInArray)
			in.path = append(in.path, NewAtomNull())
			return WalkItem {
				Value: []Atom{NewAtomTerminal(ArrayBegin)},
				Path: in.path[:len(in.path) - 1],
			}, nil
	}
	if isNumberRune(r) {
		var builder strings.Builder
		builder.WriteRune(r)
		for {
			r, _, err = in.reader.ReadRune()
			if err != nil || !isNumberRune(r) {
				break
			}
			builder.WriteRune(r)
		}
		in.reader.UnreadRune()
		number, parseError := strconv.ParseFloat(builder.String(), 64)
		if parseError != nil {
			panic("Invalid number")
		}
		in.structure = append(in.structure, JSONInValueEnd)
		return WalkItem {
			Value: []Atom{NewAtomNumber(number)},
			Path: in.path,
		}, nil
	}
	panic("Invalid JSON value")
}

func (in *JSONIn) AssertDone() {
	if len(in.structure) != 2 || in.structure[0] != JSONInRoot || in.structure[1] != JSONInValueEnd {
		panic("Input ended on incomplete JSON root")
	}
}