package server

import (
	"bytes"
	"testing"
	"time"

	"git.smesh.lol/smesh/pkg/broadcast"
	"git.smesh.lol/smesh/pkg/nostr/envelope"
	"git.smesh.lol/smesh/pkg/nostr/event"
	"git.smesh.lol/smesh/pkg/nostr/signer/p8k"
	"git.smesh.lol/smesh/pkg/nostr/tag"
	"git.smesh.lol/smesh/pkg/relay/config"
)

// stSubSrv builds a server with a live broadcast worker (its own spawn domain)
// and a fake database. Callers must close the broadcaster.
func stSubSrv(c *config.C) (s *Server, b *broadcast.Broadcaster) {
	s = stServer(c)
	s.db = stFakeDB()
	b = broadcast.New()
	s.bcast = b
	return
}

// stAuthEvent builds and signs a NIP-42 AUTH event. A nil tag slice omits that
// tag, so the missing-tag failure paths are reachable.
func stAuthEvent(t *testing.T, sg *p8k.Signer, k uint16, ch, rl []byte) (ev *event.E) {
	t.Helper()
	ev = event.New()
	ev.CreatedAt = time.Now().Unix()
	ev.Kind = k
	var chT *tag.T
	var rlT *tag.T
	if ch != nil {
		chT = tag.NewFromBytesSlice([]byte("challenge"), ch)
	}
	if rl != nil {
		rlT = tag.NewFromBytesSlice([]byte("relay"), rl)
	}
	if chT != nil && rlT != nil {
		ev.Tags = tag.NewS(chT, rlT)
	} else if chT != nil {
		ev.Tags = tag.NewS(chT)
	} else if rlT != nil {
		ev.Tags = tag.NewS(rlT)
	}
	if serr := ev.Sign(sg); serr != nil {
		t.Fatal(serr)
		return
	}
	return
}

func stAuthMsg(ev *event.E) (msg []byte) {
	ar := &envelope.AuthResponse{Event: ev}
	return ar.Marshal(nil)
}

func stNewSigner(t *testing.T) (sg *p8k.Signer) {
	t.Helper()
	sg = p8k.MustNew()
	if gerr := sg.Generate(); gerr != nil {
		t.Fatal(gerr)
		return
	}
	return
}

// --- subscription registry ---

func TestHandleReqRegistersSubscription(t *testing.T) {
	s, b := stSubSrv(stCfg())
	defer b.Close()
	fd := int32(4)
	s.conns[fd] = &cstate{subs: map[string]*sub{}}

	req := []byte("[\"REQ\",\"s1\",{}]")
	s.handleReq(fd, req)
	cs := s.conns[fd]
	if len(cs.subs) != 1 {
		t.Fatalf("subs = %d, want 1", len(cs.subs))
		return
	}
	sb := cs.subs["s1"]
	if sb == nil {
		t.Fatal("subscription s1 must be registered")
		return
	}
	if sb.id != "s1" {
		t.Fatalf("subscription id = %s", sb.id)
		return
	}
	if !bytes.Equal(sb.rawReq, req) {
		t.Fatal("subscription must retain the raw REQ frame")
		return
	}
	if len(s.dbPending) != 1 {
		t.Fatalf("dbPending = %d, want a history backfill", len(s.dbPending))
		return
	}
}

func TestHandleReqSubscriptionCap(t *testing.T) {
	c := stCfg()
	c.MaxSubscriptions = 1
	s, b := stSubSrv(c)
	defer b.Close()
	fd := int32(4)
	s.conns[fd] = &cstate{subs: map[string]*sub{}}

	s.handleReq(fd, []byte("[\"REQ\",\"s1\",{}]"))
	s.handleReq(fd, []byte("[\"REQ\",\"s2\",{}]"))
	if len(s.conns[fd].subs) != 1 {
		t.Fatalf("cap must hold at 1, got %d", len(s.conns[fd].subs))
		return
	}
	if _, has := s.conns[fd].subs["s2"]; has {
		t.Fatal("s2 must have been refused")
		return
	}
	// Re-subscribing an existing id is a replace, not a second slot.
	s.handleReq(fd, []byte("[\"REQ\",\"s1\",{}]"))
	if len(s.conns[fd].subs) != 1 {
		t.Fatal("replacing an existing subscription must not grow the set")
		return
	}
}

func TestHandleReqAuthRequired(t *testing.T) {
	c := stCfg()
	c.AuthRequired = true
	s, b := stSubSrv(c)
	defer b.Close()
	fd := int32(4)
	cs := &cstate{subs: map[string]*sub{}}
	s.conns[fd] = cs

	s.handleReq(fd, []byte("[\"REQ\",\"s1\",{}]"))
	if len(cs.subs) != 0 {
		t.Fatal("an unauthed REQ must be refused under AuthRequired")
		return
	}

	cs.authedPubkey = []byte{0xAB}
	s.handleReq(fd, []byte("[\"REQ\",\"s1\",{}]"))
	if len(cs.subs) != 1 {
		t.Fatal("an authed REQ must be accepted")
		return
	}
}

func TestHandleReqMalformedAndTainted(t *testing.T) {
	s, b := stSubSrv(stCfg())
	defer b.Close()
	fd := int32(4)
	cs := &cstate{subs: map[string]*sub{}}
	s.conns[fd] = cs

	s.handleReq(99, []byte("[\"REQ\",\"s1\",{}]"))

	s.handleReq(fd, []byte("[\"REQ\",]"))
	if len(cs.subs) != 0 {
		t.Fatal("a malformed REQ must not register")
		return
	}

	// An undersized hex id taints the filter and jails the connection.
	s.handleReq(fd, []byte("[\"REQ\",\"s2\",{\"ids\":[\"abcd\"]}]"))
	if !cs.jailed {
		t.Fatal("an undersized hex id must jail the connection")
		return
	}
	if len(cs.subs) != 0 {
		t.Fatal("a tainted REQ must not register")
		return
	}
}

func TestHandleClose(t *testing.T) {
	s, b := stSubSrv(stCfg())
	defer b.Close()
	fd := int32(4)
	cs := &cstate{subs: map[string]*sub{}}
	s.conns[fd] = cs

	s.handleReq(fd, []byte("[\"REQ\",\"s1\",{}]"))
	if len(cs.subs) != 1 {
		t.Fatal("setup: s1 must be registered")
		return
	}
	s.handleClose(fd, []byte("[\"CLOSE\",\"s1\"]"))
	if len(cs.subs) != 0 {
		t.Fatal("CLOSE must remove the subscription")
		return
	}
	// Unknown connection and malformed CLOSE are no-ops.
	s.handleClose(99, []byte("[\"CLOSE\",\"s1\"]"))
	s.handleClose(fd, []byte("[\"CLOSE\"]"))
}

func TestOnWSClose(t *testing.T) {
	s, b := stSubSrv(stCfg())
	defer b.Close()
	fd := int32(4)
	cs := &cstate{subs: map[string]*sub{}}
	cs.subs["s1"] = &sub{id: "s1"}
	cs.challenge = []byte("c")
	cs.authedPubkey = []byte("p")
	s.conns[fd] = cs
	s.pendingReq[1] = fd
	s.pendingReq[2] = 99
	s.dbPending[3] = fd
	s.dbPending[4] = 99

	s.OnWSClose(fd)

	if s.conns[fd] != nil {
		t.Fatal("OnWSClose must drop the connection record")
		return
	}
	if _, ok := s.pendingReq[1]; ok {
		t.Fatal("pendingReq entries for the closed fd must be dropped")
		return
	}
	if _, ok := s.pendingReq[2]; !ok {
		t.Fatal("pendingReq entries for other fds must survive")
		return
	}
	if _, ok := s.dbPending[3]; ok {
		t.Fatal("dbPending entries for the closed fd must be dropped")
		return
	}
	if _, ok := s.dbPending[4]; !ok {
		t.Fatal("dbPending entries for other fds must survive")
		return
	}
}

// --- NIP-42 auth ---

func TestHandleAuthSuccess(t *testing.T) {
	c := stCfg()
	c.RelayURL = "wss://relay.example.com"
	s, b := stSubSrv(c)
	defer b.Close()
	fd := int32(4)
	cs := &cstate{subs: map[string]*sub{}}
	cs.challenge = []byte("chal-123")
	s.conns[fd] = cs
	sg := stNewSigner(t)

	ev := stAuthEvent(t, sg, 22242, []byte("chal-123"), []byte("relay.example.com"))
	s.handleAuth(fd, stAuthMsg(ev))

	if len(cs.authedPubkey) != 32 {
		t.Fatalf("authedPubkey length = %d, want 32", len(cs.authedPubkey))
		return
	}
	if !bytes.Equal(cs.authedPubkey, ev.Pubkey) {
		t.Fatal("authedPubkey must be the auth event's pubkey")
		return
	}
}

func TestHandleAuthFailures(t *testing.T) {
	c := stCfg()
	c.RelayURL = "wss://relay.example.com"
	s := stServer(c)
	fd := int32(4)
	cs := &cstate{subs: map[string]*sub{}}
	cs.challenge = []byte("chal-123")
	s.conns[fd] = cs
	sg := stNewSigner(t)

	// No connection, then an unparseable AUTH frame.
	s.handleAuth(99, []byte("[\"AUTH\",]"))
	s.handleAuth(fd, []byte("[\"AUTH\",]"))

	// Wrong event kind.
	s.handleAuth(fd, stAuthMsg(stAuthEvent(t, sg, 1, []byte("chal-123"), []byte("relay.example.com"))))
	// Missing challenge tag.
	s.handleAuth(fd, stAuthMsg(stAuthEvent(t, sg, 22242, nil, []byte("relay.example.com"))))
	// Wrong challenge.
	s.handleAuth(fd, stAuthMsg(stAuthEvent(t, sg, 22242, []byte("nope"), []byte("relay.example.com"))))
	// Missing relay tag.
	s.handleAuth(fd, stAuthMsg(stAuthEvent(t, sg, 22242, []byte("chal-123"), nil)))
	// Relay URL mismatch.
	s.handleAuth(fd, stAuthMsg(stAuthEvent(t, sg, 22242, []byte("chal-123"), []byte("other.example.com"))))

	// Timestamp out of range.
	old := stAuthEvent(t, sg, 22242, []byte("chal-123"), []byte("relay.example.com"))
	old.CreatedAt = time.Now().Unix() - 1000
	if serr := old.Sign(sg); serr != nil {
		t.Fatal(serr)
		return
	}
	s.handleAuth(fd, stAuthMsg(old))

	// Invalid signature.
	bad := stAuthEvent(t, sg, 22242, []byte("chal-123"), []byte("relay.example.com"))
	bad.Sig[0] ^= 0xFF
	s.handleAuth(fd, stAuthMsg(bad))

	if len(cs.authedPubkey) != 0 {
		t.Fatal("no failure path may authenticate the connection")
		return
	}
}

// --- broadcast domain ---

func TestHandleBroadcastFrameNoFrames(t *testing.T) {
	s, b := stSubSrv(stCfg())
	defer b.Close()
	s.handleBroadcastFrame()
	if !b.Valid() {
		t.Fatal("the broadcast worker must stay alive")
		return
	}
}

func TestSendBroadcastFlags(t *testing.T) {
	c := stCfg()
	c.RelayURL = "wss://relay.example"
	c.PrivilegedOpen = false
	c.NIP70Enforce = true
	c.MarmotOpen = true
	s, b := stSubSrv(c)
	defer b.Close()
	s.sendBroadcast([]byte("[\"EVENT\",{}]"), 3)
}

// --- worker lifecycle ---

func TestWorkerStartupAndRespawn(t *testing.T) {
	s := stServer(stCfg())

	s.startIngestWorkers(1)
	if s.workers.Len() != 1 || len(s.workerIn) != 1 {
		t.Fatal("ingest pool must be started")
		return
	}
	if !workerAlive(s.workerDone[0]) {
		t.Fatal("ingest worker must be alive")
		return
	}
	oldIngest := s.workerIn[0]
	s.workers.Busy[0] = true
	s.respawnIngestWorker(0)
	if s.workers.Busy[0] {
		t.Fatal("ingest respawn must clear the busy flag")
		return
	}
	if !workerAlive(s.workerDone[0]) {
		t.Fatal("respawned ingest worker must be alive")
		return
	}
	close(oldIngest)
	close(s.workerIn[0])

	s.startProxyWorkers(1)
	if s.proxyPool.Len() != 1 || len(s.proxyBusyTime) != 1 {
		t.Fatal("proxy pool must be started")
		return
	}
	oldProxy := s.proxyIn[0]
	s.proxyPool.Busy[0] = true
	s.proxyBusyTime[0] = 123
	s.respawnProxyWorker(0)
	if s.proxyPool.Busy[0] || s.proxyBusyTime[0] != 0 {
		t.Fatal("proxy respawn must reset busy state")
		return
	}
	if !workerAlive(s.proxyDone[0]) {
		t.Fatal("respawned proxy worker must be alive")
		return
	}
	close(oldProxy)
	close(s.proxyIn[0])

	s.startBlossomWorkers(1)
	if s.blossomPool.Len() != 1 {
		t.Fatal("blossom pool must be started")
		return
	}
	oldBlossom := s.blossomIn[0]
	s.blossomPool.Busy[0] = true
	s.respawnBlossomWorker(0)
	if s.blossomPool.Busy[0] {
		t.Fatal("blossom respawn must clear the busy flag")
		return
	}
	if !workerAlive(s.blossomDone[0]) {
		t.Fatal("respawned blossom worker must be alive")
		return
	}
	close(oldBlossom)
	close(s.blossomIn[0])
}

func TestServerCloseIdempotent(t *testing.T) {
	c := stCfg()
	s, b := stSubSrv(c)
	if b == nil {
		t.Fatal("setup: broadcaster must exist")
		return
	}
	s.Close()
	s.Close()
	if !s.closed {
		t.Fatal("Close must set the closed flag")
		return
	}
	// Close's effect on the broadcast worker is not asserted: in this build a
	// spawn channel's close() is not observed by the child and the spawn done
	// channel never signals the child's exit, so Valid() stays true forever.
	// Reported; the closed flag (idempotence) is what Close can be held to.

	// A server with no database still closes cleanly and idempotently.
	s2 := stServer(stCfg())
	s2.bcast = broadcast.New()
	s2.Close()
	s2.Close()
	if !s2.closed {
		t.Fatal("second server Close must set the closed flag")
		return
	}
}