// Package e2e runs a real agent connection against a real hub over TLS with
// a pinned key, in one process, with in-memory storage and fake Docker.
package e2e

import (
	"context"
	"crypto/tls"
	"io"
	"log"
	"log/slog"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

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

type memStore struct {
	mu        sync.Mutex
	agents    map[string]model.Agent
	lookups   int
	snapshots int
	services  map[string]protocol.Service
}

func (m *memStore) AgentByTokenHash(_ context.Context, hash string) (model.Agent, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.lookups++
	a, ok := m.agents[hash]
	if !ok {
		return model.Agent{}, model.ErrNotFound
	}
	return a, nil
}

func (m *memStore) RecordHello(context.Context, int64, protocol.Hello, time.Time) error { return nil }

func (m *memStore) ApplySnapshot(_ context.Context, _ int64, _ protocol.Hello, s []protocol.Service, _ time.Time) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.snapshots++
	m.services = map[string]protocol.Service{}
	for _, x := range s {
		m.services[x.Key()] = x
	}
	return nil
}

func (m *memStore) UpsertService(_ context.Context, _ int64, _ string, s protocol.Service, _ time.Time) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.services[s.Key()] = s
	return nil
}

func (m *memStore) RemoveService(_ context.Context, _ int64, _, key string, _ time.Time) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.services, key)
	return nil
}

func (m *memStore) MarkSeen(context.Context, int64, time.Time) error { return nil }

func (m *memStore) keys() []string {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out []string
	for k := range m.services {
		out = append(out, k)
	}
	return out
}

type staticSource struct {
	mu       sync.Mutex
	services []protocol.Service
}

func (s *staticSource) Hello(context.Context) (protocol.Hello, error) {
	return protocol.Hello{Hostname: "node-2"}, nil
}

func (s *staticSource) Snapshot(context.Context) ([]protocol.Service, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]protocol.Service(nil), s.services...), nil
}

func (s *staticSource) set(v ...protocol.Service) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.services = v
}

// testMaster is a running TLS test server with the hub mounted on /agent.
type testMaster struct {
	url         string
	fingerprint string
	hub         *hub.Hub
	attempts    *atomic.Int64 // TCP connections accepted by the server
}

func startMaster(t *testing.T, store *memStore) testMaster {
	t.Helper()
	key, err := pin.GenerateKey()
	require.NoError(t, err)
	cert, err := pin.Certificate(key)
	require.NoError(t, err)
	fingerprint, err := pin.Fingerprint(key.Public())
	require.NoError(t, err)

	h := hub.New(store, events.NewBroker(), slog.New(slog.NewTextHandler(io.Discard, nil)))
	mux := http.NewServeMux()
	mux.Handle("GET /agent", h)
	attempts := &atomic.Int64{}
	srv := httptest.NewUnstartedServer(mux)
	srv.Config.ErrorLog = log.New(io.Discard, "", 0)
	srv.Config.ConnState = func(_ net.Conn, s http.ConnState) {
		if s == http.StateNew {
			attempts.Add(1)
		}
	}
	srv.TLS = &tls.Config{Certificates: []tls.Certificate{cert}, MinVersion: tls.VersionTLS13}
	srv.StartTLS()
	t.Cleanup(srv.Close)
	return testMaster{url: "wss" + strings.TrimPrefix(srv.URL, "https") + "/agent", fingerprint: fingerprint, hub: h, attempts: attempts}
}

func runAgent(t *testing.T, url, fingerprint, token string, src conn.Source) *conn.Agent {
	t.Helper()
	a := conn.New(conn.Config{URL: url, Token: token, Fingerprint: fingerprint, Version: "e2e",
		MinBackoff: 10 * time.Millisecond, MaxBackoff: 50 * time.Millisecond,
	}, src, slog.New(slog.NewTextHandler(io.Discard, nil)))
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() { _ = a.Run(ctx); close(done) }()
	t.Cleanup(func() { cancel(); <-done })
	return a
}

var (
	web   = protocol.Service{Kind: protocol.KindCompose, ComposeProject: "slusk", ComposeService: "web", Image: "slusk:1", Status: "running"}
	cache = protocol.Service{Kind: protocol.KindContainer, ContainerName: "cache", Image: "redis:alpine", Status: "running"}
)

func TestInventoryFlowsFromAgentToStore(t *testing.T) {
	store := &memStore{agents: map[string]model.Agent{agenttoken.Hash("pua_ok"): {ID: 2, Name: "node-2"}}}
	m := startMaster(t, store)
	src := &staticSource{}
	src.set(web)
	agent := runAgent(t, m.url, m.fingerprint, "pua_ok", src)

	require.Eventually(t, func() bool { return len(store.keys()) == 1 }, 2*time.Second, 10*time.Millisecond)

	src.set(cache)
	agent.Notify()
	require.Eventually(t, func() bool {
		k := store.keys()
		return len(k) == 1 && k[0] == "container|cache"
	}, 2*time.Second, 10*time.Millisecond)
}

func TestAgentWithWrongKeyNeverReachesHub(t *testing.T) {
	store := &memStore{agents: map[string]model.Agent{agenttoken.Hash("pua_ok"): {ID: 2, Name: "node-2"}}}
	m := startMaster(t, store)
	other, err := pin.GenerateKey()
	require.NoError(t, err)
	wrong, err := pin.Fingerprint(other.Public())
	require.NoError(t, err)

	runAgent(t, m.url, wrong, "pua_ok", &staticSource{})
	require.Eventually(t, func() bool { return m.attempts.Load() >= 1 }, 2*time.Second, 10*time.Millisecond, "the agent must actually try to connect")
	time.Sleep(200 * time.Millisecond)
	store.mu.Lock()
	defer store.mu.Unlock()
	assert.Zero(t, store.lookups, "TLS must fail before the token is ever presented")
}

func TestRevokedAgentIsRejected(t *testing.T) {
	store := &memStore{agents: map[string]model.Agent{
		agenttoken.Hash("pua_old"): {ID: 3, Name: "old", RevokedAt: func() *time.Time { n := time.Now(); return &n }()},
	}}
	m := startMaster(t, store)
	src := &staticSource{}
	src.set(web)
	runAgent(t, m.url, m.fingerprint, "pua_old", src)

	require.Eventually(t, func() bool {
		store.mu.Lock()
		defer store.mu.Unlock()
		return store.lookups >= 2
	}, 2*time.Second, 10*time.Millisecond, "the agent keeps retrying")
	assert.Empty(t, store.keys())
}

func TestAgentReconnectsAndResendsSnapshot(t *testing.T) {
	store := &memStore{agents: map[string]model.Agent{agenttoken.Hash("pua_ok"): {ID: 2, Name: "node-2"}}}
	m := startMaster(t, store)
	src := &staticSource{}
	src.set(web)
	runAgent(t, m.url, m.fingerprint, "pua_ok", src)

	snapshots := func() int {
		store.mu.Lock()
		defer store.mu.Unlock()
		return store.snapshots
	}
	require.Eventually(t, func() bool { return snapshots() >= 1 }, 2*time.Second, 10*time.Millisecond)

	m.hub.Disconnect(2)

	require.Eventually(t, func() bool { return snapshots() >= 2 }, 2*time.Second, 10*time.Millisecond, "agent reconnects and sends a fresh snapshot")
	assert.Len(t, store.keys(), 1)
}