// Package relay is a small NIP-01 client for one relay connection.
//
// It owns the text frames, the subscription bookkeeping and the event
// signature check. Reconnection is the caller's policy: a Client is one
// connection, and when Run returns every subscription has ended.
package relay

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"git.smesh.lol/nostrino/pkg/nostr"
	"git.smesh.lol/nostrino/pkg/ws"
)

// Transport is one WebSocket connection carrying text messages. The ws
// package implements it for native builds; a syscall/js adapter can implement
// it for wasm.
type Transport interface {
	WriteText([]byte) error
	ReadMessage() ([]byte, error)
	Close() error
}

// DialTransport opens a ws:// or wss:// transport.
func DialTransport(rawURL string) (Transport, error) {
	return ws.Dial(rawURL)
}

// Options bound a relay connection.
type Options struct {
	// EventBuffer is the per-subscription queue length. Zero means 64.
	EventBuffer int
	// OnDrop is called when a subscription queue is full. It must not block.
	OnDrop func(subID string, dropped int)
	// RelayURL is the URL signed into the NIP-42 auth event. Empty means the
	// relay URL is not available and AUTH challenges cannot be answered
	// correctly; a relay that requires auth will reject the connection.
	RelayURL string
}

// Client is one authenticated relay connection.
type Client struct {
	t    Transport
	opts Options

	key [32]byte
	pub [32]byte

	writeMu sync.Mutex
	mu      sync.Mutex
	subs    map[string]*Subscription
	oks     map[string]chan string
	nextID  uint64
	closed  bool
	err     error
	done    chan struct{}

	onClose func(error)

	relayURL  string
	authState int // 0 unseen, 1 pending, 2 authenticated, 3 failed
	authID    string
	authReady chan struct{}
	authDone  bool
	authErr   error
}

// New starts a relay client and its read loop. The caller must not use t
// afterwards. key signs the client's own events.
func New(t Transport, key [32]byte, opts Options) (*Client, error) {
	pub, ok := nostrPubKey(key)
	if !ok {
		return nil, errors.New("relay: invalid secret key")
	}
	if opts.EventBuffer <= 0 {
		opts.EventBuffer = 64
	}
	c := &Client{
		t:        t,
		opts:     opts,
		key:      key,
		pub:      pub,
		subs:     make(map[string]*Subscription),
		oks:      make(map[string]chan string),
		done:     make(chan struct{}),
		relayURL: opts.RelayURL,
	}
	go c.readLoop()
	return c, nil
}

func nostrPubKey(key [32]byte) ([32]byte, bool) {
	var zero [32]byte
	// The nostr package validates and derives the key through its own helper.
	h, err := nostr.PubKeyHex(key)
	if err != nil || h == "" {
		return zero, false
	}
	var pub [32]byte
	if err := nostr.Decode32(h, pub[:]); err != nil {
		return zero, false
	}
	return pub, true
}

// PubKey is the client's public key.
func (c *Client) PubKey() [32]byte { return c.pub }

// PubKeyHex is the client's public key as lowercase hex.
func (c *Client) PubKeyHex() string { return nostr.Hex(c.pub[:]) }

// Done is closed when the read loop ends.
func (c *Client) Done() <-chan struct{} { return c.done }

// Err returns the read-loop error after Done is closed.
func (c *Client) Err() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.err
}

// OnClose installs a callback that runs once when the connection ends.
func (c *Client) OnClose(fn func(error)) {
	c.mu.Lock()
	c.onClose = fn
	c.mu.Unlock()
}

// Subscription is one REQ and its delivered events.
type Subscription struct {
	ID     string
	Filter nostr.Filter

	Events chan *nostr.Event
	EOSE   chan struct{}

	once     sync.Once
	eoseOnce sync.Once
	done     chan struct{}

	client *Client
}

// Done is closed when the subscription is closed or the connection ends.
func (s *Subscription) Done() <-chan struct{} { return s.done }

// Close sends CLOSE and removes the subscription.
func (s *Subscription) Close() error { return s.client.unsubscribe(s.ID) }

// Err returns the connection error once Done is closed.
func (s *Subscription) Err() error { return s.client.Err() }

func (s *Subscription) finish() {
	s.once.Do(func() {
		close(s.done)
		s.signalEOSE()
	})
}

func (s *Subscription) signalEOSE() {
	s.eoseOnce.Do(func() { close(s.EOSE) })
}

// Subscribe sends a REQ and returns the subscription. It does not wait for
// EOSE; call WaitEOSE when the subscription must be live before the peer is
// addressed.
// awaitAuth blocks while a NIP-42 authentication is in flight. It returns nil
// when no auth has been seen or auth succeeded.
func (c *Client) awaitAuth(ctx context.Context) error {
	c.mu.Lock()
	state, ch, err := c.authState, c.authReady, c.authErr
	c.mu.Unlock()
	if state != 1 {
		return err
	}
	if ch == nil {
		return nil
	}
	select {
	case <-ch:
		c.mu.Lock()
		err := c.authErr
		c.mu.Unlock()
		return err
	case <-ctx.Done():
		return ctx.Err()
	case <-c.done:
		return c.Err()
	}
}

func (c *Client) sendAuth(challenge string) {
	if c.relayURL == "" {
		return
	}
	ev := AuthEvent(c.relayURL, challenge, c.key)
	if ev == nil {
		return
	}
	c.mu.Lock()
	c.authState = 1
	c.authID = nostr.Hex(ev.ID[:])
	c.authReady = make(chan struct{})
	c.authDone = false
	c.authErr = nil
	c.mu.Unlock()

	msg := append([]byte(`["AUTH",`), ev.AppendJSON(nil)...)
	msg = append(msg, ']')
	if err := c.write(msg); err != nil {
		c.authFinished(err)
	}
}

func (c *Client) authFinished(err error) {
	c.mu.Lock()
	if c.authDone {
		c.mu.Unlock()
		return
	}
	c.authDone = true
	c.authErr = err
	if err != nil {
		c.authState = 3
	} else {
		c.authState = 2
	}
	ch := c.authReady
	c.mu.Unlock()
	if ch != nil {
		close(ch)
	}
}

// AuthEvent builds and signs a NIP-42 authentication event.
func AuthEvent(relayURL, challenge string, key [32]byte) *nostr.Event {
	ev := &nostr.Event{
		Kind:      22242,
		CreatedAt: time.Now().Unix(),
		Tags: nostr.Tags{
			{"relay", relayURL},
			{"challenge", challenge},
		},
	}
	if err := ev.Sign(key); err != nil {
		return nil
	}
	return ev
}

func (c *Client) Subscribe(f nostr.Filter) (*Subscription, error) {
	c.mu.Lock()
	if c.closed {
		err := c.err
		c.mu.Unlock()
		if err == nil {
			err = errors.New("relay: closed")
		}
		return nil, err
	}
	c.nextID++
	id := fmt.Sprintf("n%x", c.nextID)
	s := &Subscription{
		ID:     id,
		Filter: f,
		Events: make(chan *nostr.Event, c.opts.EventBuffer),
		EOSE:   make(chan struct{}),
		done:   make(chan struct{}),
		client: c,
	}
	c.subs[id] = s
	c.mu.Unlock()

	msg := []byte(`["REQ","` + id + `",`)
	msg = f.AppendJSON(msg)
	msg = append(msg, ']')
	if err := c.write(msg); err != nil {
		c.mu.Lock()
		delete(c.subs, id)
		c.mu.Unlock()
		return nil, err
	}
	return s, nil
}

// WaitEOSE waits until the relay has replayed the stored matches, or ctx ends.
func (c *Client) WaitEOSE(ctx context.Context, s *Subscription) error {
	select {
	case <-s.EOSE:
		return nil
	case <-s.done:
		if err := c.Err(); err != nil {
			return err
		}
		return errors.New("relay: subscription closed")
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (c *Client) unsubscribe(id string) error {
	c.mu.Lock()
	s, ok := c.subs[id]
	if ok {
		delete(c.subs, id)
	}
	c.mu.Unlock()
	if !ok {
		return nil
	}
	err := c.write([]byte(`["CLOSE","` + id + `"]`))
	s.finish()
	return err
}

// Send signs and writes an event without waiting for the relay's OK.
// It is the data-plane path: the tunnel's own acknowledgements decide
// delivery, and waiting for a per-frame OK would add a round trip per frame.
func (c *Client) Send(ev *nostr.Event) error {
	actx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	err := c.awaitAuth(actx)
	cancel()
	if err != nil {
		return err
	}
	if err := ev.Sign(c.key); err != nil {
		return err
	}
	msg := append([]byte(`["EVENT",`), ev.AppendJSON(nil)...)
	msg = append(msg, ']')
	return c.write(msg)
}

// Publish signs and writes an event and waits for the relay's OK.
func (c *Client) Publish(ctx context.Context, ev *nostr.Event) error {
	if err := c.awaitAuth(ctx); err != nil {
		return err
	}
	if err := ev.Sign(c.key); err != nil {
		return err
	}
	ch := make(chan string, 1)
	id := nostr.Hex(ev.ID[:])
	c.mu.Lock()
	c.oks[id] = ch
	c.mu.Unlock()
	defer func() {
		c.mu.Lock()
		delete(c.oks, id)
		c.mu.Unlock()
	}()

	msg := append([]byte(`["EVENT",`), ev.AppendJSON(nil)...)
	msg = append(msg, ']')
	if err := c.write(msg); err != nil {
		return err
	}
	select {
	case reason := <-ch:
		if reason != "" {
			return errors.New("relay: rejected: " + reason)
		}
		return nil
	case <-c.done:
		if err := c.Err(); err != nil {
			return err
		}
		return errors.New("relay: closed")
	case <-ctx.Done():
		return ctx.Err()
	}
}

// Close closes the transport and ends the read loop.
func (c *Client) Close() error {
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return nil
	}
	c.closed = true
	c.mu.Unlock()
	err := c.t.Close()
	// Release waiters at once rather than waiting for the read loop to observe
	// the closed socket. fail is idempotent, so the read loop's own report is
	// harmless.
	c.fail(errors.New("relay: closed"))
	return err
}

func (c *Client) write(msg []byte) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	if err := c.t.WriteText(msg); err != nil {
		c.fail(err)
		return err
	}
	return nil
}

func (c *Client) readLoop() {
	for {
		raw, err := c.t.ReadMessage()
		if err != nil {
			c.fail(err)
			return
		}
		if err := c.handle(raw); err != nil {
			c.fail(err)
			return
		}
	}
}

func (c *Client) fail(err error) {
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return
	}
	c.closed = true
	c.err = err
	subs := make([]*Subscription, 0, len(c.subs))
	for _, s := range c.subs {
		subs = append(subs, s)
	}
	c.subs = map[string]*Subscription{}
	oks := c.oks
	c.oks = map[string]chan string{}
	authCh := c.authReady
	authDone := c.authDone
	if !authDone && authCh != nil {
		c.authDone = true
		c.authState = 3
		c.authErr = err
	}
	fn := c.onClose
	c.mu.Unlock()
	if !authDone && authCh != nil {
		close(authCh)
	}

	close(c.done)
	for _, s := range subs {
		s.finish()
	}
	for id, ch := range oks {
		select {
		case ch <- "connection closed":
		default:
			_ = id
		}
	}
	c.t.Close()
	if fn != nil {
		fn(err)
	}
}

func (c *Client) handle(raw []byte) error {
	parts, err := nostr.SplitArray(raw)
	if err != nil {
		return fmt.Errorf("relay: bad message: %w", err)
	}
	if len(parts) == 0 {
		return errors.New("relay: empty message")
	}
	typ, err := nostr.ParseJSONString(parts[0])
	if err != nil {
		return fmt.Errorf("relay: bad message type: %w", err)
	}
	switch typ {
	case "EVENT":
		if len(parts) != 3 {
			return errors.New("relay: malformed EVENT")
		}
		subID, err := nostr.ParseJSONString(parts[1])
		if err != nil {
			return fmt.Errorf("relay: malformed EVENT subscription: %w", err)
		}
		ev, err := nostr.ParseEvent(parts[2])
		if err != nil {
			return fmt.Errorf("relay: malformed event: %w", err)
		}
		if err := ev.Verify(); err != nil {
			return fmt.Errorf("relay: event verification: %w", err)
		}
		c.mu.Lock()
		s := c.subs[subID]
		c.mu.Unlock()
		if s == nil {
			return nil
		}
		if !s.Filter.Matches(ev) {
			return nil
		}
		select {
		case s.Events <- ev:
		default:
			if c.opts.OnDrop != nil {
				c.opts.OnDrop(subID, 1)
			}
		}
		return nil
	case "EOSE":
		if len(parts) < 2 {
			return errors.New("relay: malformed EOSE")
		}
		subID, err := nostr.ParseJSONString(parts[1])
		if err != nil {
			return err
		}
		c.mu.Lock()
		s := c.subs[subID]
		c.mu.Unlock()
		if s != nil {
			s.signalEOSE()
		}
		return nil
	case "OK":
		if len(parts) < 3 {
			return errors.New("relay: malformed OK")
		}
		id, err := nostr.ParseJSONString(parts[1])
		if err != nil {
			return err
		}
		accepted := string(parts[2]) == "true"
		reason := ""
		if !accepted {
			reason = "rejected"
			if len(parts) >= 4 {
				if m, err := nostr.ParseJSONString(parts[3]); err == nil && m != "" {
					reason = m
				}
			}
		}
		c.mu.Lock()
		ch := c.oks[id]
		authID := c.authID
		c.mu.Unlock()
		if id == authID && authID != "" {
			if reason == "" {
				c.authFinished(nil)
			} else {
				c.authFinished(errors.New("relay: authentication rejected: " + reason))
			}
		}
		if ch != nil {
			select {
			case ch <- reason:
			default:
			}
		}
		return nil
	case "AUTH":
		if len(parts) < 2 {
			return errors.New("relay: malformed AUTH")
		}
		challenge, err := nostr.ParseJSONString(parts[1])
		if err != nil {
			return err
		}
		c.mu.Lock()
		state := c.authState
		c.mu.Unlock()
		if state != 1 && state != 2 {
			c.sendAuth(challenge)
		}
		return nil
	case "NOTICE", "CLOSED":
		// NOTICE is advisory; CLOSED ends one subscription; AUTH is not used
		// by the public data plane. Torly closes only invalid subscriptions,
		// which the tunnel treats as a connection failure through its
		// acknowledgements anyway.
		if typ == "CLOSED" && len(parts) >= 2 {
			if subID, err := nostr.ParseJSONString(parts[1]); err == nil {
				c.mu.Lock()
				s := c.subs[subID]
				delete(c.subs, subID)
				c.mu.Unlock()
				if s != nil {
					s.finish()
				}
			}
		}
		return nil
	default:
		// Unknown message types are ignored, as NIP-01 requires.
		return nil
	}
}