// 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).
package extension

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

	"github.com/fxamacker/cbor/v2"

	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
}

// Wire is the CBOR map of one extension (spec §54). Data is the content of the
// byte string at key 2; nil omits the key.
type Wire struct {
	ID      string `cbor:"0,keyasint"`
	Version uint64 `cbor:"1,keyasint"`
	Data    []byte `cbor:"2,keyasint,omitempty"`
}

// UnmarshalCBOR decodes one extension map. 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 here, explicitly, and not left to the re-encoding
// check of the containing object (spec §54, §58.1).
func (w *Wire) UnmarshalCBOR(b []byte) error {
	var raw struct {
		ID      string          `cbor:"0,keyasint"`
		Version uint64          `cbor:"1,keyasint"`
		Data    cbor.RawMessage `cbor:"2,keyasint,omitempty"`
	}
	if err := codec.Unmarshal(b, &raw); err != nil {
		return err
	}
	*w = Wire{ID: raw.ID, Version: raw.Version}
	if len(raw.Data) == 0 {
		return nil
	}
	const majorByteString = 2
	if major := raw.Data[0] >> 5; major != majorByteString {
		return fmt.Errorf("extension %q: data is CBOR major type %d, not a byte string: %w", raw.ID, major, datekeys.ErrNonCanonicalCBOR)
	}
	var data []byte
	if err := codec.Unmarshal(raw.Data, &data); err != nil {
		return fmt.Errorf("extension %q: data: %w", raw.ID, err)
	}
	if len(data) == 0 {
		return fmt.Errorf("extension %q: data is present but empty; an extension without data omits key 2: %w", raw.ID, datekeys.ErrNonCanonicalCBOR)
	}
	w.Data = data
	return 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
}

// Set is a simple Registry. It does not validate data.
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 canonical
// order; 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
}

// Encode validates one extension array and returns its canonical wire form,
// sorted by the UTF-8 bytes of extension_id. An empty input yields nil, so
// that the array key is omitted (spec §58.1).
func Encode(exts []Extension) ([]Wire, error) {
	if len(exts) == 0 {
		return nil, nil
	}
	if len(exts) > MaxExtensions {
		return nil, fmt.Errorf("extension: %d extensions in one array, at most %d: %w", len(exts), MaxExtensions, datekeys.ErrNonCanonicalCBOR)
	}
	sorted := slices.Clone(exts)
	slices.SortFunc(sorted, compare)
	out := make([]Wire, 0, len(sorted))
	for i, e := range sorted {
		if err := validate(e); err != nil {
			return nil, err
		}
		if i > 0 && sorted[i-1].ID == e.ID {
			return nil, fmt.Errorf("extension %s: appears more than once: %w", e.ID, datekeys.ErrNonCanonicalCBOR)
		}
		out = append(out, Wire{ID: e.ID, Version: e.Version, Data: bytes.Clone(e.Data)})
	}
	return out, nil
}

// Decode validates one decoded extension array: at most 64 entries, each one
// valid, in canonical order and with no repeated identifier. The data is
// copied, never decoded.
func Decode(ws []Wire) ([]Extension, error) {
	if len(ws) == 0 {
		return nil, nil
	}
	if len(ws) > MaxExtensions {
		return nil, fmt.Errorf("extension: %d extensions in one array, at most %d: %w", len(ws), MaxExtensions, datekeys.ErrNonCanonicalCBOR)
	}
	out := make([]Extension, 0, len(ws))
	for i, w := range ws {
		e := Extension{ID: w.ID, Version: w.Version}
		if w.Data != nil {
			if len(w.Data) == 0 {
				return nil, fmt.Errorf("extension %q: data is present but empty: %w", w.ID, datekeys.ErrNonCanonicalCBOR)
			}
			e.Data = bytes.Clone(w.Data)
		}
		if err := validate(e); 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
}

// 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 Decode returns them, are
// merged in one linear pass; other input is sorted first.
func CheckDisjoint(critical, noncritical []Extension) error {
	critical, noncritical = canonical(critical), canonical(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
}

// canonical returns exts itself when it is in canonical order, and a sorted
// copy otherwise.
func canonical(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.
func CheckCritical(critical []Extension, reg Registry) 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)
		}
	}
	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).
func CheckNoncritical(noncritical []Extension, reg Registry) []Unusable {
	var out []Unusable
	for _, n := range noncritical {
		if reg == nil || !reg.Known(n.ID, n.Version) {
			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
}