// Package extension implements the single generic extension mechanism shared
// by PUBLIC_HEADER, CONTROL_CBOR and .dkk (spec §31, §44, §54, §72).
//
// The base protocol does not interpret extension data. It enforces the
// structural rules only: valid UTF-8 identifiers, no identifier repeated
// within an object (V1 registers no schema that allows multiplicity), no
// identifier in both the critical and the noncritical array, canonical order
// by the UTF-8 bytes of extension_id and then by version, rejection of unknown
// critical extensions, and omission of empty arrays (spec §58.1).
package extension

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

	"github.com/fxamacker/cbor/v2"

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

// MaxIDLen bounds extension_id. It is an implementation limit (spec §74).
const MaxIDLen = 256

// 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 Deterministic CBOR encoding of the data item (key 2), or
	// nil when the extension carries no data and the key is omitted.
	Data []byte
}

// New builds an extension whose data is the deterministic encoding of value.
func New(id string, version uint64, value any) (Extension, error) {
	b, err := codec.Marshal(value)
	if err != nil {
		return Extension{}, err
	}
	return Extension{ID: id, Version: version, Data: b}, nil
}

// Wire is the CBOR map of one extension (spec §54).
type Wire struct {
	ID      string          `cbor:"0,keyasint"`
	Version uint64          `cbor:"1,keyasint"`
	Data    cbor.RawMessage `cbor:"2,keyasint,omitempty"`
}

// Registry tells which critical 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
}

// Set is a simple Registry.
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) }

func compare(a, b Extension) int {
	if c := bytes.Compare([]byte(a.ID), []byte(b.ID)); c != 0 {
		return c
	}
	return cmp.Compare(a.Version, b.Version)
}

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.Data != nil {
		if err := codec.Valid(e.Data); err != nil {
			return fmt.Errorf("extension %s: data: %w", e.ID, err)
		}
	}
	return nil
}

// Encode validates one extension array and returns its canonical wire form,
// sorted by extension_id bytes and then version. 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
	}
	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: canonical order, no repeated
// identifier and canonical data.
func Decode(ws []Wire) ([]Extension, error) {
	if len(ws) == 0 {
		return nil, nil
	}
	out := make([]Extension, 0, len(ws))
	for i, w := range ws {
		e := Extension{ID: w.ID, Version: w.Version}
		if w.Data != nil {
			e.Data = bytes.Clone(w.Data)
		}
		if err := validate(e); err != nil {
			return nil, err
		}
		if i > 0 {
			prev := out[i-1]
			if prev.ID == e.ID {
				return nil, fmt.Errorf("extension %s: appears more than once: %w", e.ID, datekeys.ErrNonCanonicalCBOR)
			}
			if compare(prev, e) > 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).
func CheckDisjoint(critical, noncritical []Extension) error {
	for _, c := range critical {
		for _, n := range noncritical {
			if c.ID == n.ID {
				return fmt.Errorf("extension %s: both critical and noncritical: %w", c.ID, datekeys.ErrNonCanonicalCBOR)
			}
		}
	}
	return nil
}

// CheckCritical rejects every critical extension unknown to reg (spec §54,
// §70). Unknown noncritical extensions may be ignored and are not checked.
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)
		}
	}
	return nil
}