package walk

import (
	"fmt"
	"strings"
	"math"
	"unicode/utf8"
	"bufio"
	"strconv"
)

// int or string
type PathSegment interface {}
func stringPathSegment(segment PathSegment) string {
	return fmt.Sprintf("%v", segment)
}
type Path []PathSegment
func (path Path) ToWalkValues() []WalkValue {
	var values []WalkValue
	for _, segment := range path {
		switch s := segment.(type) {
			case int:
				values = append(values, ValueNumber(s))
			case string:
				values = append(values, ValueString(s))
			default:
				panic("Invalid PathSegment")
		}
	}
	return values
}

func PathFromWalkValues(values []WalkValue) Path {
	var segments []PathSegment
	for _, value := range values {
		switch v := value.(type) {
			case ValueNumber:
				segments = append(segments, int(math.Round(float64(v))))
			case ValueString:
				segments = append(segments, string(v))
			default:
				panic("Invalid value in path")
		}
	}
	return segments
}

type TerminalValue int
const (
	ArrayBegin TerminalValue = iota
	ArrayEnd
	MapBegin
	MapEnd
)
func (value TerminalValue) Atomise(in []Atom) []Atom {
	return append(in, NewAtomTerminal(value))
}
func (value TerminalValue) String() string {
	switch value {
		case ArrayBegin:
			return "["
		case ArrayEnd:
			return "]"
		case MapBegin:
			return "{"
		case MapEnd:
			return "}"
		default:
			panic("Unknown TerminalValue")
	}
}

type ValueNull struct {}
func (value ValueNull) Atomise(in []Atom) []Atom {
	return append(in, NewAtomNull())
}
func (value ValueNull) String() string {
	return "null"
}

type ValueBool bool
func (value ValueBool) Atomise(in []Atom) []Atom {
	return append(in, NewAtomBool(bool(value)))
}
func (value ValueBool) String() string {
	if value {
		return "true"
	} else {
		return "false"
	}
}

type ValueNumber float64
func (value ValueNumber) Atomise(in []Atom) []Atom {
	return append(in, NewAtomNumber(float64(value)))
}
func (value ValueNumber) String() string {
	v := float64(value)
	return fmt.Sprintf("%f", v)
}

type ValueString string
func (value ValueString) Atomise(in []Atom) []Atom {
	in = append(in, NewAtomStringTerminal())
	for _, char := range value {
		in = append(in, NewAtomStringRune(char))
	}
	in = append(in, NewAtomStringTerminal())
	return in
}
func (value ValueString) String() string {
	return fmt.Sprintf("\"%s\"", string(value))
}

type AtomType int64
const (
	AtomNull AtomType = iota
	AtomBool
	AtomNumber
	AtomTerminal
	AtomStringTerminal
	AtomStringRune
)
type Atom struct {
	Typ AtomType
	data uint64
}
func NewAtomNull() Atom {
	return Atom {
		Typ: AtomNull,
		data: 0,
	}
}
func NewAtomBool(v bool) Atom {
	if v {
		return Atom {
			Typ: AtomBool,
			data: 1,
		}
	} else {
		return Atom {
			Typ: AtomBool,
			data: 0,
		}
	}
}
func NewAtomNumber(v float64) Atom {
	return Atom {
		Typ: AtomNumber,
		data: math.Float64bits(v),
	}
}
func NewAtomTerminal(v TerminalValue) Atom {
	return Atom {
		Typ: AtomTerminal,
		data: uint64(v),
	}
}
func NewAtomStringTerminal() Atom {
	return Atom {
		Typ: AtomStringTerminal,
		data: 0,
	}
}
func NewAtomStringRune(v rune) Atom {
	return Atom {
		Typ: AtomStringRune,
		data: uint64(v),
	}
}
func (v Atom) String() string {
	switch v.Typ {
		case AtomNull:
			return "null"
		case AtomBool:
			if v.data == 0 {
				return "false"
			}
			return "true"
		case AtomNumber:
			return fmt.Sprintf("%v", math.Float64frombits(v.data))
		case AtomTerminal:
			switch TerminalValue(v.data) {
				case MapBegin:
					return "{"
				case MapEnd:
					return "}"
				case ArrayBegin:
					return "["
				case ArrayEnd:
					return "]"
				default:
					panic("Invalid terminal atom")
			}
		case AtomStringTerminal:
			return "\""
		case AtomStringRune:
			return string(rune(v.data))
		default:
			panic("Invalid atom type")
	}
}

type WalkValue interface {
	// Append this values atoms to the input
	Atomise(in []Atom) []Atom
	String() string
}

type WalkItem struct {
	Value []Atom
	Path []Atom
}

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) {
	item, err := in.read()
	if err != nil {
		return item, err
	}
	return WalkItem {
		Value: item.Value,
		Path: append([]Atom{}, item.Path...),
	}, err
}

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")
	}
}

type JSONOutStructure int
const (
	JSONOutRoot JSONOutStructure = iota
	JSONOutMap
	JSONOutArray
	JSONOutString
	JSONOutValueEnd
)

type JSONOut struct {
	structure []JSONOutStructure
	writer *bufio.Writer
}

func (out *JSONOut) indent(adjust int) {
	fmt.Fprint(out.writer, strings.Repeat("\t", len(out.structure) - 1 + adjust))
}

func (out *JSONOut) atomOut(key string, atom Atom) {
	state := out.structure[len(out.structure) - 1]
	switch state {
		case JSONOutRoot, JSONOutMap, JSONOutArray:
			switch atom.Typ {
				case AtomNull, AtomBool, AtomNumber:
					out.indent(0)
					if state == JSONOutMap {
						fmt.Fprintf(out.writer, "%q: ", key)
					}
					fmt.Fprint(out.writer, atom.String())
					out.structure = append(out.structure, JSONOutValueEnd)
				case AtomStringTerminal:
					out.indent(0)
					if state == JSONOutMap {
						fmt.Fprintf(out.writer, "%q: ", key)
					}
					fmt.Fprint(out.writer, "\"")
					out.structure = append(out.structure, JSONOutString)
				case AtomTerminal:
					switch TerminalValue(atom.data) {
						case MapBegin:
							out.indent(0)
							if state == JSONOutMap {
								fmt.Fprintf(out.writer, "%q: ", key)
							}
							fmt.Fprint(out.writer, "{\n")
							out.structure = append(out.structure, JSONOutMap)
						case ArrayBegin:
							out.indent(0)
							if state == JSONOutMap {
								fmt.Fprintf(out.writer, "%q: ", key)
							}
							fmt.Fprint(out.writer, "[\n")
							out.structure = append(out.structure, JSONOutArray)
						case MapEnd:
							out.indent(-1)
							if state != JSONOutMap {
								panic("Map ended while not inside a map")
							}
							fmt.Fprint(out.writer, "}")
							out.structure[len(out.structure) - 1] = JSONOutValueEnd
						case ArrayEnd:
							out.indent(-1)
							if state != JSONOutArray {
								panic("Array ended while not inside a array")
							}
							fmt.Fprint(out.writer, "]")
							out.structure[len(out.structure) - 1] = JSONOutValueEnd
						default:
							panic("Invalid TerminalValue")
					}
				default:
					panic("Invalid AtomType in root value")
			}
		case JSONOutValueEnd:
			out.structure = out.structure[:len(out.structure) - 1]
			underState := out.structure[len(out.structure) - 1]
			if underState == JSONOutMap && atom.Typ == AtomTerminal && TerminalValue(atom.data) == MapEnd {
				fmt.Fprint(out.writer, "\n")
				out.indent(-1)
				fmt.Fprint(out.writer, "}")
				out.structure[len(out.structure) - 1] = JSONOutValueEnd
			} else if underState == JSONOutArray && atom.Typ == AtomTerminal && TerminalValue(atom.data) == ArrayEnd {
				fmt.Fprint(out.writer, "\n")
				out.indent(-1)
				fmt.Fprint(out.writer, "]")
				out.structure[len(out.structure) - 1] = JSONOutValueEnd
			} else if underState == JSONOutRoot {
				panic("Tried to output JSON after root value has concluded")
			} else {
				fmt.Fprint(out.writer, ",\n")
				out.atomOut(key, atom)
			}
		case JSONOutString:
			if atom.Typ == AtomStringTerminal {
				fmt.Fprint(out.writer, "\"")
				out.structure[len(out.structure) - 1] = JSONOutValueEnd
			} else {
				fmt.Fprint(out.writer, atom.String())
			}
		default:
			panic("Invalid JSONOutState")
	}
}

func (out *JSONOut) Print(path Path, values []Atom) {
	var segment PathSegment
	if len(path) > 0 {
		segment = path[len(path) - 1]
	}
	segmentString := stringPathSegment(segment)
	for _, atom := range values {
		out.atomOut(segmentString, atom)
	}
}

func (out *JSONOut) AssertDone() {
	out.writer.Flush()
	if len(out.structure) != 2 || out.structure[0] != JSONOutRoot || out.structure[1] != JSONOutValueEnd {
		panic("Program ended with incomplete JSON output")
	}
}

func NewJSONOut(writer *bufio.Writer) JSONOut {
	return JSONOut {
		structure: []JSONOutStructure{JSONOutRoot},
		writer: writer,
	}
}

func ConcatData(first []Atom, second []Atom) []Atom {
	res := make([]Atom, 0, len(first) + len(second))
	res = append(res, first...)
	res = append(res, second...)
	return res
}

func Atomise(in []WalkValue) (out []Atom) {
	numAtoms := 0
	for _, value := range in {
		switch v := value.(type) {
			case TerminalValue, ValueNull, ValueBool, ValueNumber:
				numAtoms++
			case ValueString:
				numAtoms += utf8.RuneCountInString(string(v)) + 2
			default:
				panic("Invalid WalkValue")
		}
	}
	out = make([]Atom, 0, numAtoms)
	for _, value := range in {
		out = value.Atomise(out)
	}
	return out
}

type CompoundError int

const (
	CompoundRuneOutsideString CompoundError = iota
	CompoundUnknownAtom
	CompoundMissingEnd
	CompoundInvalidStringAtom
)

func (err CompoundError) Error() string {
	switch err {
		case CompoundRuneOutsideString:
			return "Compound Error: Rune Outside String"
		case CompoundUnknownAtom:
			return "Compound Error: Unknown Atom"
		case CompoundMissingEnd:
			return "Compound Error: Missing End"
		case CompoundInvalidStringAtom:
			return "Compound Error: Invalid String Atom"
		default:
			panic("Invalid CompoundError")
	}
}

type CompoundResult struct {
	value WalkValue
	error error
}

func Compound(in []Atom) (out []WalkValue, error error) {
	numValues := 0
	i := 0
	inString := false
	for _, atom := range in {
		switch atom.Typ {
			case AtomNull, AtomBool, AtomNumber, AtomTerminal:
				if !inString {
					numValues++
				}
			case AtomStringTerminal:
				if inString {
					numValues++
				}
				inString = !inString
		}
	}
	i = 0
	out = make([]WalkValue, 0, numValues)
	for {
		if i >= len(in) {
			break
		}
		atom := in[i]
		i++
		switch atom.Typ {
			case AtomNull:
				out = append(out, ValueNull{})
				continue
			case AtomBool:
				out = append(out, ValueBool(atom.data != 0))
				continue
			case AtomNumber:
				out = append(out, ValueNumber(math.Float64frombits(atom.data)))
				continue
			case AtomTerminal:
				out = append(out, TerminalValue(atom.data))
				continue
			case AtomStringRune:
				return nil, CompoundRuneOutsideString
			case AtomStringTerminal:
			default:
				return nil, CompoundUnknownAtom
		}
		// Handle string start
		var builder strings.Builder
		for {
			if i >= len(in) {
				return nil, CompoundMissingEnd
			}
			atom := in[i]
			i++
			if atom.Typ == AtomStringTerminal {
				break
			}
			builder.WriteString(atom.String())
		}
		out = append(out, ValueString(builder.String()))
	}
	return out, nil
}