package hub

import (
	"context"
	"errors"
	"strconv"
	"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/master/events"
	"github.com/shcizo/package-updater/internal/master/model"
	"github.com/shcizo/package-updater/internal/protocol"
)

var t0 = time.Date(2026, 10, 12, 12, 0, 0, 0, time.UTC)

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

func readEnv(t *testing.T, c *websocket.Conn) protocol.Envelope {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	var env protocol.Envelope
	require.NoError(t, wsjson.Read(ctx, c, &env))
	return env
}

func waitRoute(t *testing.T, h *Hub) model.Route {
	t.Helper()
	svc := model.Service{Kind: protocol.KindCompose, AgentID: ptr(node2.ID)}
	var r model.Route
	require.Eventually(t, func() bool {
		var err error
		r, err = h.Route(svc, protocol.MinJobVersion)
		return err == nil
	}, 2*time.Second, 5*time.Millisecond)
	return r
}

func TestHub_ReconcilesBeforeRoutable(t *testing.T) {
	store := newFakeStore(node2)
	store.lostGate = make(chan struct{})
	store.lostResult = []int64{5}
	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", ProtocolVersion: 2, KnownJobs: []string{"7", "9", "x"}})
	svc := model.Service{Kind: protocol.KindCompose, AgentID: ptr(node2.ID)}
	time.Sleep(50 * time.Millisecond) // hello is now blocked inside FailLostJobs
	_, err = h.Route(svc, protocol.MinJobVersion)
	assert.ErrorIs(t, err, model.ErrAgentOffline)

	close(store.lostGate)
	waitRoute(t, h)

	store.mu.Lock()
	defer store.mu.Unlock()
	assert.Equal(t, [][]int64{{7, 9}}, store.lostCalls)
	assert.True(t, pub.has(events.Jobs))
	assert.True(t, pub.has(events.JobTopic(5)))
}

func TestHub_StartJobSendsEnvelope(t *testing.T) {
	store := newFakeStore(node2)
	h, _, 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{ProtocolVersion: 2})
	r := waitRoute(t, h)

	start := protocol.JobStart{Kind: protocol.KindCompose, ComposeProject: "slusk", ComposeService: "web", Ref: "r:1"}
	require.NoError(t, h.StartJob(r, 42, start))

	env := readEnv(t, c)
	assert.Equal(t, protocol.TypeJobStart, env.Type)
	assert.Equal(t, "42", env.ID)
	got, err := protocol.Decode[protocol.JobStart](env)
	require.NoError(t, err)
	assert.Equal(t, start, got)
}

func TestHub_StartJobOnReplacedConnectionFails(t *testing.T) {
	store := newFakeStore(node2)
	h, _, srv := newHub(t, store)
	first, _, err := dial(t, srv, "pua_node-2")
	require.NoError(t, err)
	defer first.CloseNow()
	send(t, first, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
	old := waitRoute(t, h)

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

	// dial returns before the hub has registered the connection, so wait for
	// the replacement to take effect: the old route is then refused. Until the
	// new connection says hello, nothing is routable.
	require.Eventually(t, func() bool {
		return errors.Is(h.StartJob(old, 1, protocol.JobStart{}), model.ErrAgentOffline)
	}, 2*time.Second, 5*time.Millisecond)
	send(t, second, protocol.TypeHello, protocol.Hello{ProtocolVersion: 2})
	fresh := waitRoute(t, h)
	require.NotEqual(t, old.ConnID, fresh.ConnID)
	assert.ErrorIs(t, h.StartJob(old, 1, protocol.JobStart{}), model.ErrAgentOffline)

	require.NoError(t, h.StartJob(fresh, 2, protocol.JobStart{Ref: "r:2"}))
	env := readEnv(t, second)
	assert.Equal(t, "2", env.ID, "the new connection got only the job sent to it")
}

func TestHub_ProgressAndResult(t *testing.T) {
	store := newFakeStore(node2)
	store.progressOK, store.finishOK = true, true
	_, 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{ProtocolVersion: 2})

	sendJob(t, c, protocol.TypeJobProgress, "42", protocol.JobProgress{Output: "out"})
	require.Eventually(t, func() bool { return pub.has(events.JobTopic(42)) }, time.Second, 5*time.Millisecond)
	store.mu.Lock()
	assert.Equal(t, []string{"2|42|out"}, store.progress)
	store.mu.Unlock()

	pub.reset()
	res := protocol.JobResult{Status: protocol.ResultSucceeded, Ref: "r:1@sha256:d", Output: "done", StartedAt: t0, FinishedAt: t0}
	sendJob(t, c, protocol.TypeJobResult, "42", res)
	ack := readEnv(t, c)
	assert.Equal(t, protocol.TypeJobAck, ack.Type)
	assert.Equal(t, "42", ack.ID)
	assert.True(t, pub.has(events.Services))
	assert.True(t, pub.has(events.JobTopic(42)))
	store.mu.Lock()
	require.Len(t, store.finishes, 1)
	assert.Equal(t, model.JobSucceeded, store.finishes[0].Status)
	assert.Equal(t, "r:1@sha256:d", store.finishes[0].Ref)
	store.mu.Unlock()

	// Not applied (already recorded, or not this agent's): still acked.
	store.mu.Lock()
	store.finishOK = false
	store.mu.Unlock()
	pub.reset()
	sendJob(t, c, protocol.TypeJobResult, "43", res)
	assert.Equal(t, "43", readEnv(t, c).ID)
	assert.False(t, pub.has(events.JobTopic(43)))

	// Non-numeric ID: acked, no store call.
	store.mu.Lock()
	before := len(store.finishes)
	store.mu.Unlock()
	sendJob(t, c, protocol.TypeJobResult, "abc", res)
	assert.Equal(t, "abc", readEnv(t, c).ID)
	store.mu.Lock()
	assert.Len(t, store.finishes, before)
	store.mu.Unlock()

	// Unknown status is stored as a failure.
	store.mu.Lock()
	store.finishOK = true
	store.mu.Unlock()
	sendJob(t, c, protocol.TypeJobResult, "44", protocol.JobResult{Status: "weird", Ref: "r:9"})
	assert.Equal(t, "44", readEnv(t, c).ID)
	store.mu.Lock()
	defer store.mu.Unlock()
	last := store.finishes[len(store.finishes)-1]
	assert.Equal(t, model.JobFailed, last.Status)
	assert.Equal(t, `agent sent invalid status "weird"`, last.Error)
	assert.Empty(t, last.Ref)
}

func TestHub_ResultStoreErrorClosesConnection(t *testing.T) {
	store := newFakeStore(node2)
	store.finishErr = assert.AnError
	_, _, 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{ProtocolVersion: 2})
	sendJob(t, c, protocol.TypeJobResult, "42", protocol.JobResult{Status: protocol.ResultFailed})

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	_, _, err = c.Read(ctx)
	assert.Equal(t, websocket.StatusPolicyViolation, websocket.CloseStatus(err), "no ack, connection closed")
}

func TestHub_IgnoresUnknownType(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{ProtocolVersion: 2})
	require.NoError(t, wsjson.Write(context.Background(), c, protocol.Envelope{Type: "future.thing"}))
	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)
}

func TestHub_LargeResultArrives(t *testing.T) {
	store := newFakeStore(node2)
	store.finishOK = true
	_, _, 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{ProtocolVersion: 2})

	out := strings.Repeat("\x1b", protocol.MaxJobOutput)
	sendJob(t, c, protocol.TypeJobResult, "42", protocol.JobResult{Status: protocol.ResultFailed, Output: out})
	assert.Equal(t, "42", readEnv(t, c).ID)
	store.mu.Lock()
	defer store.mu.Unlock()
	require.Len(t, store.finishes, 1)
	assert.Equal(t, out, store.finishes[0].Output)
}

func TestHub_ConcurrentWritesAreSerialised(t *testing.T) {
	const n = 50
	store := newFakeStore(node2)
	store.finishOK = true
	h, _, 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{ProtocolVersion: 2})
	r := waitRoute(t, h)

	var wg sync.WaitGroup
	for i := 0; i < n; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			assert.NoError(t, h.StartJob(r, int64(1000+i), protocol.JobStart{Ref: "r:1"}))
		}()
	}
	go func() {
		for i := 0; i < n; i++ {
			sendJob(t, c, protocol.TypeJobResult, strconv.Itoa(i), protocol.JobResult{Status: protocol.ResultFailed})
		}
	}()

	starts, acks := 0, 0
	for starts+acks < 2*n {
		switch env := readEnv(t, c); env.Type {
		case protocol.TypeJobStart:
			starts++
		case protocol.TypeJobAck:
			acks++
		default:
			t.Fatalf("unexpected %q", env.Type)
		}
	}
	wg.Wait()
	assert.Equal(t, n, starts)
	assert.Equal(t, n, acks)
}