// Package relay implements a minimal NIP-01 Nostr relay.
//
// The relay accepts WebSocket connections from clients, handles EVENT
// (store + broadcast), REQ (subscribe + replay), and CLOSE (unsubscribe)
// messages. Events are stored in memory and matched against subscription
// filters. This is Stage 6 of the developmental plan — the organism
// learns to serve, not just consume.
package relay

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"net"
	"net/http"
	"slices"
	"strings"
	"sync"

	"git.mleku.dev/mleku/dendrite/pkg/nostr"
	"github.com/coder/websocket"
)

// Relay is a NIP-01 relay server with in-memory event storage.
type Relay struct {
	// Name is the relay's display name for NIP-11.
	Name string

	// RequireAuth enables NIP-42 authentication. When true, clients
	// must authenticate before publishing events.
	RequireAuth bool

	// ContentFilter is an optional coherence filter that rejects events
	// whose content doesn't bond to any known lattice structure.
	ContentFilter *CoherenceFilter

	mu     sync.RWMutex
	events map[string]*nostr.Event // id → event
	order  []string               // insertion order for limit queries

	subMu sync.RWMutex
	subs  map[*conn]map[string][]nostr.Filter // conn → subID → filters

	srv *http.Server
	ln  net.Listener
}

// conn wraps a WebSocket connection with a write mutex.
type conn struct {
	ws        *websocket.Conn
	mu        sync.Mutex
	challenge string // NIP-42 auth challenge
	authedAs  string // pubkey of authenticated user (empty if not authed)
}

// New creates a relay with empty storage.
func New(name string) *Relay {
	return &Relay{
		Name:   name,
		events: make(map[string]*nostr.Event),
		subs:   make(map[*conn]map[string][]nostr.Filter),
	}
}

// EventCount returns the number of stored events.
func (r *Relay) EventCount() int {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return len(r.events)
}

// Listen starts the relay on the given address (e.g., "127.0.0.1:0").
// Returns the actual address (useful when port is 0).
func (r *Relay) Listen(addr string) (string, error) {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return "", err
	}
	r.ln = ln

	mux := http.NewServeMux()
	mux.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		// NIP-11: relay information document.
		if req.Header.Get("Accept") == "application/nostr+json" {
			r.handleNIP11(w)
			return
		}
		// WebSocket upgrade.
		ws, err := websocket.Accept(w, req, &websocket.AcceptOptions{
			OriginPatterns: []string{"*"},
		})
		if err != nil {
			return
		}
		r.handleConn(req.Context(), ws)
	})

	r.srv = &http.Server{Handler: mux}
	go r.srv.Serve(ln)

	return ln.Addr().String(), nil
}

// Shutdown stops the relay gracefully.
func (r *Relay) Shutdown(ctx context.Context) error {
	if r.srv != nil {
		return r.srv.Shutdown(ctx)
	}
	return nil
}

// Addr returns the listener address, or "" if not listening.
func (r *Relay) Addr() string {
	if r.ln != nil {
		return r.ln.Addr().String()
	}
	return ""
}

// handleNIP11 serves the NIP-11 relay information document.
func (r *Relay) handleNIP11(w http.ResponseWriter) {
	info := map[string]any{
		"name":           r.Name,
		"description":    "dendrite lattice relay",
		"supported_nips": []int{1, 9, 11, 42},
		"software":       "dendrite",
		"version":        "0.1.0",
	}
	w.Header().Set("Content-Type", "application/nostr+json")
	json.NewEncoder(w).Encode(info)
}

// handleConn manages one WebSocket client session.
func (r *Relay) handleConn(ctx context.Context, ws *websocket.Conn) {
	c := &conn{ws: ws}
	defer func() {
		r.removeSubs(c)
		ws.CloseNow()
	}()

	// NIP-42: send auth challenge if authentication is enabled.
	if r.RequireAuth {
		var b [16]byte
		rand.Read(b[:])
		c.challenge = hex.EncodeToString(b[:])
		r.sendAuth(ctx, c, c.challenge)
	}

	for {
		_, data, err := ws.Read(ctx)
		if err != nil {
			return
		}

		var envelope []json.RawMessage
		if json.Unmarshal(data, &envelope) != nil || len(envelope) < 2 {
			r.sendNotice(ctx, c, "invalid message")
			continue
		}

		var msgType string
		if json.Unmarshal(envelope[0], &msgType) != nil {
			continue
		}

		switch msgType {
		case "EVENT":
			r.handleEvent(ctx, c, envelope)
		case "REQ":
			r.handleReq(ctx, c, envelope)
		case "CLOSE":
			r.handleClose(ctx, c, envelope)
		case "AUTH":
			r.handleAuth(ctx, c, envelope)
		}
	}
}

// handleEvent processes an incoming EVENT message.
func (r *Relay) handleEvent(ctx context.Context, c *conn, envelope []json.RawMessage) {
	if len(envelope) < 2 {
		return
	}
	var event nostr.Event
	if err := json.Unmarshal(envelope[1], &event); err != nil {
		r.sendOK(ctx, c, "", false, "invalid: bad JSON")
		return
	}

	// NIP-42: require authentication before accepting events.
	if r.RequireAuth && c.authedAs == "" {
		r.sendOK(ctx, c, event.ID, false, "auth-required: must authenticate")
		return
	}

	// Validate.
	if !event.Valid() {
		r.sendOK(ctx, c, event.ID, false, "invalid: bad id or signature")
		return
	}

	// Coherence filter: reject events that don't bond to any known structure.
	if r.ContentFilter != nil && !r.ContentFilter.Admits(&event) {
		r.sendOK(ctx, c, event.ID, false, "blocked: content incoherent with relay structure")
		return
	}

	// NIP-09: kind 5 deletion events remove referenced events by the same author.
	if event.Kind == 5 {
		r.handleDeletion(ctx, c, &event)
		return
	}

	// Store.
	r.mu.Lock()
	_, exists := r.events[event.ID]
	if !exists {
		r.events[event.ID] = &event
		r.order = append(r.order, event.ID)
	}
	r.mu.Unlock()

	r.sendOK(ctx, c, event.ID, true, "")

	// Broadcast to matching subscriptions (skip duplicate).
	if !exists {
		r.broadcast(ctx, &event)
	}
}

// handleAuth processes a NIP-42 AUTH response from a client.
// The client sends ["AUTH", <signed kind 22242 event>] with the
// challenge in a "challenge" tag and the relay URL in a "relay" tag.
func (r *Relay) handleAuth(ctx context.Context, c *conn, envelope []json.RawMessage) {
	if len(envelope) < 2 {
		return
	}
	var event nostr.Event
	if err := json.Unmarshal(envelope[1], &event); err != nil {
		r.sendOK(ctx, c, "", false, "invalid: bad auth event JSON")
		return
	}

	// Must be kind 22242.
	if event.Kind != 22242 {
		r.sendOK(ctx, c, event.ID, false, "invalid: auth event must be kind 22242")
		return
	}

	// Must have a valid signature.
	if !event.Valid() {
		r.sendOK(ctx, c, event.ID, false, "invalid: bad auth event signature")
		return
	}

	// Must contain the challenge we sent.
	challengeMatch := false
	for _, tag := range event.Tags {
		if len(tag) >= 2 && tag[0] == "challenge" && tag[1] == c.challenge {
			challengeMatch = true
			break
		}
	}
	if !challengeMatch {
		r.sendOK(ctx, c, event.ID, false, "invalid: challenge mismatch")
		return
	}

	c.authedAs = event.Pubkey
	r.sendOK(ctx, c, event.ID, true, "")
}

// handleDeletion processes a NIP-09 kind 5 deletion event.
// Only events authored by the same pubkey are deleted.
func (r *Relay) handleDeletion(ctx context.Context, c *conn, del *nostr.Event) {
	r.mu.Lock()
	deleted := 0
	for _, tag := range del.Tags {
		if len(tag) >= 2 && tag[0] == "e" {
			targetID := tag[1]
			if existing, ok := r.events[targetID]; ok {
				if existing.Pubkey == del.Pubkey {
					delete(r.events, targetID)
					deleted++
				}
			}
		}
	}
	// Store the deletion event itself.
	if _, exists := r.events[del.ID]; !exists {
		r.events[del.ID] = del
		r.order = append(r.order, del.ID)
	}
	r.mu.Unlock()

	r.sendOK(ctx, c, del.ID, true, "")
}

// handleReq processes a REQ message: ["REQ", "<sub_id>", <filter>, ...]
func (r *Relay) handleReq(ctx context.Context, c *conn, envelope []json.RawMessage) {
	if len(envelope) < 3 {
		return
	}
	var subID string
	if json.Unmarshal(envelope[1], &subID) != nil {
		return
	}

	var filters []nostr.Filter
	for _, raw := range envelope[2:] {
		var f nostr.Filter
		if json.Unmarshal(raw, &f) == nil {
			// Parse tag filters from raw JSON.
			var m map[string]json.RawMessage
			if json.Unmarshal(raw, &m) == nil {
				for k, v := range m {
					if strings.HasPrefix(k, "#") {
						var vals []string
						if json.Unmarshal(v, &vals) == nil {
							if f.Tags == nil {
								f.Tags = make(map[string][]string)
							}
							f.Tags[k[1:]] = vals
						}
					}
				}
			}
			filters = append(filters, f)
		}
	}

	// Register subscription.
	r.subMu.Lock()
	if r.subs[c] == nil {
		r.subs[c] = make(map[string][]nostr.Filter)
	}
	r.subs[c][subID] = filters
	r.subMu.Unlock()

	// Replay stored events matching the filters.
	r.mu.RLock()
	var matched []*nostr.Event
	for _, id := range r.order {
		ev := r.events[id]
		if matchesAny(ev, filters) {
			matched = append(matched, ev)
		}
	}
	r.mu.RUnlock()

	// Apply limit (from last filter that has one).
	limit := len(matched)
	for i := len(filters) - 1; i >= 0; i-- {
		if filters[i].Limit != nil {
			limit = *filters[i].Limit
			break
		}
	}
	if limit < len(matched) {
		// Return the most recent (last inserted) events.
		matched = matched[len(matched)-limit:]
	}

	for _, ev := range matched {
		r.sendEvent(ctx, c, subID, ev)
	}

	// EOSE.
	r.sendEOSE(ctx, c, subID)
}

// handleClose processes a CLOSE message: ["CLOSE", "<sub_id>"]
func (r *Relay) handleClose(ctx context.Context, c *conn, envelope []json.RawMessage) {
	if len(envelope) < 2 {
		return
	}
	var subID string
	if json.Unmarshal(envelope[1], &subID) != nil {
		return
	}

	r.subMu.Lock()
	if subs, ok := r.subs[c]; ok {
		delete(subs, subID)
	}
	r.subMu.Unlock()
}

// removeSubs cleans up all subscriptions for a disconnected client.
func (r *Relay) removeSubs(c *conn) {
	r.subMu.Lock()
	delete(r.subs, c)
	r.subMu.Unlock()
}

// broadcast sends a new event to all subscriptions that match it.
// Pre-marshals per subID to avoid redundant JSON encoding.
func (r *Relay) broadcast(ctx context.Context, ev *nostr.Event) {
	r.subMu.RLock()
	defer r.subMu.RUnlock()

	// Cache pre-marshaled messages per subID to avoid redundant encoding.
	cache := make(map[string][]byte)

	for c, subs := range r.subs {
		for subID, filters := range subs {
			if matchesAny(ev, filters) {
				data, ok := cache[subID]
				if !ok {
					data, _ = json.Marshal([]any{"EVENT", subID, ev})
					cache[subID] = data
				}
				c.mu.Lock()
				c.ws.Write(ctx, websocket.MessageText, data)
				c.mu.Unlock()
			}
		}
	}
}

// matchesAny returns true if the event matches any of the filters.
func matchesAny(ev *nostr.Event, filters []nostr.Filter) bool {
	for _, f := range filters {
		if matches(ev, f) {
			return true
		}
	}
	return false
}

// matches evaluates a single filter against an event.
func matches(ev *nostr.Event, f nostr.Filter) bool {
	if len(f.IDs) > 0 && !containsPrefix(f.IDs, ev.ID) {
		return false
	}
	if len(f.Authors) > 0 && !containsPrefix(f.Authors, ev.Pubkey) {
		return false
	}
	if len(f.Kinds) > 0 && !containsInt(f.Kinds, ev.Kind) {
		return false
	}
	if f.Since != nil && ev.CreatedAt < *f.Since {
		return false
	}
	if f.Until != nil && ev.CreatedAt > *f.Until {
		return false
	}
	// Tag filters.
	for tagName, values := range f.Tags {
		if !eventHasTagValue(ev, tagName, values) {
			return false
		}
	}
	return true
}

// containsPrefix returns true if any element of haystack is a prefix of needle (or equal).
func containsPrefix(haystack []string, needle string) bool {
	for _, h := range haystack {
		if strings.HasPrefix(needle, h) {
			return true
		}
	}
	return false
}

func containsInt(haystack []int, needle int) bool {
	return slices.Contains(haystack, needle)
}

// eventHasTagValue returns true if the event has a tag with the given name
// whose value matches any of the provided values.
func eventHasTagValue(ev *nostr.Event, tagName string, values []string) bool {
	for _, tag := range ev.Tags {
		if len(tag) >= 2 && tag[0] == tagName {
			for _, v := range values {
				if tag[1] == v {
					return true
				}
			}
		}
	}
	return false
}

// Wire message helpers.

func (r *Relay) sendOK(ctx context.Context, c *conn, eventID string, accepted bool, msg string) {
	data, _ := json.Marshal([]any{"OK", eventID, accepted, msg})
	c.mu.Lock()
	defer c.mu.Unlock()
	c.ws.Write(ctx, websocket.MessageText, data)
}

func (r *Relay) sendEvent(ctx context.Context, c *conn, subID string, ev *nostr.Event) {
	data, _ := json.Marshal([]any{"EVENT", subID, ev})
	c.mu.Lock()
	defer c.mu.Unlock()
	c.ws.Write(ctx, websocket.MessageText, data)
}

func (r *Relay) sendEOSE(ctx context.Context, c *conn, subID string) {
	data, _ := json.Marshal([]string{"EOSE", subID})
	c.mu.Lock()
	defer c.mu.Unlock()
	c.ws.Write(ctx, websocket.MessageText, data)
}

func (r *Relay) sendNotice(ctx context.Context, c *conn, msg string) {
	data, _ := json.Marshal([]string{"NOTICE", msg})
	c.mu.Lock()
	defer c.mu.Unlock()
	c.ws.Write(ctx, websocket.MessageText, data)
}

func (r *Relay) sendAuth(ctx context.Context, c *conn, challenge string) {
	data, _ := json.Marshal([]string{"AUTH", challenge})
	c.mu.Lock()
	defer c.mu.Unlock()
	c.ws.Write(ctx, websocket.MessageText, data)
}

// Stats returns relay operational statistics.
type Stats struct {
	Events        int
	Connections   int
	Subscriptions int
}

// Stats returns current relay statistics.
func (r *Relay) Stats() Stats {
	r.mu.RLock()
	events := len(r.events)
	r.mu.RUnlock()

	r.subMu.RLock()
	conns := len(r.subs)
	subs := 0
	for _, s := range r.subs {
		subs += len(s)
	}
	r.subMu.RUnlock()

	return Stats{
		Events:        events,
		Connections:   conns,
		Subscriptions: subs,
	}
}