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

package lowmemjson

import (
	"bytes"

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

type reEncodeCompactWSIfUnder struct {
	out reEncoderModule

	// CompactWSIfUnder runs uses reEncodeCompactWScauses for
	// individual elements if doing so would cause that element to
	// be under this number of bytes.
	//
	// This has O(2^min(CompactWSIfUnder, depth)) time overhead,
	// so set with caution.
	CompactWSIfUnder int

	// state
	compactor        reEncodeWrite
	compacted        bytes.Buffer
	full             []handleRuneCall
	endWhenStackSize int
}

var _ reEncoderModule = (*reEncodeCompactWSIfUnder)(nil)

type handleRuneCall struct {
	c         rune
	t         jsonparse.RuneType
	escape    BackslashEscapeMode
	stackSize int
}

func (enc *reEncodeCompactWSIfUnder) reset() {
	enc.compactor = reEncodeWrite{}
	enc.compacted.Reset()
	enc.full = enc.full[:0]
	enc.endWhenStackSize = 0
}

func (enc *reEncodeCompactWSIfUnder) PopWriteBarrier() {
	enc.out.PopWriteBarrier()
}

func (enc *reEncodeCompactWSIfUnder) HandleRune(c rune, t jsonparse.RuneType, escape BackslashEscapeMode, stackSize int) error {
	if enc.compactor.out == nil { // not speculating
		switch t {
		case jsonparse.RuneTypeObjectBeg, jsonparse.RuneTypeArrayBeg: // start speculating
			enc.endWhenStackSize = stackSize - 1
			enc.compactor = reEncodeWrite{
				out: &enc.compacted,
			}
			enc.full = append(enc.full, handleRuneCall{
				c:         c,
				t:         t,
				escape:    escape,
				stackSize: stackSize,
			})
			return enc.compactor.HandleRune(c, t, escape, stackSize)
		default:
			return enc.out.HandleRune(c, t, escape, stackSize)
		}
	} else { // speculating
		enc.full = append(enc.full, handleRuneCall{
			c:         c,
			t:         t,
			escape:    escape,
			stackSize: stackSize,
		})
		if t != jsonparse.RuneTypeSpace {
			if err := enc.compactor.HandleRune(c, t, escape, stackSize); err != nil {
				return err
			}
		}
		switch {
		case enc.compacted.Len() >= enc.CompactWSIfUnder: // stop speculating; use indent
			buf := append([]handleRuneCall(nil), enc.full...)
			enc.reset()
			if err := enc.out.HandleRune(buf[0].c, buf[0].t, buf[0].escape, buf[0].stackSize); err != nil {
				return err
			}
			for _, tuple := range buf[1:] {
				if err := enc.HandleRune(tuple.c, tuple.t, tuple.escape, tuple.stackSize); err != nil {
					return err
				}
			}
		case stackSize == enc.endWhenStackSize: // stop speculating; use compact
			for _, tuple := range enc.full {
				if tuple.t == jsonparse.RuneTypeSpace {
					continue
				}
				if err := enc.out.HandleRune(tuple.c, tuple.t, tuple.escape, tuple.stackSize); err != nil {
					return err
				}
			}
			enc.reset()
		}
		return nil
	}
}