// Package axon implements the Axon protocol core: event signing, verification,
// canonical payload construction, and related helpers.
package axon

import (
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"fmt"
	"sort"
)

// Event kind constants matching the registry in PROTOCOL.md.
const (
	KindProfile     uint16 = 0    // Identity metadata
	KindMessage     uint16 = 1000 // Plain text note
	KindDM          uint16 = 2000 // Encrypted direct message
	KindProgress    uint16 = 3000 // Ephemeral progress/status indicator
	KindJobRequest  uint16 = 5000 // Request for agent work
	KindJobResult   uint16 = 6000 // Completed job output
	KindJobFeedback uint16 = 7000 // In-progress status / error
)

// Tag is a named list of string values attached to an Event.
type Tag struct {
	Name   string   `msgpack:"name" json:"name"`
	Values []string `msgpack:"values" json:"values"`
}

// TagFilter selects events that have a tag with the given name and any of the
// given values. An empty Values slice matches any value.
type TagFilter struct {
	Name   string   `msgpack:"name"`
	Values []string `msgpack:"values"`
}

// Filter selects a subset of events. All non-empty fields are ANDed together;
// multiple entries within a slice field are ORed.
//
// IDs and Authors support prefix matching: a []byte shorter than 32 bytes
// matches any event whose ID (or pubkey) starts with those bytes.
type Filter struct {
	IDs     [][]byte    `msgpack:"ids"`
	Authors [][]byte    `msgpack:"authors"`
	Kinds   []uint16    `msgpack:"kinds"`
	Since   int64       `msgpack:"since"` // inclusive lower bound on created_at
	Until   int64       `msgpack:"until"` // inclusive upper bound on created_at
	Limit   int32       `msgpack:"limit"` // max events to return (0 = no limit)
	Tags    []TagFilter `msgpack:"tags"`
}

// Event is the core Axon data structure. All fields use their wire types.
// id, pubkey and sig are raw 32/64-byte slices, not hex.
// content is opaque bytes (msgpack bin type).
type Event struct {
	ID        []byte `msgpack:"id"`         // 32 bytes, SHA256 of canonical payload
	PubKey    []byte `msgpack:"pubkey"`     // 32 bytes, Ed25519 public key
	CreatedAt int64  `msgpack:"created_at"` // Unix timestamp
	Kind      uint16 `msgpack:"kind"`
	Content   []byte `msgpack:"content"` // opaque; msgpack bin type
	Sig       []byte `msgpack:"sig"`     // 64 bytes, Ed25519 signature over id
	Tags      []Tag  `msgpack:"tags"`
}

// KeyPair holds an Ed25519 private and public key.
type KeyPair struct {
	PrivKey ed25519.PrivateKey
	PubKey  ed25519.PublicKey
}

// NewKeyPair generates a fresh Ed25519 keypair.
func NewKeyPair() (KeyPair, error) {
	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return KeyPair{}, fmt.Errorf("axon: generate key: %w", err)
	}
	return KeyPair{PrivKey: priv, PubKey: pub}, nil
}

// NewKeyPairFromSeed derives a keypair from a 32-byte seed (the canonical
// private key representation). Panics if seed is not 32 bytes.
func NewKeyPairFromSeed(seed []byte) KeyPair {
	if len(seed) != ed25519.SeedSize {
		panic(fmt.Sprintf("axon: seed must be %d bytes, got %d", ed25519.SeedSize, len(seed)))
	}
	priv := ed25519.NewKeyFromSeed(seed)
	return KeyPair{PrivKey: priv, PubKey: priv.Public().(ed25519.PublicKey)}
}

// CanonicalTags encodes tags into their canonical binary representation and
// returns the raw bytes (before hashing). This is exposed so callers can
// inspect the encoding in tests; normally you want CanonicalTagsHash.
//
// Encoding:
//
//	uint16(num_tags)
//	for each tag (sorted by name, then first value for ties):
//	  uint16(len(name)) || utf8(name)
//	  uint16(num_values)
//	  for each value:
//	    uint32(len(value)) || utf8(value)
//
// Returns an error if two tags share the same name and first value (protocol error).
func CanonicalTags(tags []Tag) ([]byte, error) {
	// Sort a copy so we don't mutate the caller's slice.
	sorted := make([]Tag, len(tags))
	copy(sorted, tags)
	sort.SliceStable(sorted, func(i, j int) bool {
		if sorted[i].Name != sorted[j].Name {
			return sorted[i].Name < sorted[j].Name
		}
		vi := ""
		if len(sorted[i].Values) > 0 {
			vi = sorted[i].Values[0]
		}
		vj := ""
		if len(sorted[j].Values) > 0 {
			vj = sorted[j].Values[0]
		}
		return vi < vj
	})

	// Detect duplicates (same name + same first value).
	for i := 1; i < len(sorted); i++ {
		prev, cur := sorted[i-1], sorted[i]
		if prev.Name != cur.Name {
			continue
		}
		prevFirst := ""
		if len(prev.Values) > 0 {
			prevFirst = prev.Values[0]
		}
		curFirst := ""
		if len(cur.Values) > 0 {
			curFirst = cur.Values[0]
		}
		if prevFirst == curFirst {
			return nil, fmt.Errorf("axon: duplicate tag (name=%q first_value=%q)", cur.Name, curFirst)
		}
	}

	// Estimate capacity to avoid repeated allocations.
	buf := make([]byte, 0, 2+len(sorted)*16)
	var hdr [4]byte

	binary.BigEndian.PutUint16(hdr[:2], uint16(len(sorted)))
	buf = append(buf, hdr[:2]...)

	for _, tag := range sorted {
		name := []byte(tag.Name)
		binary.BigEndian.PutUint16(hdr[:2], uint16(len(name)))
		buf = append(buf, hdr[:2]...)
		buf = append(buf, name...)

		binary.BigEndian.PutUint16(hdr[:2], uint16(len(tag.Values)))
		buf = append(buf, hdr[:2]...)

		for _, v := range tag.Values {
			vb := []byte(v)
			binary.BigEndian.PutUint32(hdr[:4], uint32(len(vb)))
			buf = append(buf, hdr[:4]...)
			buf = append(buf, vb...)
		}
	}

	return buf, nil
}

// CanonicalTagsHash returns SHA256(canonical_tags encoding).
func CanonicalTagsHash(tags []Tag) ([32]byte, error) {
	enc, err := CanonicalTags(tags)
	if err != nil {
		return [32]byte{}, err
	}
	return sha256.Sum256(enc), nil
}

// CanonicalPayload constructs the deterministic byte payload that is hashed to
// produce the event ID.
//
// Layout:
//
//	[0:2]       uint16 = 32          pubkey length (always 32)
//	[2:34]      bytes                pubkey
//	[34:42]     uint64               created_at
//	[42:44]     uint16               kind
//	[44:48]     uint32               content length
//	[48:48+n]   bytes                content
//	[48+n:80+n] bytes                SHA256(canonical_tags), 32 bytes
func CanonicalPayload(pubkey []byte, createdAt int64, kind uint16, content []byte, tags []Tag) ([]byte, error) {
	if len(pubkey) != 32 {
		return nil, fmt.Errorf("axon: pubkey must be 32 bytes, got %d", len(pubkey))
	}
	if len(content) > 65536 {
		return nil, errors.New("axon: content exceeds 65536 byte limit")
	}

	tagsHash, err := CanonicalTagsHash(tags)
	if err != nil {
		return nil, err
	}

	n := len(content)
	// Total size: 2 + 32 + 8 + 2 + 4 + n + 32
	payload := make([]byte, 80+n)

	binary.BigEndian.PutUint16(payload[0:2], 32)
	copy(payload[2:34], pubkey)
	binary.BigEndian.PutUint64(payload[34:42], uint64(createdAt))
	binary.BigEndian.PutUint16(payload[42:44], kind)
	binary.BigEndian.PutUint32(payload[44:48], uint32(n))
	copy(payload[48:48+n], content)
	copy(payload[48+n:80+n], tagsHash[:])

	return payload, nil
}

// EventID computes the canonical SHA256 event ID for the given fields.
func EventID(pubkey []byte, createdAt int64, kind uint16, content []byte, tags []Tag) ([]byte, error) {
	payload, err := CanonicalPayload(pubkey, createdAt, kind, content, tags)
	if err != nil {
		return nil, err
	}
	h := sha256.Sum256(payload)
	return h[:], nil
}

// Sign fills in e.ID and e.Sig using kp. It also sets e.PubKey from kp.
// The caller should populate all other fields (CreatedAt, Kind, Content, Tags)
// before calling Sign.
func Sign(e *Event, kp KeyPair) error {
	e.PubKey = []byte(kp.PubKey)

	id, err := EventID(e.PubKey, e.CreatedAt, e.Kind, e.Content, e.Tags)
	if err != nil {
		return err
	}
	e.ID = id
	e.Sig = ed25519.Sign(kp.PrivKey, id)
	return nil
}

// Verify checks that e.Sig is a valid Ed25519 signature of e.ID using e.PubKey,
// and that e.ID matches the canonical payload derived from the event fields.
// Returns nil if both checks pass.
func Verify(e *Event) error {
	if len(e.PubKey) != 32 {
		return fmt.Errorf("axon: pubkey must be 32 bytes, got %d", len(e.PubKey))
	}
	if len(e.Sig) != ed25519.SignatureSize {
		return fmt.Errorf("axon: sig must be %d bytes, got %d", ed25519.SignatureSize, len(e.Sig))
	}

	expectedID, err := EventID(e.PubKey, e.CreatedAt, e.Kind, e.Content, e.Tags)
	if err != nil {
		return fmt.Errorf("axon: compute expected id: %w", err)
	}
	if len(e.ID) != 32 {
		return fmt.Errorf("axon: id must be 32 bytes, got %d", len(e.ID))
	}
	for i := range expectedID {
		if expectedID[i] != e.ID[i] {
			return errors.New("axon: event id does not match canonical payload")
		}
	}

	pub := ed25519.PublicKey(e.PubKey)
	if !ed25519.Verify(pub, e.ID, e.Sig) {
		return errors.New("axon: invalid signature")
	}
	return nil
}

// SignChallenge signs the relay authentication challenge:
//
//	sig = ed25519.Sign(privkey, SHA256(nonce || utf8(relay_url)))
func SignChallenge(kp KeyPair, nonce []byte, relayURL string) []byte {
	h := sha256.New()
	h.Write(nonce)
	h.Write([]byte(relayURL))
	digest := h.Sum(nil)
	return ed25519.Sign(kp.PrivKey, digest)
}

// VerifyChallenge verifies a challenge signature.
func VerifyChallenge(pubkey []byte, nonce []byte, relayURL string, sig []byte) bool {
	h := sha256.New()
	h.Write(nonce)
	h.Write([]byte(relayURL))
	digest := h.Sum(nil)
	return ed25519.Verify(ed25519.PublicKey(pubkey), digest, sig)
}