// Package memrelay is an in-process NIP-01 relay used by tests. It implements
// just REQ/EVENT/CLOSE with local filter matching, and can deterministically
// drop every Nth delivery to exercise a client's gap recovery.
package memrelay

import (
	"errors"
	"io"
	"sync"

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

type Relay struct {
	mu        sync.Mutex
	conns     map[*memTransport]map[string]nostr.Filter
	dropEvery int
	drop      map[*memTransport]int
	closed    bool
}

func New() *Relay {
	return &Relay{
		conns: make(map[*memTransport]map[string]nostr.Filter),
		drop:  make(map[*memTransport]int),
	}
}

func (r *Relay) Dial(string) (relay.Transport, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.closed {
		return nil, errors.New("memrelay: closed")
	}
	t := &memTransport{
		relay: r,
		in:    make(chan []byte, 512),
		done:  make(chan struct{}),
	}
	r.conns[t] = make(map[string]nostr.Filter)
	return t, nil
}

func (r *Relay) SetDrop(n int) {
	r.mu.Lock()
	r.dropEvery = n
	r.drop = make(map[*memTransport]int)
	r.mu.Unlock()
}

func (r *Relay) Close() {
	r.mu.Lock()
	r.closed = true
	conns := make([]*memTransport, 0, len(r.conns))
	for c := range r.conns {
		conns = append(conns, c)
	}
	r.mu.Unlock()
	for _, c := range conns {
		c.Close()
	}
}

type memTransport struct {
	relay *Relay
	in    chan []byte
	done  chan struct{}
	once  sync.Once
}

func (t *memTransport) WriteText(raw []byte) error {
	parts, err := nostr.SplitArray(raw)
	if err != nil || len(parts) == 0 {
		return errors.New("memrelay: bad message")
	}
	typ, err := nostr.ParseJSONString(parts[0])
	if err != nil {
		return err
	}
	switch typ {
	case "REQ":
		if len(parts) < 3 {
			return errors.New("memrelay: REQ")
		}
		id, err := nostr.ParseJSONString(parts[1])
		if err != nil {
			return err
		}
		f, err := nostr.ParseFilter(parts[2])
		if err != nil {
			return err
		}
		t.relay.mu.Lock()
		if subs := t.relay.conns[t]; subs != nil {
			subs[id] = *f
		}
		t.relay.mu.Unlock()
		t.deliver([]byte(`["EOSE","` + id + `"]`))
		return nil
	case "CLOSE":
		if len(parts) < 2 {
			return errors.New("memrelay: CLOSE")
		}
		id, _ := nostr.ParseJSONString(parts[1])
		t.relay.mu.Lock()
		if subs := t.relay.conns[t]; subs != nil {
			delete(subs, id)
		}
		t.relay.mu.Unlock()
		return nil
	case "EVENT":
		if len(parts) < 2 {
			return errors.New("memrelay: EVENT")
		}
		ev, err := nostr.ParseEvent(parts[1])
		if err != nil {
			return err
		}
		if err := ev.Verify(); err != nil {
			return err
		}
		id := nostr.Hex(ev.ID[:])
		t.deliver([]byte(`["OK","` + id + `",true,""]`))
		t.relay.route(t, ev)
		return nil
	default:
		return errors.New("memrelay: unsupported " + typ)
	}
}

func (r *Relay) route(from *memTransport, ev *nostr.Event) {
	type target struct {
		t  *memTransport
		id string
	}
	r.mu.Lock()
	var targets []target
	drop := make(map[target]bool)
	for t, subs := range r.conns {
		for id, f := range subs {
			if !f.Matches(ev) {
				continue
			}
			tg := target{t: t, id: id}
			if r.dropEvery > 0 {
				r.drop[t]++
				if r.drop[t]%r.dropEvery == 0 {
					drop[tg] = true
				}
			}
			targets = append(targets, tg)
		}
	}
	r.mu.Unlock()
	raw := ev.AppendJSON(nil)
	for _, tg := range targets {
		if drop[tg] {
			continue
		}
		msg := make([]byte, 0, len(raw)+len(tg.id)+16)
		msg = append(msg, `["EVENT","`...)
		msg = append(msg, tg.id...)
		msg = append(msg, `",`...)
		msg = append(msg, raw...)
		msg = append(msg, ']')
		tg.t.deliver(msg)
	}
	_ = from
}

func (t *memTransport) deliver(msg []byte) {
	select {
	case t.in <- msg:
	case <-t.done:
	default:
		// A real relay has a bounded send buffer; the client re-syncs or the
		// tunnel retransmits, so drop rather than block the whole relay.
	}
}

func (t *memTransport) ReadMessage() ([]byte, error) {
	select {
	case msg := <-t.in:
		return msg, nil
	case <-t.done:
		return nil, io.EOF
	}
}

func (t *memTransport) Close() error {
	t.once.Do(func() {
		close(t.done)
		t.relay.mu.Lock()
		delete(t.relay.conns, t)
		t.relay.mu.Unlock()
	})
	return nil
}

var _ relay.Transport = (*memTransport)(nil)