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

// Package json is a wrapper around lowmemjson that is a (mostly)
// drop-in replacement for the standard library's encoding/json.
package json

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"strconv"
	"unicode/utf8"

	"git.lukeshu.com/go/lowmemjson"
	"git.lukeshu.com/go/lowmemjson/internal/jsonstring"
)

//nolint:stylecheck // ST1021 False positive; these aren't comments on individual types.
type (
	Number      = json.Number
	RawMessage  = json.RawMessage
	Marshaler   = json.Marshaler
	Unmarshaler = json.Unmarshaler

	// low-level decode errors.
	UnmarshalTypeError = json.UnmarshalTypeError
	// SyntaxError        = json.SyntaxError // Duplicated to access a private field.

	// high-level decode errors.
	InvalidUnmarshalError = json.InvalidUnmarshalError

	// marshal errors.
	UnsupportedTypeError  = json.UnsupportedTypeError
	UnsupportedValueError = json.UnsupportedValueError
	// MarshalerError        = json.MarshalerError // Duplicated to access a private field.
)

// Error conversion //////////////////////////////////////////////////

func convertError(err error, isUnmarshal bool) error {
	switch err := err.(type) {
	case nil:
		return nil
	case *lowmemjson.DecodeArgumentError:
		return err
	case *lowmemjson.DecodeError:
		switch suberr := err.Err.(type) {
		case *lowmemjson.DecodeReadError:
			return err
		case *lowmemjson.DecodeSyntaxError:
			if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
				if isUnmarshal {
					return &SyntaxError{
						msg:    "unexpected end of JSON input",
						Offset: suberr.Offset,
					}
				}
				return suberr.Err
			}
			return &SyntaxError{
				msg:    suberr.Err.Error(),
				Offset: suberr.Offset + 1,
			}
		case *lowmemjson.DecodeTypeError:
			switch subsuberr := suberr.Err.(type) {
			case *UnmarshalTypeError:
				// Populate the .Struct and .Field members.
				subsuberr.Struct = err.FieldParent
				subsuberr.Field = err.FieldName
				return subsuberr
			default:
				switch {
				case errors.Is(err, lowmemjson.ErrDecodeNonEmptyInterface),
					errors.Is(err, strconv.ErrSyntax),
					errors.Is(err, strconv.ErrRange):
					return &UnmarshalTypeError{
						Value:  suberr.JSONType,
						Type:   suberr.GoType,
						Offset: suberr.Offset,
						Struct: err.FieldParent,
						Field:  err.FieldName,
					}
				default:
					return subsuberr
				}
			case nil, *lowmemjson.DecodeArgumentError:
				return &UnmarshalTypeError{
					Value:  suberr.JSONType,
					Type:   suberr.GoType,
					Offset: suberr.Offset,
					Struct: err.FieldParent,
					Field:  err.FieldName,
				}
			}
		default:
			panic(fmt.Errorf("should not happen: unexpected lowmemjson.DecodeError sub-type: %T: %w", suberr, err))
		}
	case *lowmemjson.EncodeWriteError:
		return err
	case *lowmemjson.EncodeTypeError:
		return err
	case *lowmemjson.EncodeValueError:
		return err
	case *lowmemjson.EncodeMethodError:
		return &MarshalerError{
			Type:       err.Type,
			Err:        err.Err,
			sourceFunc: err.SourceFunc,
		}
	case *lowmemjson.ReEncodeWriteError:
		return err
	case *lowmemjson.ReEncodeSyntaxError:
		ret := &SyntaxError{
			msg:    err.Err.Error(),
			Offset: err.Offset + 1,
		}
		if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
			ret.msg = "unexpected end of JSON input"
		}
		return ret
	default:
		panic(fmt.Errorf("should not happen: unexpected lowmemjson error type: %T: %w", err, err))
	}
}

// Encode wrappers ///////////////////////////////////////////////////

func marshal(v any, cfg lowmemjson.ReEncoderConfig) ([]byte, error) {
	var buf bytes.Buffer
	if err := convertError(lowmemjson.NewEncoder(lowmemjson.NewReEncoder(&buf, cfg)).Encode(v), false); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

func MarshalIndent(v any, prefix, indent string) ([]byte, error) {
	return marshal(v, lowmemjson.ReEncoderConfig{
		Indent: indent,
		Prefix: prefix,
	})
}

func Marshal(v any) ([]byte, error) {
	return marshal(v, lowmemjson.ReEncoderConfig{
		Compact: true,
	})
}

type Encoder struct {
	out io.Writer
	buf bytes.Buffer

	cfg lowmemjson.ReEncoderConfig

	encoder   *lowmemjson.Encoder
	formatter *lowmemjson.ReEncoder
}

func NewEncoder(w io.Writer) *Encoder {
	ret := &Encoder{
		out: w,

		cfg: lowmemjson.ReEncoderConfig{
			AllowMultipleValues: true,

			Compact:               true,
			ForceTrailingNewlines: true,
		},
	}
	ret.refreshConfig()
	return ret
}

func (enc *Encoder) refreshConfig() {
	enc.formatter = lowmemjson.NewReEncoder(&enc.buf, enc.cfg)
	enc.encoder = lowmemjson.NewEncoder(enc.formatter)
}

func (enc *Encoder) Encode(v any) error {
	if err := convertError(enc.encoder.Encode(v), false); err != nil {
		enc.buf.Reset()
		return err
	}
	if _, err := enc.buf.WriteTo(enc.out); err != nil {
		return err
	}
	return nil
}

func (enc *Encoder) SetEscapeHTML(on bool) {
	if on {
		enc.cfg.BackslashEscape = lowmemjson.EscapeDefault
	} else {
		enc.cfg.BackslashEscape = lowmemjson.EscapeDefaultNonHTMLSafe
	}
	enc.refreshConfig()
}

func (enc *Encoder) SetIndent(prefix, indent string) {
	enc.cfg.Compact = prefix == "" && indent == ""
	enc.cfg.Prefix = prefix
	enc.cfg.Indent = indent
	enc.refreshConfig()
}

// ReEncode wrappers /////////////////////////////////////////////////

func HTMLEscape(dst *bytes.Buffer, src []byte) {
	for n := 0; n < len(src); {
		c, size := utf8.DecodeRune(src[n:])
		if c == utf8.RuneError && size == 1 {
			dst.WriteByte(src[n])
		} else {
			mode := lowmemjson.EscapeHTMLSafe(c, lowmemjson.BackslashEscapeNone)
			switch mode {
			case lowmemjson.BackslashEscapeNone:
				dst.WriteRune(c)
			case lowmemjson.BackslashEscapeUnicode:
				_ = jsonstring.WriteStringUnicodeEscape(dst, c, mode)
			default:
				panic(fmt.Errorf("lowmemjson.EscapeHTMLSafe returned an unexpected escape mode=%d", mode))
			}
		}
		n += size
	}
}

func reencode(dst io.Writer, src []byte, cfg lowmemjson.ReEncoderConfig) error {
	formatter := lowmemjson.NewReEncoder(dst, cfg)
	_, err := formatter.Write(src)
	if err == nil {
		err = formatter.Close()
	}
	return convertError(err, false)
}

func Compact(dst *bytes.Buffer, src []byte) error {
	start := dst.Len()
	err := reencode(dst, src, lowmemjson.ReEncoderConfig{
		Compact:         true,
		InvalidUTF8:     lowmemjson.InvalidUTF8Preserve,
		BackslashEscape: lowmemjson.EscapePreserve,
	})
	if err != nil {
		dst.Truncate(start)
	}
	return err
}

func isSpace(c byte) bool {
	switch c {
	case 0x0020, 0x000A, 0x000D, 0x0009:
		return true
	default:
		return false
	}
}

func Indent(dst *bytes.Buffer, src []byte, prefix, indent string) error {
	start := dst.Len()
	err := reencode(dst, src, lowmemjson.ReEncoderConfig{
		Indent:          indent,
		Prefix:          prefix,
		InvalidUTF8:     lowmemjson.InvalidUTF8Preserve,
		BackslashEscape: lowmemjson.EscapePreserve,
	})
	if err != nil {
		dst.Truncate(start)
		return err
	}

	// Preserve trailing whitespace.
	lastNonWS := len(src) - 1
	for ; lastNonWS >= 0 && isSpace(src[lastNonWS]); lastNonWS-- {
	}
	if _, err := dst.Write(src[lastNonWS+1:]); err != nil {
		return err
	}

	return nil
}

func Valid(data []byte) bool {
	formatter := lowmemjson.NewReEncoder(io.Discard, lowmemjson.ReEncoderConfig{
		Compact:     true,
		InvalidUTF8: lowmemjson.InvalidUTF8Error,
	})
	if _, err := formatter.Write(data); err != nil {
		return false
	}
	if err := formatter.Close(); err != nil {
		return false
	}
	return true
}

// Decode wrappers ///////////////////////////////////////////////////

type decodeValidator struct{}

func (*decodeValidator) DecodeJSON(r io.RuneScanner) error {
	for {
		if _, _, err := r.ReadRune(); err != nil {

			if err == io.EOF {
				return nil
			}
			return err
		}
	}
}

var _ lowmemjson.Decodable = (*decodeValidator)(nil)

func Unmarshal(data []byte, ptr any) error {
	if err := convertError(lowmemjson.NewDecoder(bytes.NewReader(data)).DecodeThenEOF(&decodeValidator{}), true); err != nil {
		return err
	}
	if err := convertError(lowmemjson.NewDecoder(bytes.NewReader(data)).DecodeThenEOF(ptr), true); err != nil {
		return err
	}
	return nil
}

type teeRuneScanner struct {
	src interface {
		io.RuneScanner
		io.ByteScanner
	}
	dst      *bytes.Buffer
	lastSize int
}

func (tee *teeRuneScanner) ReadRune() (r rune, size int, err error) {
	r, size, err = tee.src.ReadRune()
	if err == nil {
		if r == utf8.RuneError && size == 1 {
			_ = tee.src.UnreadRune()
			b, _ := tee.src.ReadByte()
			_ = tee.dst.WriteByte(b)
		} else {
			_, _ = tee.dst.WriteRune(r)
		}
	}
	tee.lastSize = size
	return
}

func (tee *teeRuneScanner) UnreadRune() error {
	if tee.lastSize == 0 {
		return lowmemjson.ErrInvalidUnreadRune
	}
	_ = tee.src.UnreadRune()
	tee.dst.Truncate(tee.dst.Len() - tee.lastSize)
	tee.lastSize = 0
	return nil
}

func (tee *teeRuneScanner) ReadByte() (b byte, err error) {
	b, err = tee.src.ReadByte()
	if err == nil {
		_ = tee.dst.WriteByte(b)
		tee.lastSize = 1
	}
	return
}

func (tee *teeRuneScanner) UnreadByte() error {
	if tee.lastSize != 1 {
		return lowmemjson.ErrInvalidUnreadRune
	}
	_ = tee.src.UnreadByte()
	tee.dst.Truncate(tee.dst.Len() - tee.lastSize)
	tee.lastSize = 0
	return nil
}

type Decoder struct {
	validatorBuf *bufio.Reader
	validator    *lowmemjson.Decoder

	decoderBuf bytes.Buffer
	*lowmemjson.Decoder
}

func NewDecoder(r io.Reader) *Decoder {
	br, ok := r.(*bufio.Reader)
	if !ok {
		br = bufio.NewReader(r)
	}
	ret := &Decoder{
		validatorBuf: br,
	}
	ret.validator = lowmemjson.NewDecoder(&teeRuneScanner{
		src: ret.validatorBuf,
		dst: &ret.decoderBuf,
	})
	ret.Decoder = lowmemjson.NewDecoder(&ret.decoderBuf)
	return ret
}

func (dec *Decoder) Decode(ptr any) error {
	if err := convertError(dec.validator.Decode(&decodeValidator{}), false); err != nil {
		return err
	}
	if err := convertError(dec.Decoder.Decode(ptr), false); err != nil {
		return err
	}
	return nil
}

func (dec *Decoder) Buffered() io.Reader {
	dat, _ := dec.validatorBuf.Peek(dec.validatorBuf.Buffered())
	return bytes.NewReader(dat)
}

// func (dec *Decoder) Token() (Token, error)