// Package extension implements the single generic extension mechanism shared
// by PUBLIC_HEADER, CONTROL_CBOR and .dkk (spec §31, §44, §54, §72).
//
// Extension data is opaque bytes: the base protocol never decodes or
// validates its content, and the validity of the containing object never
// depends on it. The package enforces the structural rules only: valid UTF-8
// identifiers, extension_version at most 2^32-1, data that is absent or a
// non-empty byte string, 1 to 64 extensions per array, no identifier repeated
// within an object, no identifier in both the critical and the noncritical
// array, canonical order by the UTF-8 bytes of extension_id, rejection of
// unknown critical extensions, and omission of empty arrays (spec §58.1).
//
// Only an application that knows an extension interprets its data. A
// Registry that also implements DataValidator checks the data of the
// extensions it knows: invalid data rejects a critical extension with
// ErrExtensionDataInvalid and makes a noncritical one Unusable (spec §54).
// A Registry that also implements Placement tells in which objects and
// arrays each extension is registered: elsewhere a known extension is
// treated as unknown (spec §54, §72).
package extension

import (
	"bytes"
	"fmt"
	"slices"
	"strings"
	"unicode/utf8"

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

// Limits of one extension array and of one extension (spec §31, §54, §57).
const (
	// MaxIDLen bounds extension_id. It is an implementation limit (spec §74).
	MaxIDLen = 256
	// MaxExtensions is the largest number of extensions in one array.
	MaxExtensions = 64
	// MaxVersion is the largest extension_version, 2^32-1.
	MaxVersion = 1<<32 - 1
	// MaxDataLen is the largest data: the largest frame of spec §57,
	// SEALED_CONTROL. The frame of the containing object is the effective
	// bound.
	MaxDataLen = 64 << 20
)

// Extension is one entry of an extension array.
type Extension struct {
	ID      string // key 0, extension_id
	Version uint64 // key 1, extension_version
	// Data is the opaque content of key 2, at least one byte, or nil when the
	// extension carries no data and key 2 is omitted. An empty non-nil slice
	// is invalid: an empty byte string never stands for absence (spec §58.1).
	Data []byte
}

// New returns an extension that carries data, of which it keeps a copy. data
// must hold at least one byte. An extension without data has no constructor:
// it is the literal Extension{ID: id, Version: version}, which omits key 2.
func New(id string, version uint64, data []byte) (Extension, error) {
	if data == nil {
		return Extension{}, fmt.Errorf("extension %q: New needs data; an extension without data is Extension{ID, Version}: %w", id, datekeys.ErrNonCanonicalCBOR)
	}
	e := Extension{ID: id, Version: version, Data: bytes.Clone(data)}
	if err := validate(e); err != nil {
		return Extension{}, err
	}
	return e, nil
}

// Registry tells which extensions the application implements. A nil Registry
// knows none, which is the state of the base protocol V1.
type Registry interface {
	Known(id string, version uint64) bool
}

// DataValidator is an optional interface of a Registry. ValidateData reports
// whether the data of e (nil when e carries none) follows the registered
// schema of (e.ID, e.Version) (spec §72). It is called only for extensions
// the Registry knows.
type DataValidator interface {
	ValidateData(e Extension) error
}

// Object names an object that carries extension arrays (spec §54).
type Object int

// The objects that carry extensions.
const (
	PublicHeader Object = iota + 1 // PUBLIC_HEADER, keys 5 and 6 (spec §24)
	Control                        // CONTROL_CBOR, keys 4 and 5 (spec §31)
	AccessKey                      // the body of a .dkk, keys 7 and 8 (spec §41)
	Head                           // the head of a format 3 capsule, keys 6 and 7 (spec §29.4)
)

// String returns the name of the object in the specification.
func (o Object) String() string {
	switch o {
	case PublicHeader:
		return "PUBLIC_HEADER"
	case Control:
		return "CONTROL_CBOR"
	case AccessKey:
		return ".dkk"
	case Head:
		return "head"
	}
	return fmt.Sprintf("Object(%d)", int(o))
}

// Array names one of the two extension arrays of an object.
type Array int

// The two extension arrays of an object.
const (
	Critical    Array = iota + 1 // critical_extensions
	Noncritical                  // noncritical_extensions
)

// String returns the name of the array in the specification.
func (a Array) String() string {
	switch a {
	case Critical:
		return "critical_extensions"
	case Noncritical:
		return "noncritical_extensions"
	}
	return fmt.Sprintf("Array(%d)", int(a))
}

// Placement is an optional interface of a Registry. RegisteredIn reports
// whether (id, version) is registered for the array arr of the object obj:
// the registration of each extension declares the objects and arrays where
// it may appear (spec §72). It is called only for extensions the Registry
// knows. A known extension that appears in an object or array it is not
// registered for is treated there as unknown (spec §54): a critical one is
// rejected with ErrExtensionCriticalUnknown and a noncritical one is ignored,
// its data neither checked nor interpreted. A Registry that does not
// implement Placement knows each of its extensions in every object and array.
//
// An encoder must not write a registered extension where it is not
// registered (spec §72): CheckWrite is that rule, which the writers of this
// module apply with the Registry of the extensions that the specification
// itself registers.
type Placement interface {
	RegisteredIn(id string, version uint64, obj Object, arr Array) bool
}

// CheckWrite applies the rules of an encoder of spec §72 to the extensions
// exts that it writes in the array arr of obj: an extension that reg knows
// goes only where reg registers it, and with data that reg validates when it
// is a DataValidator. The extensions that reg does not know are the
// application's own, and the application answers for them.
func CheckWrite(reg Registry, obj Object, arr Array, exts []Extension) error {
	if reg == nil {
		return nil
	}
	for _, e := range exts {
		if !reg.Known(e.ID, e.Version) {
			continue
		}
		if !registered(reg, e.ID, e.Version, obj, arr) {
			return fmt.Errorf("extension: %s version %d is not registered for %s of %s: an encoder must not write it there (spec §72)", e.ID, e.Version, arr, obj)
		}
		if v, ok := reg.(DataValidator); ok {
			if err := v.ValidateData(e); err != nil {
				return fmt.Errorf("extension: %s version %d: %w", e.ID, e.Version, err)
			}
		}
	}
	return nil
}

// KnownIn reports whether reg knows (id, version) in the array arr of obj:
// reg knows it and, when reg is a Placement, registers it there (spec §54,
// §72). A nil Registry knows none. It is the rule of CheckCriticalIn and
// CheckNoncriticalIn, for an application that interprets the extensions of
// an object it has read: an extension unknown there is not interpreted.
func KnownIn(reg Registry, id string, version uint64, obj Object, arr Array) bool {
	return reg != nil && reg.Known(id, version) && registered(reg, id, version, obj, arr)
}

// registered reports whether reg, which knows (id, version), registers it for
// the array arr of obj: always when reg is not a Placement.
func registered(reg Registry, id string, version uint64, obj Object, arr Array) bool {
	p, ok := reg.(Placement)
	return !ok || p.RegisteredIn(id, version, obj, arr)
}

// Set is a simple Registry. It does not validate data, and it knows each of
// its extensions in every object and array.
type Set map[string][]uint64

// Known reports whether (id, version) is in the set.
func (s Set) Known(id string, version uint64) bool { return slices.Contains(s[id], version) }

// compare orders extensions by the UTF-8 bytes of extension_id, the order of
// spec §54: strings.Compare compares the bytes as unsigned values, and a
// proper prefix sorts first. Within one object an identifier appears at most
// once.
func compare(a, b Extension) int { return strings.Compare(a.ID, b.ID) }

func validate(e Extension) error {
	if e.ID == "" || len(e.ID) > MaxIDLen || !utf8.ValidString(e.ID) {
		return fmt.Errorf("extension: invalid extension_id %q: %w", e.ID, datekeys.ErrNonCanonicalCBOR)
	}
	if e.Version > MaxVersion {
		return fmt.Errorf("extension %s: extension_version %d exceeds %d: %w", e.ID, e.Version, uint64(MaxVersion), datekeys.ErrNonCanonicalCBOR)
	}
	if e.Data != nil && (len(e.Data) == 0 || len(e.Data) > MaxDataLen) {
		return fmt.Errorf("extension %s: data of %d bytes outside 1..%d: %w", e.ID, len(e.Data), MaxDataLen, datekeys.ErrNonCanonicalCBOR)
	}
	return nil
}

// Canonical validates one extension array and returns it in canonical order,
// sorted by the UTF-8 bytes of extension_id: 1 to 64 valid extensions, no
// identifier repeated. An empty input yields nil, so that the array key is
// omitted (spec §58.1).
func Canonical(exts []Extension) ([]Extension, error) {
	if len(exts) == 0 {
		return nil, nil
	}
	if len(exts) > MaxExtensions {
		return nil, errTooMany(len(exts))
	}
	out := slices.Clone(exts)
	slices.SortFunc(out, compare)
	if err := checkArray(out); err != nil {
		return nil, err
	}
	return out, nil
}

func errTooMany(n int) error {
	return fmt.Errorf("extension: %d extensions in one array, at most %d: %w", n, MaxExtensions, datekeys.ErrNonCanonicalCBOR)
}

// checkArray applies the rules of DecodeArray to an array to be written: 1 to
// 64 valid extensions in canonical order, no identifier repeated.
func checkArray(exts []Extension) error {
	switch {
	case len(exts) == 0:
		return fmt.Errorf("extension: empty array; an absent array omits its key: %w", datekeys.ErrNonCanonicalCBOR)
	case len(exts) > MaxExtensions:
		return errTooMany(len(exts))
	}
	for i, e := range exts {
		if err := validate(e); err != nil {
			return err
		}
		if i == 0 {
			continue
		}
		switch c := compare(exts[i-1], e); {
		case c == 0:
			return fmt.Errorf("extension %s: appears more than once: %w", e.ID, datekeys.ErrNonCanonicalCBOR)
		case c > 0:
			return fmt.Errorf("extension %s: array is not in canonical order: %w", e.ID, datekeys.ErrNonCanonicalCBOR)
		}
	}
	return nil
}

// EncodeArray writes a non-empty extension array as Canonical returns it:
// each extension is the map {0: extension_id, 1: extension_version} with
// key 2, the data as a byte string, only when the extension carries data
// (spec §54). An array that DecodeArray would reject is not written: its
// error is recorded in e, whose Out returns it.
func EncodeArray(e *codec.Encoder, exts []Extension) {
	if err := checkArray(exts); err != nil {
		e.Fail(err)
		return
	}
	e.Array(len(exts))
	for _, x := range exts {
		if x.Data == nil {
			e.Map(2)
		} else {
			e.Map(3)
		}
		e.Uint(0)
		e.Text(x.ID)
		e.Uint(1)
		e.Uint(x.Version)
		if x.Data != nil {
			e.Uint(2)
			e.Bstr(x.Data)
		}
	}
}

// DecodeArray reads one extension array. The array holds 1 to 64 entries,
// which its head declares before any is read; each entry is a map with
// key 0, a non-empty UTF-8 extension_id of at most MaxIDLen bytes, key 1, an
// extension_version of at most MaxVersion, and optionally key 2, a byte
// string of at least one byte whose content is copied and never decoded
// (spec §54, §58.1). Entries are in canonical order with no identifier
// repeated. Every failure wraps ErrNonCanonicalCBOR.
func DecodeArray(d *codec.Decoder) ([]Extension, error) {
	n, err := d.Array(MaxExtensions)
	if err != nil {
		return nil, fmt.Errorf("extension: %w", err)
	}
	if n == 0 {
		return nil, fmt.Errorf("extension: empty array; an absent array omits its key: %w", datekeys.ErrNonCanonicalCBOR)
	}
	out := make([]Extension, 0, n)
	for i := range n {
		e, err := decodeOne(d)
		if err != nil {
			return nil, err
		}
		if i > 0 {
			switch c := compare(out[i-1], e); {
			case c == 0:
				return nil, fmt.Errorf("extension %s: appears more than once: %w", e.ID, datekeys.ErrNonCanonicalCBOR)
			case c > 0:
				return nil, fmt.Errorf("extension %s: array is not in canonical order: %w", e.ID, datekeys.ErrNonCanonicalCBOR)
			}
		}
		out = append(out, e)
	}
	return out, nil
}

// decodeOne reads the map of one extension. Key 2, when present, must be a
// byte string of at least one byte: the empty byte string and every other
// CBOR type are rejected explicitly (spec §54, §58.1).
func decodeOne(d *codec.Decoder) (Extension, error) {
	var e Extension
	pairs, err := d.Map(3)
	if err != nil {
		return e, fmt.Errorf("extension: %w", err)
	}
	var seen [3]bool
	for range pairs {
		k, err := d.Key()
		if err != nil {
			return e, fmt.Errorf("extension: %w", err)
		}
		switch k {
		case 0:
			e.ID, err = d.Text(MaxIDLen)
		case 1:
			e.Version, err = d.Uint(MaxVersion)
		case 2:
			if e.Data, err = d.Bstr(0, MaxDataLen); err == nil && len(e.Data) == 0 {
				return e, fmt.Errorf("extension %q: data is present but empty; an extension without data omits key 2: %w", e.ID, datekeys.ErrNonCanonicalCBOR)
			}
		default:
			return e, fmt.Errorf("extension %q: unknown key %d: %w", e.ID, k, datekeys.ErrNonCanonicalCBOR)
		}
		if err != nil {
			return e, fmt.Errorf("extension %q: key %d: %w", e.ID, k, err)
		}
		seen[k] = true
	}
	if !seen[0] || !seen[1] {
		return e, fmt.Errorf("extension %q: extension_id and extension_version are required: %w", e.ID, datekeys.ErrNonCanonicalCBOR)
	}
	if err := validate(e); err != nil {
		return e, err
	}
	return e, d.EndMap()
}

// CheckDisjoint applies the cross-array rule of one object: an extension_id
// must not appear in both critical_extensions and noncritical_extensions
// (spec §31, §54). Arrays in canonical order, as Canonical and DecodeArray
// return them, are merged in one linear pass; other input is sorted first.
func CheckDisjoint(critical, noncritical []Extension) error {
	critical, noncritical = sorted(critical), sorted(noncritical)
	for i, j := 0, 0; i < len(critical) && j < len(noncritical); {
		switch c := compare(critical[i], noncritical[j]); {
		case c == 0:
			return fmt.Errorf("extension %s: both critical and noncritical: %w", critical[i].ID, datekeys.ErrNonCanonicalCBOR)
		case c < 0:
			i++
		default:
			j++
		}
	}
	return nil
}

// sorted returns exts itself when it is in canonical order, and a sorted copy
// otherwise.
func sorted(exts []Extension) []Extension {
	if slices.IsSortedFunc(exts, compare) {
		return exts
	}
	sorted := slices.Clone(exts)
	slices.SortFunc(sorted, compare)
	return sorted
}

// CheckCritical rejects every critical extension unknown to reg with
// ErrExtensionCriticalUnknown and, when reg is a DataValidator, every known
// one whose data it rejects with ErrExtensionDataInvalid (spec §54, §70).
// An unknown extension takes precedence over invalid data. It does not know
// the object of the array and consults no Placement: CheckCriticalIn does.
func CheckCritical(critical []Extension, reg Registry) error {
	return checkCritical(critical, reg, 0)
}

// CheckCriticalIn is CheckCritical for the critical_extensions of obj: an
// extension that reg knows but, as a Placement, does not register for that
// array of obj is unknown there, ErrExtensionCriticalUnknown (spec §54,
// §72). The reading flow of package capsule checks every critical array
// with it.
func CheckCriticalIn(obj Object, critical []Extension, reg Registry) error {
	return checkCritical(critical, reg, obj)
}

// checkCritical checks the critical extensions of obj, or of any object when
// obj is 0: every unknown one first, then the data of the known ones.
func checkCritical(critical []Extension, reg Registry, obj Object) error {
	for _, c := range critical {
		if reg == nil || !reg.Known(c.ID, c.Version) {
			return fmt.Errorf("extension %s v%d: %w", c.ID, c.Version, datekeys.ErrExtensionCriticalUnknown)
		}
		if obj != 0 && !registered(reg, c.ID, c.Version, obj, Critical) {
			return fmt.Errorf("extension %s v%d: known, but not registered for the %s of %s: %w", c.ID, c.Version, Critical, obj, datekeys.ErrExtensionCriticalUnknown)
		}
	}
	for _, c := range critical {
		if err := validateData(c, reg); err != nil {
			return err
		}
	}
	return nil
}

// Unusable is a known noncritical extension whose data does not follow its
// registered schema. The object that carries it stays valid; the application
// must not use the extension, and the caller is told (spec §54).
type Unusable struct {
	ID      string
	Version uint64
	Err     error // wraps datekeys.ErrExtensionDataInvalid
}

// CheckNoncritical returns the noncritical extensions that reg knows and whose
// data it rejects. It never fails the object: unknown noncritical extensions
// are ignored, and a Registry that is not a DataValidator rejects no data
// (spec §54). It does not know the object of the array and consults no
// Placement: CheckNoncriticalIn does.
func CheckNoncritical(noncritical []Extension, reg Registry) []Unusable {
	return checkNoncritical(noncritical, reg, 0)
}

// CheckNoncriticalIn is CheckNoncritical for the noncritical_extensions of
// obj: an extension that reg knows but, as a Placement, does not register
// for that array of obj is unknown there and ignored, its data unchecked
// (spec §54, §72). The reading flow of package capsule checks every
// noncritical array with it.
func CheckNoncriticalIn(obj Object, noncritical []Extension, reg Registry) []Unusable {
	return checkNoncritical(noncritical, reg, obj)
}

// checkNoncritical checks the noncritical extensions of obj, or of any
// object when obj is 0.
func checkNoncritical(noncritical []Extension, reg Registry, obj Object) []Unusable {
	var out []Unusable
	for _, n := range noncritical {
		if reg == nil || !reg.Known(n.ID, n.Version) || obj != 0 && !registered(reg, n.ID, n.Version, obj, Noncritical) {
			continue
		}
		if err := validateData(n, reg); err != nil {
			out = append(out, Unusable{ID: n.ID, Version: n.Version, Err: err})
		}
	}
	return out
}

// validateData applies the optional DataValidator of reg to a known
// extension. The validator's own error is kept as text only, so that the
// result carries exactly one normative code.
func validateData(e Extension, reg Registry) error {
	v, ok := reg.(DataValidator)
	if !ok {
		return nil
	}
	if err := v.ValidateData(e); err != nil {
		return fmt.Errorf("extension %s v%d: data: %v: %w", e.ID, e.Version, err, datekeys.ErrExtensionDataInvalid)
	}
	return nil
}