// Package codec implements the CBOR profile of the DateKeys protocol (spec
// §58, §58.1) without reflection and without dependencies.
//
// The profile is Deterministic CBOR (RFC 8949 §4.2.1) restricted to major
// types 0 (unsigned integer), 2 (byte string), 3 (text string), 4 (array) and
// 5 (map), with unsigned integer map keys in strictly ascending order,
// integers and lengths in their shortest form, definite lengths only and
// valid UTF-8 text. Negative integers, tags, floats, simple values (false,
// true, null, undefined), indefinite lengths and every other map key are
// rejected with ErrNonCanonicalCBOR.
//
// Each schema writes its own encoding with an Encoder and reads it with a
// Decoder, a strict cursor that reads exactly what the schema asks for.
// Unmarshal runs the decoder of a schema and then re-encodes what it decoded:
// the input must be reproduced byte for byte, or it is ErrNonCanonicalCBOR.
// The same principle as dk1_ canonicality (spec §19): canonicality does not
// depend on the decoder rejecting every non-canonical form.
//
// Peek reads the type tag and the schema version of an object before strict
// decoding (spec §70). Walk checks that bytes are one data item of the
// profile; it is a helper for vectors, fuzzing and diagnostics, and never
// decides whether an object of the protocol is valid.
package codec

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"math"
	"unicode/utf8"

	datekeys "g.activething.com/go/DateKeys"
)

// MaxSafeUint is 2^53-1, the largest unsigned integer any schema of the
// protocol allows, so that every integer is exact as an IEEE 754 double
// (spec §58).
const MaxSafeUint = 1<<53 - 1

// MaxTypeTagLen bounds the type tag that Peek reads. Every type tag of V1 is
// at most 25 bytes, so a longer one is of no known schema; the bound also
// keeps an input-sized tag out of the errors. It is an implementation limit
// (spec §74).
const MaxTypeTagLen = 64

// Major types of the profile (spec §58).
const (
	majorUint  = 0
	majorBytes = 2
	majorText  = 3
	majorArray = 4
	majorMap   = 5
)

var majorNames = [8]string{
	"an unsigned integer", "a negative integer", "a byte string", "a text string",
	"an array", "a map", "a tag", "a float or simple value",
}

// errorf returns an error that wraps ErrNonCanonicalCBOR.
func errorf(format string, args ...any) error {
	return fmt.Errorf("codec: "+format+": %w", append(args, datekeys.ErrNonCanonicalCBOR)...)
}

// ---------------------------------------------------------------------------
// Encoder

// Encoder writes the deterministic encoding of data items of the profile:
// every integer and length in its shortest form, definite lengths only. The
// first error is kept and later calls do nothing; Out returns it. The zero
// value is ready to use.
//
// An Encoder does not know the schema: the caller writes the map keys, as
// unsigned integers in ascending order, and as many entries and items as it
// announced. The decoder of the schema checks both on the output (spec §72).
//
// An encoding may hold secrets, such as I_PAYLOAD or access_material: the
// Encoder wipes every buffer it outgrows, so that the output is the only
// copy, and the caller wipes the output.
type Encoder struct {
	buf []byte
	err error
}

// grow makes room for n more bytes, wiping the buffer it outgrows.
func (e *Encoder) grow(n int) {
	if cap(e.buf)-len(e.buf) >= n {
		return
	}
	b := make([]byte, len(e.buf), 2*cap(e.buf)+n)
	copy(b, e.buf)
	clear(e.buf)
	e.buf = b
}

// head appends the head of a data item: its major type and argument.
func (e *Encoder) head(major byte, arg uint64) {
	if e.err != nil {
		return
	}
	var h [9]byte
	h[0] = major << 5
	n := 1
	switch {
	case arg < 24:
		h[0] |= byte(arg)
	case arg <= math.MaxUint8:
		h[0] |= 24
		h[1] = byte(arg)
		n = 2
	case arg <= math.MaxUint16:
		h[0] |= 25
		binary.BigEndian.PutUint16(h[1:], uint16(arg))
		n = 3
	case arg <= math.MaxUint32:
		h[0] |= 26
		binary.BigEndian.PutUint32(h[1:], uint32(arg))
		n = 5
	default:
		h[0] |= 27
		binary.BigEndian.PutUint64(h[1:], arg)
		n = 9
	}
	e.grow(n)
	e.buf = append(e.buf, h[:n]...)
}

// Fail records err as the error of the encoding unless one is recorded
// already; a nil err is ignored. Every later call does nothing, and Out
// returns the first error. The encoder of a schema calls it when its value
// breaks a rule of the schema, so that bytes the decoder rejects are never
// returned.
func (e *Encoder) Fail(err error) {
	if e.err == nil {
		e.err = err
	}
}

// Map writes the head of a map of the given number of entries. The caller
// then writes each key, with Uint, followed by its value.
func (e *Encoder) Map(pairs int) {
	if pairs < 0 {
		e.Fail(errorf("map of %d entries", pairs))
		return
	}
	e.head(majorMap, uint64(pairs))
}

// Array writes the head of an array of the given number of items. The caller
// then writes each item.
func (e *Encoder) Array(items int) {
	if items < 0 {
		e.Fail(errorf("array of %d items", items))
		return
	}
	e.head(majorArray, uint64(items))
}

// Uint writes an unsigned integer.
func (e *Encoder) Uint(v uint64) { e.head(majorUint, v) }

// Bstr writes a byte string. A nil slice is the empty byte string.
func (e *Encoder) Bstr(b []byte) {
	e.head(majorBytes, uint64(len(b)))
	if e.err == nil {
		e.grow(len(b))
		e.buf = append(e.buf, b...)
	}
}

// Text writes a text string, which must be valid UTF-8.
func (e *Encoder) Text(s string) {
	if !utf8.ValidString(s) {
		e.Fail(errorf("text string %q is not valid UTF-8", s))
		return
	}
	e.head(majorText, uint64(len(s)))
	if e.err == nil {
		e.grow(len(s))
		e.buf = append(e.buf, s...)
	}
}

// Out returns the encoding, or the first error. On error the partial output
// is wiped.
func (e *Encoder) Out() ([]byte, error) {
	if e.err != nil {
		clear(e.buf)
		e.buf = nil
		return nil, e.err
	}
	return e.buf, nil
}

// ---------------------------------------------------------------------------
// Decoder

// Decoder is a strict cursor over the encoding of one data item of the
// profile. Each method reads one data item, or one head, and rejects with
// ErrNonCanonicalCBOR a major type outside the profile, a major type other
// than the one asked for, an indefinite length, an integer or length not in
// its shortest form, a length beyond the remaining input and a value outside
// the bounds the caller gives. Within each open map the keys are unsigned
// integers in strictly ascending order.
//
// The first error is kept: every later call returns it.
type Decoder struct {
	in   []byte
	off  int
	maps []openMap
	err  error
}

// openMap is the state of a map between Map and EndMap.
type openMap struct {
	left    uint64 // entries not read yet
	last    uint64 // last key read
	started bool   // at least one key was read
}

// NewDecoder returns a Decoder positioned at the start of in.
func NewDecoder(in []byte) *Decoder { return &Decoder{in: in} }

// fail records the first error, with the current offset, and returns it.
func (d *Decoder) fail(format string, args ...any) error {
	if d.err == nil {
		d.err = errorf("offset %d: "+format, append([]any{d.off}, args...)...)
	}
	return d.err
}

func (d *Decoder) remaining() int { return len(d.in) - d.off }

// head reads the head of the next data item and returns its major type and
// argument. It rejects the major types outside the profile, reserved values,
// indefinite lengths, arguments not in their shortest form and truncation.
func (d *Decoder) head() (byte, uint64, error) {
	if d.err != nil {
		return 0, 0, d.err
	}
	if d.off >= len(d.in) {
		return 0, 0, d.fail("truncated input")
	}
	b := d.in[d.off]
	major, info := b>>5, b&0x1f
	switch major {
	case 1, 6, 7:
		return 0, 0, d.fail("%s (initial byte %#02x) is outside the CBOR profile", majorNames[major], b)
	}
	switch {
	case info < 24:
		d.off++
		return major, uint64(info), nil
	case info == 31:
		return 0, 0, d.fail("indefinite length (initial byte %#02x)", b)
	case info > 27:
		return 0, 0, d.fail("reserved additional information (initial byte %#02x)", b)
	}
	n := 1 << (info - 24)
	if d.remaining() < 1+n {
		return 0, 0, d.fail("truncated input")
	}
	var arg, min uint64
	p := d.in[d.off+1 : d.off+1+n]
	switch n {
	case 1:
		arg, min = uint64(p[0]), 24
	case 2:
		arg, min = uint64(binary.BigEndian.Uint16(p)), math.MaxUint8+1
	case 4:
		arg, min = uint64(binary.BigEndian.Uint32(p)), math.MaxUint16+1
	default:
		arg, min = binary.BigEndian.Uint64(p), math.MaxUint32+1
	}
	if arg < min {
		return 0, 0, d.fail("%d is not in its shortest form (initial byte %#02x)", arg, b)
	}
	d.off += 1 + n
	return major, arg, nil
}

// expect reads the head of a data item of major type want.
func (d *Decoder) expect(want byte) (uint64, error) {
	start := d.off
	major, arg, err := d.head()
	if err != nil {
		return 0, err
	}
	if major != want {
		d.off = start
		return 0, d.fail("%s where %s was expected", majorNames[major], majorNames[want])
	}
	return arg, nil
}

// Map reads the head of a map of at most max entries and returns the number
// of entries. The caller reads each entry with Key and a value, then calls
// EndMap.
func (d *Decoder) Map(max int) (int, error) {
	n, err := d.expect(majorMap)
	if err != nil {
		return 0, err
	}
	if max < 0 || n > uint64(max) {
		return 0, d.fail("map of %d entries, at most %d", n, max)
	}
	if n > uint64(d.remaining())/2 {
		return 0, d.fail("truncated input: map of %d entries", n)
	}
	d.maps = append(d.maps, openMap{left: n})
	return int(n), nil
}

// Key reads the key of the next entry of the innermost open map: an unsigned
// integer greater than the previous key of that map.
func (d *Decoder) Key() (uint64, error) {
	if d.err != nil {
		return 0, d.err
	}
	if len(d.maps) == 0 {
		return 0, d.fail("map key outside a map")
	}
	m := &d.maps[len(d.maps)-1]
	if m.left == 0 {
		return 0, d.fail("map key after the last entry")
	}
	start := d.off
	k, err := d.expect(majorUint)
	if err != nil {
		return 0, err
	}
	if m.started && k <= m.last {
		d.off = start
		return 0, d.fail("map key %d after key %d: keys must be strictly ascending", k, m.last)
	}
	m.left--
	m.last, m.started = k, true
	return k, nil
}

// EndMap closes the innermost open map, all of whose entries must have been
// read.
func (d *Decoder) EndMap() error {
	if d.err != nil {
		return d.err
	}
	if len(d.maps) == 0 {
		return d.fail("end of a map outside a map")
	}
	if left := d.maps[len(d.maps)-1].left; left != 0 {
		return d.fail("%d map entries not read", left)
	}
	d.maps = d.maps[:len(d.maps)-1]
	return nil
}

// Array reads the head of an array of at most max items and returns the
// number of items, which the caller then reads.
func (d *Decoder) Array(max int) (int, error) {
	n, err := d.expect(majorArray)
	if err != nil {
		return 0, err
	}
	if max < 0 || n > uint64(max) {
		return 0, d.fail("array of %d items, at most %d", n, max)
	}
	if n > uint64(d.remaining()) {
		return 0, d.fail("truncated input: array of %d items", n)
	}
	return int(n), nil
}

// Uint reads an unsigned integer of at most max.
func (d *Decoder) Uint(max uint64) (uint64, error) {
	v, err := d.expect(majorUint)
	if err != nil {
		return 0, err
	}
	if v > max {
		return 0, d.fail("unsigned integer %d above %d", v, max)
	}
	return v, nil
}

// content reads a string of major type want and returns its content, a
// subslice of the input. The length is checked against the remaining input
// and then against min and max.
func (d *Decoder) content(want byte, min, max int) ([]byte, error) {
	n, err := d.expect(want)
	if err != nil {
		return nil, err
	}
	if n > uint64(d.remaining()) {
		return nil, d.fail("truncated input: %s of %d bytes", majorNames[want], n)
	}
	if int(n) < min || int(n) > max {
		return nil, d.fail("%s of %d bytes outside %d..%d", majorNames[want], n, min, max)
	}
	b := d.in[d.off : d.off+int(n)]
	d.off += int(n)
	return b, nil
}

// Bstr reads a byte string of min to max bytes and returns a copy of its
// content. The length is checked before anything is copied.
func (d *Decoder) Bstr(min, max int) ([]byte, error) {
	b, err := d.content(majorBytes, min, max)
	if err != nil {
		return nil, err
	}
	return bytes.Clone(b), nil
}

// Text reads a text string of at most max bytes of valid UTF-8.
func (d *Decoder) Text(max int) (string, error) {
	start := d.off
	b, err := d.content(majorText, 0, max)
	if err != nil {
		return "", err
	}
	if !utf8.Valid(b) {
		d.off = start
		return "", d.fail("text string is not valid UTF-8")
	}
	return string(b), nil
}

// Done checks that every map was closed and that no byte follows the data
// item.
func (d *Decoder) Done() error {
	if d.err != nil {
		return d.err
	}
	if len(d.maps) != 0 {
		return d.fail("%d maps not closed", len(d.maps))
	}
	if d.off != len(d.in) {
		return d.fail("%d trailing bytes", len(d.in)-d.off)
	}
	return nil
}

// next returns the major type of the next data item, or majorUint at the end
// of the input, where reading it reports the truncation.
func (d *Decoder) next() byte {
	if d.off >= len(d.in) {
		return majorUint
	}
	return d.in[d.off] >> 5
}

// ---------------------------------------------------------------------------
// Objects

// Unmarshal decodes one object from in with decode, checks that the whole
// input was read, and re-encodes the decoded value with encode: the result
// must reproduce in byte for byte. decode and encode are the two halves of
// one schema and work on the same value.
//
// Errors of the Decoder and of the re-encoding wrap ErrNonCanonicalCBOR; any
// other error of decode is returned as it is.
//
// The re-encoding equals in on success, so it may hold secrets such as
// I_PAYLOAD or access_material; it is wiped on every path, and so is every
// buffer the Encoder outgrows.
func Unmarshal(in []byte, decode func(*Decoder) error, encode func(*Encoder)) error {
	d := NewDecoder(in)
	if err := decode(d); err != nil {
		return err
	}
	if err := d.Done(); err != nil {
		return err
	}
	e := Encoder{buf: make([]byte, 0, len(in))}
	encode(&e)
	re, err := e.Out()
	defer clear(re)
	if err != nil || !bytes.Equal(re, in) {
		return errorf("input is not the deterministic encoding of its value")
	}
	return nil
}

// Peek reads the type tag (key 0, a text string of at most MaxTypeTagLen
// bytes) and the schema version (key 1, an unsigned integer) of the map at the
// start of in, before strict decoding, so that an unknown schema version is
// reported as such (spec §70). The map must start with keys 0 and 1, in the
// profile; nothing after them is read. The result must never be used as the
// decoded object.
func Peek(in []byte) (typeTag string, version uint64, err error) {
	d := NewDecoder(in)
	// The input bounds the map.
	pairs, err := d.Map(math.MaxInt)
	if err != nil {
		return "", 0, err
	}
	if pairs < 2 {
		return "", 0, d.fail("map without a type tag and a schema version")
	}
	for want := range uint64(2) {
		k, err := d.Key()
		if err != nil {
			return "", 0, err
		}
		if k != want {
			return "", 0, d.fail("map key %d where key %d was expected", k, want)
		}
		if want == 0 {
			typeTag, err = d.Text(MaxTypeTagLen)
		} else {
			version, err = d.Uint(MaxSafeUint)
		}
		if err != nil {
			return "", 0, err
		}
	}
	return typeTag, version, nil
}

// CheckSchema reads the type tag and the schema version of an object with
// Peek and requires the expected values. A different type tag is
// ErrNonCanonicalCBOR; a different version is ErrUnsupportedVersion.
func CheckSchema(in []byte, typeTag string, version uint64) error {
	tag, v, err := Peek(in)
	if err != nil {
		return err
	}
	if tag != typeTag {
		return errorf("type %q, want %q", tag, typeTag)
	}
	if v != version {
		return fmt.Errorf("codec: %s schema version %d, want %d: %w", typeTag, v, version, datekeys.ErrUnsupportedVersion)
	}
	return nil
}

// Walk checks that in is exactly one data item of the profile, with
// containers nested at most maxDepth deep (a scalar has depth 0) and every
// string and container at most maxLen long. It reads iteratively, so deep
// input cannot exhaust the stack.
//
// Walk is a helper for vectors, fuzzing and diagnostics, and for registered
// extensions whose data is CBOR (spec §72). It never decides whether an
// object of the protocol is valid: the decoder of its schema does.
func Walk(in []byte, maxDepth, maxLen int) error {
	d := NewDecoder(in)
	var open []walkLevel
	for first := true; first || len(open) > 0; first = false {
		var err error
		if n := len(open); n > 0 && open[n-1].left == 0 {
			// The innermost container is complete.
			if open[n-1].isMap {
				err = d.EndMap()
			}
			open = open[:n-1]
		} else {
			if n > 0 {
				open[n-1].left--
				if open[n-1].isMap {
					_, err = d.Key()
				}
			}
			if err == nil {
				open, err = d.walkItem(open, maxDepth, maxLen)
			}
		}
		if err != nil {
			return err
		}
	}
	return d.Done()
}

// walkLevel is an open container of Walk.
type walkLevel struct {
	left  int // entries of a map or items of an array not read yet
	isMap bool
}

// walkItem reads one data item for Walk: a scalar, or the head of a
// container, which it pushes on open.
func (d *Decoder) walkItem(open []walkLevel, maxDepth, maxLen int) ([]walkLevel, error) {
	var err error
	switch major := d.next(); major {
	case majorMap, majorArray:
		if len(open) >= maxDepth {
			return open, d.fail("containers nested deeper than %d", maxDepth)
		}
		l := walkLevel{isMap: major == majorMap}
		if l.isMap {
			l.left, err = d.Map(maxLen)
		} else {
			l.left, err = d.Array(maxLen)
		}
		open = append(open, l)
	case majorBytes:
		_, err = d.content(majorBytes, 0, maxLen)
	case majorText:
		_, err = d.Text(maxLen)
	default:
		// An unsigned integer, or the error of whatever is there.
		_, err = d.Uint(math.MaxUint64)
	}
	return open, err
}