//go:build integration

package store

import (
	"context"
	"strings"
	"sync"
	"testing"
	"time"

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

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

// jobFixture seeds one agent with one Compose service and returns their IDs.
func jobFixture(t *testing.T, s *Store, agentName string) (agentID, serviceID int64) {
	t.Helper()
	ctx := context.Background()
	agentID = newAgent(t, s, agentName)
	require.NoError(t, s.ApplySnapshot(ctx, agentID, worker, []protocol.Service{slusk}, t0))
	svcs, err := s.ListServices(ctx)
	require.NoError(t, err)
	for _, svc := range svcs {
		if svc.AgentID != nil && *svc.AgentID == agentID {
			return agentID, svc.ID
		}
	}
	t.Fatal("service not found")
	return
}

func newJob(serviceID, agentID int64) model.NewJob {
	return model.NewJob{
		ServiceID: serviceID, AgentID: &agentID,
		FromRef: "r:1@sha256:a", ToRef: "r:1", Trigger: model.TriggerManual,
		RequestedBy: "sub", CreatedAt: t0,
	}
}

func TestCreateJob_OneActivePerService(t *testing.T) {
	s := openTestStore(t)
	ctx := context.Background()
	agentID, svcID := jobFixture(t, s, "a")

	var wg sync.WaitGroup
	jobs := make([]model.Job, 2)
	errs := make([]error, 2)
	for i := range jobs {
		wg.Add(1)
		go func() {
			defer wg.Done()
			jobs[i], errs[i] = s.CreateJob(ctx, newJob(svcID, agentID))
		}()
	}
	wg.Wait()

	assert.Equal(t, jobs[0].ID, jobs[1].ID)
	var failures int
	for _, err := range errs {
		if err != nil {
			assert.ErrorIs(t, err, model.ErrJobActive)
			failures++
		}
	}
	assert.Equal(t, 1, failures)

	ok, err := s.FailJob(ctx, jobs[0].ID, "x", t0)
	require.NoError(t, err)
	assert.True(t, ok)
	ok, err = s.FailJob(ctx, jobs[0].ID, "x", t0)
	require.NoError(t, err)
	assert.False(t, ok)

	next, err := s.CreateJob(ctx, newJob(svcID, agentID))
	require.NoError(t, err)
	assert.NotEqual(t, jobs[0].ID, next.ID)
}

func TestFailLostJobs(t *testing.T) {
	s := openTestStore(t)
	ctx := context.Background()
	a, svcA := jobFixture(t, s, "a")
	b := newAgent(t, s, "b")
	require.NoError(t, s.ApplySnapshot(ctx, b, worker, []protocol.Service{adhoc}, t0))
	all, err := s.ListServices(ctx)
	require.NoError(t, err)
	var svcB int64
	for _, svc := range all {
		if svc.ID != svcA {
			svcB = svc.ID
		}
	}

	// Only one active job per service, so A's jobs go through the lifecycle.
	j3, err := s.CreateJob(ctx, newJob(svcA, a))
	require.NoError(t, err)
	ok, err := s.FinishJob(ctx, a, j3.ID, model.JobFinish{Status: model.JobSucceeded, Ref: "r:1@sha256:d", StartedAt: t0, FinishedAt: t0})
	require.NoError(t, err)
	require.True(t, ok)
	j1, err := s.CreateJob(ctx, newJob(svcA, a))
	require.NoError(t, err)
	j4, err := s.CreateJob(ctx, newJob(svcB, b))
	require.NoError(t, err)
	_, err = s.JobProgress(ctx, b, j4.ID, "x", t0)
	require.NoError(t, err)

	ids, err := s.FailLostJobs(ctx, a, []int64{j3.ID}, t0)
	require.NoError(t, err)
	assert.Equal(t, []int64{j1.ID}, ids)

	got, err := s.GetJob(ctx, j1.ID)
	require.NoError(t, err)
	assert.Equal(t, model.JobFailed, got.Status)
	assert.Contains(t, got.Error, "lost")
	got, err = s.GetJob(ctx, j3.ID)
	require.NoError(t, err)
	assert.Equal(t, model.JobSucceeded, got.Status)
	got, err = s.GetJob(ctx, j4.ID)
	require.NoError(t, err)
	assert.Equal(t, model.JobRunning, got.Status)

	// A running job the agent still knows is kept.
	j5, err := s.CreateJob(ctx, newJob(svcA, a))
	require.NoError(t, err)
	_, err = s.JobProgress(ctx, a, j5.ID, "", t0)
	require.NoError(t, err)
	ids, err = s.FailLostJobs(ctx, a, []int64{j5.ID}, t0)
	require.NoError(t, err)
	assert.Empty(t, ids)
}

func TestJobProgress_OwnerOnly(t *testing.T) {
	s := openTestStore(t)
	ctx := context.Background()
	a, svc := jobFixture(t, s, "a")
	b := newAgent(t, s, "b")
	j, err := s.CreateJob(ctx, newJob(svc, a))
	require.NoError(t, err)

	started := t0.Add(time.Second)
	ok, err := s.JobProgress(ctx, a, j.ID, "hello ", started)
	require.NoError(t, err)
	assert.True(t, ok)
	got, err := s.GetJob(ctx, j.ID)
	require.NoError(t, err)
	assert.Equal(t, model.JobRunning, got.Status)
	require.NotNil(t, got.StartedAt)
	assert.True(t, started.Equal(*got.StartedAt))

	ok, err = s.JobProgress(ctx, b, j.ID, "evil", t0.Add(time.Minute))
	require.NoError(t, err)
	assert.False(t, ok)
	got, err = s.GetJob(ctx, j.ID)
	require.NoError(t, err)
	assert.Equal(t, "hello ", got.Output)

	ok, err = s.JobProgress(ctx, a, j.ID, "world", t0.Add(time.Minute))
	require.NoError(t, err)
	assert.True(t, ok)
	got, err = s.GetJob(ctx, j.ID)
	require.NoError(t, err)
	assert.Equal(t, "hello world", got.Output)
	assert.True(t, started.Equal(*got.StartedAt))

	_, err = s.FailJob(ctx, j.ID, "x", t0)
	require.NoError(t, err)
	ok, err = s.JobProgress(ctx, a, j.ID, "late", t0)
	require.NoError(t, err)
	assert.False(t, ok)
}

func TestJobProgress_KeepsTail(t *testing.T) {
	s := openTestStore(t)
	ctx := context.Background()
	a, svc := jobFixture(t, s, "a")
	j, err := s.CreateJob(ctx, newJob(svc, a))
	require.NoError(t, err)

	_, err = s.JobProgress(ctx, a, j.ID, strings.Repeat("a", protocol.MaxJobOutput), t0)
	require.NoError(t, err)
	got, err := s.GetJob(ctx, j.ID)
	require.NoError(t, err)
	assert.False(t, got.OutputTruncated)

	_, err = s.JobProgress(ctx, a, j.ID, "TAIL", t0)
	require.NoError(t, err)
	got, err = s.GetJob(ctx, j.ID)
	require.NoError(t, err)
	assert.True(t, got.OutputTruncated)
	assert.Len(t, got.Output, protocol.MaxJobOutput)
	assert.True(t, strings.HasSuffix(got.Output, "aaTAIL"))
}

func TestFinishJob_OnlyOnceAndOnlyByOwner(t *testing.T) {
	s := openTestStore(t)
	ctx := context.Background()
	a, svc := jobFixture(t, s, "a")
	b := newAgent(t, s, "b")
	j, err := s.CreateJob(ctx, newJob(svc, a))
	require.NoError(t, err)
	_, err = s.JobProgress(ctx, a, j.ID, "partial", t0)
	require.NoError(t, err)

	fin := model.JobFinish{Status: model.JobSucceeded, Ref: "r:1@sha256:d", Output: "full", StartedAt: t0.Add(time.Hour), FinishedAt: t0.Add(2 * time.Hour)}
	ok, err := s.FinishJob(ctx, b, j.ID, fin)
	require.NoError(t, err)
	assert.False(t, ok)

	ok, err = s.FinishJob(ctx, a, j.ID, fin)
	require.NoError(t, err)
	assert.True(t, ok)
	got, err := s.GetJob(ctx, j.ID)
	require.NoError(t, err)
	assert.Equal(t, model.JobSucceeded, got.Status)
	assert.Equal(t, "r:1@sha256:d", got.ToRef)
	assert.Equal(t, "full", got.Output)
	assert.True(t, fin.StartedAt.Equal(*got.StartedAt), "the agent's started_at wins over the master's progress clock")
	assert.True(t, fin.FinishedAt.Equal(*got.FinishedAt))
	service, err := s.GetService(ctx, svc)
	require.NoError(t, err)
	assert.Equal(t, "r:1@sha256:d", service.DesiredRef)

	again := fin
	again.Ref = "r:2@sha256:e"
	again.Output = "other"
	ok, err = s.FinishJob(ctx, a, j.ID, again)
	require.NoError(t, err)
	assert.False(t, ok)
	got, err = s.GetJob(ctx, j.ID)
	require.NoError(t, err)
	assert.Equal(t, "full", got.Output)
	service, err = s.GetService(ctx, svc)
	require.NoError(t, err)
	assert.Equal(t, "r:1@sha256:d", service.DesiredRef)

	j2, err := s.CreateJob(ctx, newJob(svc, a))
	require.NoError(t, err)
	ok, err = s.FinishJob(ctx, a, j2.ID, model.JobFinish{Status: model.JobFailed, Error: "boom", Output: "o", StartedAt: t0, FinishedAt: t0})
	require.NoError(t, err)
	assert.True(t, ok)
	got, err = s.GetJob(ctx, j2.ID)
	require.NoError(t, err)
	assert.Equal(t, model.JobFailed, got.Status)
	assert.Equal(t, "boom", got.Error)
	assert.Equal(t, "r:1", got.ToRef)
	service, err = s.GetService(ctx, svc)
	require.NoError(t, err)
	assert.Equal(t, "r:1@sha256:d", service.DesiredRef)
}

func TestGetJobAndListJobs(t *testing.T) {
	s := openTestStore(t)
	ctx := context.Background()
	a, svcA := jobFixture(t, s, "a")
	b := newAgent(t, s, "b")
	require.NoError(t, s.ApplySnapshot(ctx, b, worker, []protocol.Service{adhoc}, t0))
	all, err := s.ListServices(ctx)
	require.NoError(t, err)
	var svcB int64
	for _, svc := range all {
		if svc.ID != svcA {
			svcB = svc.ID
		}
	}

	j1, err := s.CreateJob(ctx, newJob(svcA, a))
	require.NoError(t, err)
	_, err = s.JobProgress(ctx, a, j1.ID, "out", t0)
	require.NoError(t, err)
	_, err = s.FailJob(ctx, j1.ID, "x", t0)
	require.NoError(t, err)
	j2, err := s.CreateJob(ctx, newJob(svcA, a))
	require.NoError(t, err)
	j3, err := s.CreateJob(ctx, newJob(svcB, b))
	require.NoError(t, err)

	got, err := s.GetJob(ctx, j1.ID)
	require.NoError(t, err)
	assert.Equal(t, "slusk/web", got.ServiceName)
	assert.Equal(t, "a", got.AgentName)
	assert.Equal(t, "out", got.Output)

	_, err = s.GetJob(ctx, 99999)
	assert.ErrorIs(t, err, model.ErrNotFound)

	list, err := s.ListJobs(ctx, model.JobFilter{})
	require.NoError(t, err)
	require.Len(t, list, 3)
	assert.Equal(t, []int64{j3.ID, j2.ID, j1.ID}, []int64{list[0].ID, list[1].ID, list[2].ID})
	for _, j := range list {
		assert.Empty(t, j.Output)
	}

	list, err = s.ListJobs(ctx, model.JobFilter{ServiceID: svcA})
	require.NoError(t, err)
	assert.Len(t, list, 2)

	list, err = s.ListJobs(ctx, model.JobFilter{Limit: 1})
	require.NoError(t, err)
	assert.Len(t, list, 1)
}

func TestFinishJob_ZeroStartedAtKeepsProgressClock(t *testing.T) {
	s := openTestStore(t)
	ctx := context.Background()
	a, svc := jobFixture(t, s, "a")
	j, err := s.CreateJob(ctx, newJob(svc, a))
	require.NoError(t, err)
	_, err = s.JobProgress(ctx, a, j.ID, "x", t0)
	require.NoError(t, err)

	ok, err := s.FinishJob(ctx, a, j.ID, model.JobFinish{Status: model.JobFailed, FinishedAt: t0.Add(time.Hour)})
	require.NoError(t, err)
	require.True(t, ok)
	got, err := s.GetJob(ctx, j.ID)
	require.NoError(t, err)
	assert.True(t, t0.Equal(*got.StartedAt))
}

func TestCreateJob_RefusesRevokedAgent(t *testing.T) {
	s := openTestStore(t)
	ctx := context.Background()
	a, svc := jobFixture(t, s, "a")
	_, err := s.RevokeAgent(ctx, a, t0)
	require.NoError(t, err)

	_, err = s.CreateJob(ctx, newJob(svc, a))
	assert.ErrorIs(t, err, model.ErrAgentOffline)
	list, err := s.ListJobs(ctx, model.JobFilter{})
	require.NoError(t, err)
	assert.Empty(t, list, "no job is left behind to block the service")

	nj := newJob(svc, a)
	nj.AgentID = nil
	_, err = s.CreateJob(ctx, nj)
	assert.NoError(t, err, "a job that was never routed has no agent to check")
}

// CreateJob must wait for a revoke that is still in flight. Without the share
// lock it would read the agent as active, insert the job, and the revoke's
// job cleanup (which ran before the insert committed) would never fail it.
func TestCreateJob_WaitsForConcurrentRevoke(t *testing.T) {
	s := openTestStore(t)
	ctx := context.Background()
	a, svc := jobFixture(t, s, "a")

	tx, err := s.pool.Begin(ctx)
	require.NoError(t, err)
	defer func() { _ = tx.Rollback(ctx) }()
	_, err = tx.Exec(ctx, `UPDATE agents SET revoked_at = now() WHERE id = $1`, a)
	require.NoError(t, err)

	done := make(chan error, 1)
	go func() {
		_, err := s.CreateJob(ctx, newJob(svc, a))
		done <- err
	}()

	select {
	case err := <-done:
		t.Fatalf("CreateJob returned (%v) while the revoke was uncommitted", err)
	case <-time.After(500 * time.Millisecond):
	}

	require.NoError(t, tx.Commit(ctx))
	select {
	case err := <-done:
		assert.ErrorIs(t, err, model.ErrAgentOffline)
	case <-time.After(5 * time.Second):
		t.Fatal("CreateJob did not return after the revoke committed")
	}
	list, err := s.ListJobs(ctx, model.JobFilter{})
	require.NoError(t, err)
	assert.Empty(t, list)
}