package hub

import (
	"bytes"
	"context"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/coder/websocket"
	"github.com/coder/websocket/wsjson"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/shcizo/package-updater/internal/agenttoken"
	"github.com/shcizo/package-updater/internal/master/events"
	"github.com/shcizo/package-updater/internal/master/model"
	"github.com/shcizo/package-updater/internal/protocol"
)

type fakeStore struct {
	mu        sync.Mutex
	agents    map[string]model.Agent // by token hash
	hellos    []protocol.Hello
	snapshots [][]protocol.Service
	upserts   []string // "clusterID|key"
	removes   []string // "clusterID|key"
	seen      []int64
	lookups   int
	// revokeFromLookup, when > 0, makes AgentByTokenHash return the agent as
	// revoked from that (1-based) lookup onwards.
	revokeFromLookup int
}

func newFakeStore(agents ...model.Agent) *fakeStore {
	f := &fakeStore{agents: map[string]model.Agent{}}
	for _, a := range agents {
		f.agents[agenttoken.Hash("pua_"+a.Name)] = a
	}
	return f
}

func (f *fakeStore) AgentByTokenHash(_ context.Context, hash string) (model.Agent, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.lookups++
	a, ok := f.agents[hash]
	if !ok {
		return model.Agent{}, model.ErrNotFound
	}
	if f.revokeFromLookup > 0 && f.lookups >= f.revokeFromLookup {
		a.RevokedAt = ptr(time.Now())
	}
	return a, nil
}

func (f *fakeStore) RecordHello(_ context.Context, _ int64, h protocol.Hello, _ time.Time) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.hellos = append(f.hellos, h)
	return nil
}

func (f *fakeStore) ApplySnapshot(_ context.Context, _ int64, _ protocol.Hello, s []protocol.Service, _ time.Time) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.snapshots = append(f.snapshots, s)
	return nil
}

func (f *fakeStore) UpsertService(_ context.Context, _ int64, cluster string, s protocol.Service, _ time.Time) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.upserts = append(f.upserts, cluster+"|"+s.Key())
	return nil
}

func (f *fakeStore) RemoveService(_ context.Context, _ int64, cluster, key string, _ time.Time) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.removes = append(f.removes, cluster+"|"+key)
	return nil
}

func (f *fakeStore) MarkSeen(_ context.Context, id int64, _ time.Time) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.seen = append(f.seen, id)
	return nil
}

func (f *fakeStore) count(get func(*fakeStore) int) func() int {
	return func() int {
		f.mu.Lock()
		defer f.mu.Unlock()
		return get(f)
	}
}

type fakePub struct {
	mu     sync.Mutex
	topics []events.Topic
}

func (p *fakePub) Publish(t events.Topic) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.topics = append(p.topics, t)
}

func (p *fakePub) has(t events.Topic) bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	for _, x := range p.topics {
		if x == t {
			return true
		}
	}
	return false
}

// reset forgets every topic published so far.
func (p *fakePub) reset() {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.topics = nil
}

var (
	node2   = model.Agent{ID: 2, Name: "node-2"}
	mgr     = model.Agent{ID: 1, Name: "swarm-mgr"}
	revoked = model.Agent{ID: 3, Name: "old", RevokedAt: ptr(time.Now())}
)

func ptr[T any](v T) *T { return &v }

func newHub(t *testing.T, store *fakeStore) (*Hub, *fakePub, *httptest.Server) {
	t.Helper()
	pub := &fakePub{}
	h := New(store, pub, slog.New(slog.NewTextHandler(io.Discard, nil)))
	srv := httptest.NewServer(h)
	t.Cleanup(srv.Close)
	return h, pub, srv
}

func dial(t *testing.T, srv *httptest.Server, token string) (*websocket.Conn, *http.Response, error) {
	t.Helper()
	header := http.Header{}
	if token != "" {
		header.Set("Authorization", "Bearer "+token)
	}
	return websocket.Dial(context.Background(), "ws"+strings.TrimPrefix(srv.URL, "http"),
		&websocket.DialOptions{HTTPHeader: header})
}

func send(t *testing.T, c *websocket.Conn, typ string, payload any) {
	t.Helper()
	env, err := protocol.Encode(typ, payload)
	require.NoError(t, err)
	require.NoError(t, wsjson.Write(context.Background(), c, env))
}

func TestHub_RejectsBadTokens(t *testing.T) {
	_, _, srv := newHub(t, newFakeStore(node2, revoked))
	for name, token := range map[string]string{"missing": "", "unknown": "pua_nobody", "revoked": "pua_old"} {
		t.Run(name, func(t *testing.T) {
			_, resp, err := dial(t, srv, token)
			require.Error(t, err)
			require.NotNil(t, resp)
			assert.Equal(t, http.StatusUnauthorized, resp.StatusCode)
		})
	}
}

func TestHub_HelloAndSnapshotReachStore(t *testing.T) {
	store := newFakeStore(node2)
	h, pub, srv := newHub(t, store)
	c, _, err := dial(t, srv, "pua_node-2")
	require.NoError(t, err)
	defer c.CloseNow()

	send(t, c, protocol.TypeHello, protocol.Hello{Hostname: "node-2"})
	send(t, c, protocol.TypeSnapshot, protocol.Snapshot{Services: []protocol.Service{
		{Kind: protocol.KindContainer, ContainerName: "a", Image: "a:1", Status: "running"},
	}})

	require.Eventually(t, func() bool { return store.count(func(f *fakeStore) int { return len(f.snapshots) })() == 1 }, time.Second, 5*time.Millisecond)
	assert.True(t, h.Online(node2.ID))
	assert.True(t, pub.has(events.Agents))
	assert.True(t, pub.has(events.Services))
}

func TestHub_LargeSnapshotArrivesIntact(t *testing.T) {
	store := newFakeStore(node2)
	_, _, srv := newHub(t, store)
	c, _, err := dial(t, srv, "pua_node-2")
	require.NoError(t, err)
	defer c.CloseNow()

	services := make([]protocol.Service, 500)
	for i := range services {
		services[i] = protocol.Service{
			Kind: protocol.KindContainer, ContainerName: strings.Repeat("x", 20) + string(rune('a'+i%26)) + strings.Repeat("y", i%7),
			Image: "registry.example.com/some/long/image/name:v1.2.3@sha256:" + strings.Repeat("0", 64), Status: "running",
		}
	}
	send(t, c, protocol.TypeHello, protocol.Hello{Hostname: "node-2"})
	send(t, c, protocol.TypeSnapshot, protocol.Snapshot{Services: services})

	require.Eventually(t, func() bool { return store.count(func(f *fakeStore) int { return len(f.snapshots) })() == 1 }, time.Second, 5*time.Millisecond)
	store.mu.Lock()
	defer store.mu.Unlock()
	assert.Len(t, store.snapshots[0], 500)
}

func TestHub_UpsertAndRemoveCarryClusterID(t *testing.T) {
	store := newFakeStore(mgr)
	_, _, srv := newHub(t, store)
	c, _, err := dial(t, srv, "pua_swarm-mgr")
	require.NoError(t, err)
	defer c.CloseNow()

	gitea := protocol.Service{Kind: protocol.KindSwarm, SwarmService: "gitea_gitea", Image: "gitea/gitea:1", Status: "running"}
	send(t, c, protocol.TypeHello, protocol.Hello{Hostname: "swarm-mgr", SwarmManager: true, SwarmClusterID: "c1"})
	send(t, c, protocol.TypeUpsert, protocol.Upsert{Service: gitea})
	send(t, c, protocol.TypeRemove, protocol.Remove{Key: gitea.Key()})

	require.Eventually(t, func() bool { return store.count(func(f *fakeStore) int { return len(f.removes) })() == 1 }, time.Second, 5*time.Millisecond)
	store.mu.Lock()
	defer store.mu.Unlock()
	assert.Equal(t, []string{"c1|swarm|gitea_gitea"}, store.upserts)
	assert.Equal(t, []string{"c1|swarm|gitea_gitea"}, store.removes)
}

func TestHub_IgnoresSwarmUpsertFromWorker(t *testing.T) {
	store := newFakeStore(node2)
	_, _, srv := newHub(t, store)
	c, _, err := dial(t, srv, "pua_node-2")
	require.NoError(t, err)
	defer c.CloseNow()

	send(t, c, protocol.TypeHello, protocol.Hello{Hostname: "node-2"})
	send(t, c, protocol.TypeUpsert, protocol.Upsert{Service: protocol.Service{Kind: protocol.KindSwarm, SwarmService: "x"}})
	send(t, c, protocol.TypeUpsert, protocol.Upsert{Service: protocol.Service{Kind: protocol.KindContainer, ContainerName: "y"}})

	require.Eventually(t, func() bool { return store.count(func(f *fakeStore) int { return len(f.upserts) })() == 1 }, time.Second, 5*time.Millisecond)
	store.mu.Lock()
	defer store.mu.Unlock()
	assert.Equal(t, []string{"|container|y"}, store.upserts)
}

// Live swarm messages are gated exactly like snapshots: only a swarm manager
// with a cluster ID may upsert or remove swarm services.
func TestHub_SwarmLiveMessagesAreGatedOnManagerWithClusterID(t *testing.T) {
	tests := []struct {
		name        string
		hello       protocol.Hello
		wantReaches bool
	}{
		{"non-manager", protocol.Hello{Hostname: "swarm-mgr"}, false},
		{"non-manager with cluster id", protocol.Hello{Hostname: "swarm-mgr", SwarmClusterID: "c1"}, false},
		{"manager with empty cluster id", protocol.Hello{Hostname: "swarm-mgr", SwarmManager: true}, false},
		{"manager with cluster id", protocol.Hello{Hostname: "swarm-mgr", SwarmManager: true, SwarmClusterID: "c1"}, true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			store := newFakeStore(mgr)
			_, _, srv := newHub(t, store)
			c, _, err := dial(t, srv, "pua_swarm-mgr")
			require.NoError(t, err)
			defer c.CloseNow()

			svc := protocol.Service{Kind: protocol.KindSwarm, SwarmService: "gitea_gitea"}
			sentinel := protocol.Service{Kind: protocol.KindContainer, ContainerName: "sentinel"}
			send(t, c, protocol.TypeHello, tt.hello)
			send(t, c, protocol.TypeUpsert, protocol.Upsert{Service: svc})
			send(t, c, protocol.TypeRemove, protocol.Remove{Key: svc.Key()})
			// Messages are handled in order, so once the sentinel's remove has
			// arrived, the swarm messages before it have been handled too.
			send(t, c, protocol.TypeRemove, protocol.Remove{Key: sentinel.Key()})

			sentinelRemoved := func() bool {
				store.mu.Lock()
				defer store.mu.Unlock()
				for _, r := range store.removes {
					if strings.HasSuffix(r, "|container|sentinel") {
						return true
					}
				}
				return false
			}
			require.Eventually(t, sentinelRemoved, time.Second, 5*time.Millisecond)

			store.mu.Lock()
			defer store.mu.Unlock()
			if tt.wantReaches {
				assert.Equal(t, []string{"c1|swarm|gitea_gitea"}, store.upserts)
				assert.Contains(t, store.removes, "c1|swarm|gitea_gitea")
			} else {
				assert.Empty(t, store.upserts)
				assert.Equal(t, []string{tt.hello.SwarmClusterID + "|container|sentinel"}, store.removes)
			}
		})
	}
}

func TestHub_ClosesConnectionThatSkipsHello(t *testing.T) {
	store := newFakeStore(node2)
	_, _, srv := newHub(t, store)
	c, _, err := dial(t, srv, "pua_node-2")
	require.NoError(t, err)
	defer c.CloseNow()

	send(t, c, protocol.TypeSnapshot, protocol.Snapshot{})
	_, _, err = c.Read(context.Background())
	assert.Equal(t, websocket.StatusPolicyViolation, websocket.CloseStatus(err))
}

// lockedBuffer is a bytes.Buffer safe for the hub's handler goroutines to
// write while the test reads it.
type lockedBuffer struct {
	mu  sync.Mutex
	buf bytes.Buffer
}

func (b *lockedBuffer) Write(p []byte) (int, error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.buf.Write(p)
}

func (b *lockedBuffer) String() string {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.buf.String()
}

func TestHub_NewConnectionReplacesOld(t *testing.T) {
	store := newFakeStore(node2)
	logs := &lockedBuffer{}
	h := New(store, &fakePub{}, slog.New(slog.NewTextHandler(logs, nil)))
	srv := httptest.NewServer(h)
	t.Cleanup(srv.Close)
	first, _, err := dial(t, srv, "pua_node-2")
	require.NoError(t, err)
	defer first.CloseNow()
	send(t, first, protocol.TypeHello, protocol.Hello{Hostname: "node-2"})
	require.Eventually(t, func() bool { return h.Online(node2.ID) }, time.Second, 5*time.Millisecond)

	second, _, err := dial(t, srv, "pua_node-2")
	require.NoError(t, err)
	defer second.CloseNow()

	_, _, err = first.Read(context.Background())
	assert.Error(t, err, "the replaced connection must be closed")
	assert.True(t, h.Online(node2.ID), "the agent stays online on its new connection")
	assert.Contains(t, logs.String(), "msg=agent_replaced")
	assert.Contains(t, logs.String(), "agent=node-2")
	assert.NotContains(t, logs.String(), "pua_", "tokens are never logged")
}

func TestHub_DisconnectCutsLiveConnection(t *testing.T) {
	store := newFakeStore(node2)
	h, pub, srv := newHub(t, store)
	c, _, err := dial(t, srv, "pua_node-2")
	require.NoError(t, err)
	defer c.CloseNow()
	send(t, c, protocol.TypeHello, protocol.Hello{Hostname: "node-2"})
	require.Eventually(t, func() bool { return h.Online(node2.ID) }, time.Second, 5*time.Millisecond)
	// Online turns true before the hello is handled, so wait for the hello's
	// own Agents signal, then forget it; the assertion below can then only be
	// satisfied by the disconnect.
	require.Eventually(t, func() bool { return pub.has(events.Agents) }, time.Second, 5*time.Millisecond)
	pub.reset()

	h.Disconnect(node2.ID)

	_, _, err = c.Read(context.Background())
	assert.Error(t, err)
	require.Eventually(t, func() bool { return !h.Online(node2.ID) }, time.Second, 5*time.Millisecond)
	require.Eventually(t, func() bool { return store.count(func(f *fakeStore) int { return len(f.seen) })() == 1 }, time.Second, 5*time.Millisecond)
	require.Eventually(t, func() bool { return pub.has(events.Agents) }, time.Second, 5*time.Millisecond)
}

func TestHub_RevokeBetweenLookupAndRegisterCutsConnection(t *testing.T) {
	store := newFakeStore(node2)
	store.revokeFromLookup = 2 // first lookup passes, the re-check sees the revocation
	h, _, srv := newHub(t, store)
	c, _, err := dial(t, srv, "pua_node-2")
	require.NoError(t, err)
	defer c.CloseNow()

	_, _, err = c.Read(context.Background())
	assert.Equal(t, websocket.StatusPolicyViolation, websocket.CloseStatus(err))
	require.Eventually(t, func() bool { return !h.Online(node2.ID) }, time.Second, 5*time.Millisecond)
}

func TestHub_ClosesConnectionOnSecondHello(t *testing.T) {
	store := newFakeStore(mgr)
	_, _, srv := newHub(t, store)
	c, _, err := dial(t, srv, "pua_swarm-mgr")
	require.NoError(t, err)
	defer c.CloseNow()

	send(t, c, protocol.TypeHello, protocol.Hello{Hostname: "swarm-mgr"})
	send(t, c, protocol.TypeHello, protocol.Hello{Hostname: "swarm-mgr", SwarmManager: true, SwarmClusterID: "c1"})

	_, _, err = c.Read(context.Background())
	assert.Equal(t, websocket.StatusPolicyViolation, websocket.CloseStatus(err))
	assert.Equal(t, 1, store.count(func(f *fakeStore) int { return len(f.hellos) })())
}