// Copyright (C) 2022  Luke Shumaker <lukeshu@lukeshu.com>
//
// SPDX-License-Identifier: GPL-2.0-or-later

package lowmemjson

import (
	"bufio"
	"bytes"
	"encoding"
	"encoding/json"
	"fmt"
	"io"
	"reflect"
	"strconv"
	"strings"
)

type Decodable interface {
	DecodeJSON(io.RuneScanner) error
}

type runeBuffer interface {
	io.Writer
	WriteRune(rune) (int, error)
	Reset()
}

type Decoder struct {
	io  runeTypeScanner
	buf *bufio.Reader

	// config
	disallowUnknownFields bool
	useNumber             bool

	// state
	err   error
	stack []any
}

var forceBufio bool // used by the tests to implement .Buffered()

func NewDecoder(r io.Reader) *Decoder {
	rr, ok := r.(io.RuneReader)
	var ret Decoder
	if forceBufio || !ok {
		rr = bufio.NewReader(r)
		if forceBufio {
			ret.buf = rr.(*bufio.Reader)
		}
	}
	ret.io = &noWSRuneTypeScanner{
		inner: &runeTypeScannerImpl{
			inner: rr,
		},
	}
	return &ret
}

func (dec *Decoder) DisallowUnknownFields() { dec.disallowUnknownFields = true }
func (dec *Decoder) UseNumber()             { dec.useNumber = true }
func (dec *Decoder) InputOffset() int64     { return dec.io.InputOffset() }

func (dec *Decoder) More() bool {
	dec.io.Reset()
	_, _, t, e := dec.io.ReadRuneType()
	dec.io.UnreadRune()
	return e == nil && t != RuneTypeEOF
}

func (dec *Decoder) stackPush(idx any) {
	dec.stack = append(dec.stack, idx)
}
func (dec *Decoder) stackPop() {
	dec.stack = dec.stack[:len(dec.stack)-1]
}
func (dec *Decoder) stackStr() string {
	var buf strings.Builder
	buf.WriteString("v")
	for _, item := range dec.stack {
		fmt.Fprintf(&buf, "[%#v]", item)
	}
	return buf.String()
}

func Decode(r io.Reader, ptr any) error {
	return NewDecoder(r).Decode(ptr)
}

func (dec *Decoder) Decode(ptr any) (err error) {
	ptrVal := reflect.ValueOf(ptr)
	if ptrVal.Kind() != reflect.Pointer || ptrVal.IsNil() || !ptrVal.Elem().CanSet() {
		return &json.InvalidUnmarshalError{
			// don't use ptrVal.Type() because ptrVal might be invalid if ptr==nil
			Type: reflect.TypeOf(ptr),
		}
	}

	if dec.err != nil {
		return dec.err
	}

	dec.io.Reset()
	defer func() {
		if r := recover(); r != nil {
			if de, ok := r.(decodeError); ok {
				dec.err = de.Err
				err = dec.err
			} else {
				panic(r)
			}
		}
	}()
	dec.decode(ptrVal.Elem(), false)
	return nil
}

// io helpers //////////////////////////////////////////////////////////////////////////////////////

type decodeError struct {
	Err error
}

func (dec *Decoder) panicType(typ reflect.Type, err error) {
	panic(decodeError{fmt.Errorf("json: type mismatch error at input byte %v: %s: type %v: %w",
		dec.InputOffset(), dec.stackStr(), typ, err)})
}

func (dec *Decoder) readRune() (rune, RuneType) {
	c, _, t, e := dec.io.ReadRuneType()
	if e != nil {
		panic(decodeError{e})
	}
	return c, t
}

func (dec *Decoder) unreadRune() {
	if err := dec.io.UnreadRune(); err != nil {
		// .UnreadRune() must succeed if the previous call was
		// .ReadRune(), which it always is for this code.
		panic("should not happen")
	}
}

func (dec *Decoder) peekRuneType() RuneType {
	_, t := dec.readRune()
	dec.unreadRune()
	return t
}

func (dec *Decoder) expectRune(ec rune, et RuneType) {
	ac, at := dec.readRune()
	if ac != ec || at != et {
		panic("should not happen")
	}
}

func (dec *Decoder) expectRuneType(ec rune, et RuneType) {
	ac, at := dec.readRune()
	if ac != ec || at != et {
		dec.panicType(nil, fmt.Errorf("TODO error message"))
	}
}

type decRuneTypeScanner struct {
	dec *Decoder
}

func (sc *decRuneTypeScanner) ReadRuneType() (rune, int, RuneType, error) {
	c, s, t, e := sc.dec.io.ReadRuneType()
	if e != nil {
		panic(decodeError{e})
	}
	return c, s, t, nil
}
func (sc *decRuneTypeScanner) ReadRune() (rune, int, error) {
	r, s, t, _ := sc.ReadRuneType()
	switch t {
	case RuneTypeEOF:
		return 0, 0, io.EOF
	default:
		return r, s, nil
	}
}

func (sc *decRuneTypeScanner) UnreadRune() error  { return sc.dec.io.UnreadRune() }
func (sc *decRuneTypeScanner) InputOffset() int64 { return sc.dec.InputOffset() }
func (sc *decRuneTypeScanner) Reset()             { sc.dec.io.Reset() }

func (dec *Decoder) limitingScanner() runeTypeScanner {
	return &elemRuneTypeScanner{
		inner: &decRuneTypeScanner{
			dec: dec,
		},
	}
}

// decoder main ////////////////////////////////////////////////////////////////////////////////////

var (
	rawMessagePtrType   = reflect.TypeOf((*json.RawMessage)(nil))
	decodableType       = reflect.TypeOf((*Decodable)(nil)).Elem()
	jsonUnmarshalerType = reflect.TypeOf((*json.Unmarshaler)(nil)).Elem()
	textUnmarshalerType = reflect.TypeOf((*encoding.TextUnmarshaler)(nil)).Elem()
	boolType            = reflect.TypeOf(true)
)

var kind2bits = map[reflect.Kind]int{
	reflect.Int:   int(32 << (^uint(0) >> 63)),
	reflect.Int8:  8,
	reflect.Int16: 16,
	reflect.Int32: 32,
	reflect.Int64: 64,

	reflect.Uint:   int(32 << (^uint(0) >> 63)),
	reflect.Uint8:  8,
	reflect.Uint16: 16,
	reflect.Uint32: 32,
	reflect.Uint64: 64,

	reflect.Uintptr: int(32 << (^uintptr(0) >> 63)),

	reflect.Float32: 32,
	reflect.Float64: 64,
}

func (dec *Decoder) decode(val reflect.Value, nullOK bool) {
	typ := val.Type()
	switch {
	case val.CanAddr() && reflect.PointerTo(typ) == rawMessagePtrType:
		var buf bytes.Buffer
		dec.scan(&buf)
		if err := val.Addr().Interface().(*json.RawMessage).UnmarshalJSON(buf.Bytes()); err != nil {
			dec.panicType(typ, err)
		}
	case val.CanAddr() && reflect.PointerTo(typ).Implements(decodableType):
		obj := val.Addr().Interface().(Decodable)
		if err := obj.DecodeJSON(dec.limitingScanner()); err != nil {
			dec.panicType(typ, err)
		}
	case val.CanAddr() && reflect.PointerTo(typ).Implements(jsonUnmarshalerType):
		var buf bytes.Buffer
		dec.scan(&buf)
		obj := val.Addr().Interface().(json.Unmarshaler)
		if err := obj.UnmarshalJSON(buf.Bytes()); err != nil {
			dec.panicType(typ, err)
		}
	case val.CanAddr() && reflect.PointerTo(typ).Implements(textUnmarshalerType):
		if nullOK && dec.peekRuneType() == RuneTypeNullN {
			dec.decodeNull()
			return
		}
		var buf bytes.Buffer
		dec.decodeString(&buf)
		obj := val.Addr().Interface().(encoding.TextUnmarshaler)
		if err := obj.UnmarshalText(buf.Bytes()); err != nil {
			dec.panicType(typ, err)
		}
	default:
		kind := typ.Kind()
		switch kind {
		case reflect.Bool:
			if nullOK && dec.peekRuneType() == RuneTypeNullN {
				dec.decodeNull()
				return
			}
			val.SetBool(dec.decodeBool())
		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
			if nullOK && dec.peekRuneType() == RuneTypeNullN {
				dec.decodeNull()
				return
			}
			var buf strings.Builder
			dec.scanNumber(&buf)
			n, err := strconv.ParseInt(buf.String(), 10, kind2bits[kind])
			if err != nil {
				dec.panicType(typ, err)
			}
			val.SetInt(n)
		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
			if nullOK && dec.peekRuneType() == RuneTypeNullN {
				dec.decodeNull()
				return
			}
			var buf strings.Builder
			dec.scanNumber(&buf)
			n, err := strconv.ParseUint(buf.String(), 10, kind2bits[kind])
			if err != nil {
				dec.panicType(typ, err)
			}
			val.SetUint(n)
		case reflect.Float32, reflect.Float64:
			if nullOK && dec.peekRuneType() == RuneTypeNullN {
				dec.decodeNull()
				return
			}
			var buf strings.Builder
			dec.scanNumber(&buf)
			n, err := strconv.ParseFloat(buf.String(), kind2bits[kind])
			if err != nil {
				dec.panicType(typ, err)
			}
			val.SetFloat(n)
		case reflect.String:
			if nullOK && dec.peekRuneType() == RuneTypeNullN {
				dec.decodeNull()
				return
			}
			var buf strings.Builder
			if typ == numberType {
				dec.scanNumber(&buf)
				val.SetString(buf.String())
			} else {
				dec.decodeString(&buf)
				val.SetString(buf.String())
			}
		case reflect.Interface:
			if typ.NumMethod() > 0 {
				dec.panicType(typ, fmt.Errorf("cannot decode in to non-empty interface"))
			}
			switch dec.peekRuneType() {
			case RuneTypeNullN:
				if !val.IsNil() && val.Elem().Kind() == reflect.Pointer && val.Elem().Elem().Kind() == reflect.Pointer {
					// XXX: I can't justify this case, other than "it's what encoding/json does, but
					// I don't understand their rationale".
					dec.decode(val.Elem(), false)
				} else {
					dec.decodeNull()
					val.Set(reflect.Zero(typ))
				}
			default:
				if !val.IsNil() && val.Elem().Kind() == reflect.Pointer {
					dec.decode(val.Elem(), false)
				} else {
					val.Set(reflect.ValueOf(dec.decodeAny()))
				}
			}
		case reflect.Struct:
			if nullOK && dec.peekRuneType() == RuneTypeNullN {
				dec.decodeNull()
				return
			}
			index := indexStruct(typ)
			var nameBuf strings.Builder
			dec.decodeObject(&nameBuf, func() {
				name := nameBuf.String()
				dec.stackPush(name)
				defer dec.stackPop()
				idx, ok := index.byName[name]
				if !ok {
					if dec.disallowUnknownFields {
						dec.panicType(typ, fmt.Errorf("unknown field %q", name))
					}
					dec.scan(io.Discard)
					return
				}
				field := index.byPos[idx]
				fVal := val
				for _, idx := range field.Path {
					if fVal.Kind() == reflect.Pointer {
						if fVal.IsNil() {
							if !fVal.CanSet() { // https://golang.org/issue/21357
								dec.panicType(fVal.Type().Elem(), fmt.Errorf("cannot set embedded pointer to unexported type"))
							}
							fVal.Set(reflect.New(fVal.Type().Elem()))
						}
						fVal = fVal.Elem()
					}
					fVal = fVal.Field(idx)
				}
				if field.Quote {
					switch dec.peekRuneType() {
					case RuneTypeNullN:
						dec.decodeNull()
						switch fVal.Kind() {
						// XXX: I can't justify this list, other than "it's what encoding/json
						// does, but I don't understand their rationale".
						case reflect.Interface, reflect.Pointer, reflect.Map, reflect.Slice:
							fVal.Set(reflect.Zero(fVal.Type()))
						default:
							// do nothing???
						}
					case RuneTypeStringBeg:
						// TODO: Figure out how to do this without buffering, have correct offsets.
						var buf bytes.Buffer
						dec.decodeString(&buf)
						subD := NewDecoder(&buf)
						subD.decode(fVal, false)
					default:
						dec.panicType(typ, fmt.Errorf(",string field TODO ERROR MESSAGE"))
					}
				} else {
					dec.decode(fVal, true)
				}
			})
		case reflect.Map:
			switch dec.peekRuneType() {
			case RuneTypeNullN:
				dec.decodeNull()
				val.Set(reflect.Zero(typ))
			case RuneTypeObjectBeg:
				if val.IsNil() {
					val.Set(reflect.MakeMap(typ))
				}
				var nameBuf bytes.Buffer
				dec.decodeObject(&nameBuf, func() {
					nameValTyp := typ.Key()
					nameValPtr := reflect.New(nameValTyp)
					switch {
					case reflect.PointerTo(nameValTyp).Implements(textUnmarshalerType):
						obj := nameValPtr.Interface().(encoding.TextUnmarshaler)
						if err := obj.UnmarshalText(nameBuf.Bytes()); err != nil {
							dec.panicType(nameValTyp, err)
						}
					default:
						switch nameValTyp.Kind() {
						case reflect.String:
							nameValPtr.Elem().SetString(nameBuf.String())
						case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
							n, err := strconv.ParseInt(nameBuf.String(), 10, kind2bits[nameValTyp.Kind()])
							if err != nil {
								dec.panicType(nameValTyp, err)
							}
							nameValPtr.Elem().SetInt(n)
						case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
							n, err := strconv.ParseUint(nameBuf.String(), 10, kind2bits[nameValTyp.Kind()])
							if err != nil {
								dec.panicType(nameValTyp, err)
							}
							nameValPtr.Elem().SetUint(n)
						default:
							dec.panicType(typ, fmt.Errorf("invalid map key type: %v", nameValTyp))
						}
					}
					dec.stackPush(nameValPtr.Elem())
					defer dec.stackPop()

					fValPtr := reflect.New(typ.Elem())
					dec.decode(fValPtr.Elem(), false)

					val.SetMapIndex(nameValPtr.Elem(), fValPtr.Elem())
				})
			default:
				dec.panicType(typ, fmt.Errorf("map: TODO"))
			}
		case reflect.Slice:
			switch {
			case typ.Elem().Kind() == reflect.Uint8:
				switch dec.peekRuneType() {
				case RuneTypeNullN:
					dec.decodeNull()
					val.Set(reflect.Zero(typ))
				case RuneTypeStringBeg:
					var buf bytes.Buffer
					dec.decodeString(newBase64Decoder(&buf))
					if typ.Elem() == byteType {
						val.Set(reflect.ValueOf(buf.Bytes()))
					} else {
						bs := buf.Bytes()
						// TODO: Surely there's a better way.
						val.Set(reflect.MakeSlice(typ, len(bs), len(bs)))
						for i := 0; i < len(bs); i++ {
							val.Index(i).Set(reflect.ValueOf(bs[i]).Convert(typ.Elem()))
						}
					}
				default:
					dec.panicType(typ, fmt.Errorf("byte slice: TODO"))
				}
			default:
				switch dec.peekRuneType() {
				case RuneTypeNullN:
					dec.decodeNull()
					val.Set(reflect.Zero(typ))
				case RuneTypeArrayBeg:
					if val.IsNil() {
						val.Set(reflect.MakeSlice(typ, 0, 0))
					}
					if val.Len() > 0 {
						val.Set(val.Slice(0, 0))
					}
					i := 0
					dec.decodeArray(func() {
						dec.stackPush(i)
						defer dec.stackPop()
						mValPtr := reflect.New(typ.Elem())
						dec.decode(mValPtr.Elem(), false)
						val.Set(reflect.Append(val, mValPtr.Elem()))
						i++
					})
				default:
					dec.panicType(typ, fmt.Errorf("slice: TODO"))
				}
			}
		case reflect.Array:
			if nullOK && dec.peekRuneType() == RuneTypeNullN {
				dec.decodeNull()
				return
			}
			i := 0
			n := val.Len()
			dec.decodeArray(func() {
				dec.stackPush(i)
				defer dec.stackPop()
				if i < n {
					mValPtr := reflect.New(typ.Elem())
					dec.decode(mValPtr.Elem(), false)
					val.Index(i).Set(mValPtr.Elem())
				} else {
					dec.scan(io.Discard)
				}
				i++
			})
			for ; i < n; i++ {
				val.Index(i).Set(reflect.Zero(typ.Elem()))
			}
		case reflect.Pointer:
			switch dec.peekRuneType() {
			case RuneTypeNullN:
				dec.decodeNull()
				/*
					for typ.Elem().Kind() == reflect.Pointer {
						if val.IsNil() || !val.Elem().CanSet() {
							val.Set(reflect.New(typ.Elem()))
						}
						val = val.Elem()
						typ = val.Type()
					}
				*/
				val.Set(reflect.Zero(typ))
			default:
				if val.IsNil() {
					val.Set(reflect.New(typ.Elem()))
				}
				dec.decode(val.Elem(), false)
			}
		default:
			dec.panicType(typ, fmt.Errorf("unsupported type (kind=%v)", typ.Kind()))
		}
	}
}

func (dec *Decoder) scan(out io.Writer) {
	limiter := dec.limitingScanner()
	for {
		c, _, err := limiter.ReadRune()
		if err == io.EOF {
			return
		}
		_, _ = writeRune(out, c)
	}
}

func (dec *Decoder) scanNumber(out io.Writer) {
	if !dec.peekRuneType().IsNumber() {
		dec.panicType(numberType, fmt.Errorf("number: not a number"))
	}
	dec.scan(out)
}

func (dec *Decoder) decodeAny() any {
	c, _ := dec.readRune()
	dec.unreadRune()
	switch c {
	case '{':
		ret := make(map[string]any)
		var nameBuf strings.Builder
		dec.decodeObject(&nameBuf, func() {
			name := nameBuf.String()
			dec.stackPush(name)
			defer dec.stackPop()
			ret[name] = dec.decodeAny()
		})
		return ret
	case '[':
		ret := []any{}
		dec.decodeArray(func() {
			dec.stackPush(len(ret))
			defer dec.stackPop()
			ret = append(ret, dec.decodeAny())
		})
		return ret
	case '"':
		var buf strings.Builder
		dec.decodeString(&buf)
		return buf.String()
	case '-', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9':
		var buf strings.Builder
		dec.scan(&buf)
		num := json.Number(buf.String())
		if dec.useNumber {
			return num
		}
		f64, err := num.Float64()
		if err != nil {
			panic("should not happen")
		}
		return f64
	case 't', 'f':
		return dec.decodeBool()
	case 'n':
		dec.decodeNull()
		return nil
	default:
		panic("should not happen")
	}
}

func (dec *Decoder) decodeObject(nameBuf runeBuffer, decodeKVal func()) {
	dec.expectRuneType('{', RuneTypeObjectBeg)
	_, t := dec.readRune()
	switch t {
	case RuneTypeObjectEnd:
		return
	case RuneTypeStringBeg:
	decodeMember:
		dec.unreadRune()
		nameBuf.Reset()
		dec.decodeString(nameBuf)
		dec.expectRune(':', RuneTypeObjectColon)
		decodeKVal()
		_, t := dec.readRune()
		switch t {
		case RuneTypeObjectComma:
			dec.expectRune('"', RuneTypeStringBeg)
			goto decodeMember
		case RuneTypeObjectEnd:
			return
		default:
			panic("should not happen")
		}
	default:
		panic("should not happen")
	}
}

func (dec *Decoder) decodeArray(decodeMember func()) {
	dec.expectRuneType('[', RuneTypeArrayBeg)
	_, t := dec.readRune()
	switch t {
	case RuneTypeArrayEnd:
		return
	default:
		dec.unreadRune()
	decodeNextMember:
		decodeMember()
		_, t := dec.readRune()
		switch t {
		case RuneTypeArrayComma:
			goto decodeNextMember
		case RuneTypeArrayEnd:
			return
		default:
			panic("should not happen")
		}
	}
}

func (dec *Decoder) decodeString(out io.Writer) {
	dec.expectRuneType('"', RuneTypeStringBeg)
	var uhex [4]byte
	for {
		c, t := dec.readRune()
		switch t {
		case RuneTypeStringChar:
			_, _ = writeRune(out, c)
		case RuneTypeStringEsc, RuneTypeStringEscU:
			// do nothing
		case RuneTypeStringEsc1:
			switch c {
			case '"':
				_, _ = writeRune(out, '"')
			case '\\':
				_, _ = writeRune(out, '\\')
			case '/':
				_, _ = writeRune(out, '/')
			case 'b':
				_, _ = writeRune(out, '\b')
			case 'f':
				_, _ = writeRune(out, '\f')
			case 'n':
				_, _ = writeRune(out, '\n')
			case 'r':
				_, _ = writeRune(out, '\r')
			case 't':
				_, _ = writeRune(out, '\t')
			default:
				panic("should not happen")
			}
		case RuneTypeStringEscUA:
			uhex[0], _ = hex2int(c)
		case RuneTypeStringEscUB:
			uhex[1], _ = hex2int(c)
		case RuneTypeStringEscUC:
			uhex[2], _ = hex2int(c)
		case RuneTypeStringEscUD:
			uhex[3], _ = hex2int(c)
			c = 0 |
				rune(uhex[0])<<12 |
				rune(uhex[1])<<8 |
				rune(uhex[2])<<4 |
				rune(uhex[3])<<0
			_, _ = writeRune(out, c)
		case RuneTypeStringEnd:
			return
		default:
			panic("should not happen")
		}
	}
}

func (dec *Decoder) decodeBool() bool {
	c, _ := dec.readRune()
	switch c {
	case 't':
		dec.expectRune('r', RuneTypeTrueR)
		dec.expectRune('u', RuneTypeTrueU)
		dec.expectRune('e', RuneTypeTrueE)
		return true
	case 'f':
		dec.expectRune('a', RuneTypeFalseA)
		dec.expectRune('l', RuneTypeFalseL)
		dec.expectRune('s', RuneTypeFalseS)
		dec.expectRune('e', RuneTypeFalseE)
		return false
	default:
		dec.panicType(boolType, fmt.Errorf("bool: expected %q or %q but got %q", 't', 'f', c))
		panic("not reached")
	}
}

func (dec *Decoder) decodeNull() {
	dec.expectRuneType('n', RuneTypeNullN)
	dec.expectRune('u', RuneTypeNullU)
	dec.expectRune('l', RuneTypeNullL1)
	dec.expectRune('l', RuneTypeNullL2)
}