package e2e

import (
	"context"
	"crypto/tls"
	"errors"
	"io"
	"log/slog"
	"net/http"
	"strconv"
	"strings"
	"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/agent/conn"
	"github.com/shcizo/package-updater/internal/agent/updater"
	"github.com/shcizo/package-updater/internal/agenttoken"
	"github.com/shcizo/package-updater/internal/master/jobs"
	"github.com/shcizo/package-updater/internal/master/model"
	"github.com/shcizo/package-updater/internal/protocol"
)

const (
	slusk       = "gitea.shcizo.se/shcizo/slusk"
	agentID     = int64(2)
	waitFor     = 3 * time.Second
	pollEvery   = 10 * time.Millisecond
	lostMessage = "lost: the agent did not know this job when it reconnected"
)

// e2eDigest is the digest the fake executor pins a tag to.
const e2eDigest = "sha256:" + "e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0e2e0"

// gateExecutor writes a line of output, then blocks until release is closed
// or the job is cancelled.
type gateExecutor struct {
	started chan string // receives the job's Ref when it begins
	release chan struct{}
}

func newGateExecutor() *gateExecutor {
	return &gateExecutor{started: make(chan string, 8), release: make(chan struct{})}
}

func (g *gateExecutor) Execute(ctx context.Context, job updater.Job, out io.Writer) (string, error) {
	_, _ = io.WriteString(out, "pulling "+job.Ref+"\n")
	g.started <- job.Ref
	select {
	case <-g.release:
	case <-ctx.Done():
		return "", ctx.Err()
	}
	_, _ = io.WriteString(out, "recreated\n")
	if strings.Contains(job.Ref, "@") {
		return job.Ref, nil
	}
	return job.Ref + "@" + e2eDigest, nil
}

// jobAgent is a connected agent with a real runner and its two lifetimes: the
// connection and the runner can be stopped separately.
type jobAgent struct {
	runner   *updater.Runner
	stopConn func()
	stopRun  func()
}

// startJobAgent starts an agent. before, if given, runs on the new runner
// before the worker and the connection start, like ResumeSelfUpdates in the
// real agent.
func startJobAgent(t *testing.T, m testMaster, token string, exec updater.Executor, before ...func(*updater.Runner)) *jobAgent {
	t.Helper()
	src := &staticSource{}
	src.set(web)
	runner := updater.NewRunner(exec, time.Minute)
	for _, f := range before {
		f(runner)
	}
	runCtx, stopRun := context.WithCancel(context.Background())
	runDone := make(chan struct{})
	go func() { runner.Run(runCtx); close(runDone) }()

	a := conn.New(conn.Config{URL: m.url, Token: token, Fingerprint: m.fingerprint, Version: "e2e",
		MinBackoff: 10 * time.Millisecond, MaxBackoff: 50 * time.Millisecond, ProgressInterval: 10 * time.Millisecond,
	}, src, runner, slog.New(slog.NewTextHandler(io.Discard, nil)))
	connCtx, cancelConn := context.WithCancel(context.Background())
	connDone := make(chan struct{})
	go func() { _ = a.Run(connCtx); close(connDone) }()

	ja := &jobAgent{
		runner:   runner,
		stopConn: func() { cancelConn(); <-connDone },
		stopRun:  func() { stopRun(); <-runDone },
	}
	t.Cleanup(func() { ja.stopConn(); ja.stopRun() })
	return ja
}

// jobEnv is a master with one agent (ID 2, token "pua_ok") that owns a Compose
// service, and a jobs service wired to the master's current hub.
type jobEnv struct {
	store *memStore
	m     testMaster
	jobs  *jobs.Service
	exec  *gateExecutor
	agent *jobAgent
}

func newJobEnv(t *testing.T) *jobEnv {
	t.Helper()
	store := &memStore{agents: map[string]model.Agent{agenttoken.Hash("pua_ok"): {ID: agentID, Name: "node-2"}}}
	store.seedService(model.Service{ID: 1, Kind: protocol.KindCompose, ComposeProject: "slusk", ComposeService: "web",
		AgentID: ptr(agentID), ImageRepo: slusk, ActualRef: slusk + ":1.4.2@sha256:old"})
	m := startMaster(t, store)
	exec := newGateExecutor()
	env := &jobEnv{store: store, m: m, exec: exec,
		jobs:  jobs.New(store, m.hub, m.broker, "gitea.shcizo.se/shcizo/package-updater:latest", slog.New(slog.NewTextHandler(io.Discard, nil))),
		agent: startJobAgent(t, m, "pua_ok", exec)}
	env.waitRoutable(t)
	return env
}

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

// waitRoutable waits until the master can send jobs to the agent.
func (e *jobEnv) waitRoutable(t *testing.T) {
	t.Helper()
	svc, err := e.store.GetService(context.Background(), 1)
	require.NoError(t, err)
	require.Eventually(t, func() bool { _, err := e.m.hub.Route(svc, protocol.MinJobVersion); return err == nil }, waitFor, pollEvery)
}

func (e *jobEnv) job(t *testing.T, id int64) model.Job {
	t.Helper()
	j, err := e.store.GetJob(context.Background(), id)
	require.NoError(t, err)
	return j
}

func (e *jobEnv) waitStatus(t *testing.T, id int64, want model.JobStatus) model.Job {
	t.Helper()
	require.Eventually(t, func() bool { return e.job(t, id).Status == want }, waitFor, pollEvery,
		"job %d should become %s, is %s (%s)", id, want, e.job(t, id).Status, e.job(t, id).Error)
	return e.job(t, id)
}

func (e *jobEnv) waitStarted(t *testing.T) {
	t.Helper()
	select {
	case <-e.exec.started:
	case <-time.After(waitFor):
		t.Fatal("the agent never started the job")
	}
}

func (e *jobEnv) trigger(t *testing.T) model.Job {
	t.Helper()
	j, err := e.jobs.Trigger(context.Background(), 1, "sub-1")
	require.NoError(t, err)
	return j
}

func TestManualJobFlowsToAgentAndBack(t *testing.T) {
	e := newJobEnv(t)
	close(e.exec.release)

	j := e.trigger(t)
	done := e.waitStatus(t, j.ID, model.JobSucceeded)

	assert.Equal(t, slusk+":1.4.2@"+e2eDigest, done.ToRef)
	assert.Contains(t, done.Output, "pulling "+slusk+":1.4.2")
	assert.Contains(t, done.Output, "recreated")
	svc, err := e.store.GetService(context.Background(), 1)
	require.NoError(t, err)
	assert.Equal(t, slusk+":1.4.2@"+e2eDigest, svc.DesiredRef)
	require.Eventually(t, func() bool { return len(e.agent.runner.Known()) == 0 }, waitFor, pollEvery, "the result must be acked")
}

func TestJobFinishedWhileDisconnectedIsRecordedOnce(t *testing.T) {
	e := newJobEnv(t)
	j := e.trigger(t)
	e.waitStarted(t)

	e.m.restart() // drops the agent's connection while the executor blocks
	close(e.exec.release)

	done := e.waitStatus(t, j.ID, model.JobSucceeded)
	assert.NotEqual(t, lostMessage, done.Error)
	assert.GreaterOrEqual(t, e.m.attempts.Load(), int64(2), "the agent must have reconnected")
	require.Eventually(t, func() bool { return len(e.agent.runner.Known()) == 0 }, waitFor, pollEvery)
	e.store.mu.Lock()
	defer e.store.mu.Unlock()
	assert.Equal(t, 1, e.store.finishApplied)
}

func TestMasterRestartKeepsJobTheAgentKnows(t *testing.T) {
	e := newJobEnv(t)
	j := e.trigger(t)
	e.waitStarted(t)

	e.m.restart()
	e.waitRoutable(t) // the agent reconnected and said hello with its known jobs
	assert.True(t, e.job(t, j.ID).Active(), "the agent still has the job, so it is not lost")
	close(e.exec.release)

	done := e.waitStatus(t, j.ID, model.JobSucceeded)
	assert.Empty(t, done.Error)
}

func TestAgentRestartMarksJobLost(t *testing.T) {
	e := newJobEnv(t)
	j := e.trigger(t)
	e.waitStarted(t)

	// The agent process dies: connection first, so the cancelled job cannot
	// report a result, then the runner.
	e.agent.stopConn()
	e.agent.stopRun()
	fresh := startJobAgent(t, e.m, "pua_ok", newGateExecutor())
	_ = fresh

	lost := e.waitStatus(t, j.ID, model.JobFailed)
	assert.Equal(t, lostMessage, lost.Error)
}

// TestAgentSelfUpdateResumedByNewAgent follows a self-update: the old agent
// dies while its job runs, and the new agent finishes the job under the same
// ID before it says hello, so the master neither loses nor re-runs it.
func TestAgentSelfUpdateResumedByNewAgent(t *testing.T) {
	pinned := slusk + ":1.4.2@" + e2eDigest
	tests := []struct {
		name      string
		resume    updater.ResumeFunc
		wantError string
	}{
		{"helper succeeded", func(_ context.Context, out io.Writer) (string, error) {
			_, _ = io.WriteString(out, "helper output\n")
			return pinned, nil
		}, ""},
		{"helper failed", func(context.Context, io.Writer) (string, error) {
			return "", errors.New("docker compose up: exit status 1 (output: boom)")
		}, "docker compose up: exit status 1 (output: boom)"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			e := newJobEnv(t)
			j := e.trigger(t)
			e.waitStarted(t)
			started := time.Now().Add(-time.Second)

			e.agent.stopConn()
			e.agent.stopRun()
			startJobAgent(t, e.m, "pua_ok", newGateExecutor(), func(r *updater.Runner) {
				r.Resume(strconv.FormatInt(j.ID, 10), started, tt.resume)
			})

			if tt.wantError != "" {
				done := e.waitStatus(t, j.ID, model.JobFailed)
				assert.Equal(t, tt.wantError, done.Error)
				assert.NotEqual(t, lostMessage, done.Error)
				return
			}
			done := e.waitStatus(t, j.ID, model.JobSucceeded)
			assert.Empty(t, done.Error)
			assert.Equal(t, pinned, done.ToRef)
			assert.Contains(t, done.Output, "helper output")
			svc, err := e.store.GetService(context.Background(), 1)
			require.NoError(t, err)
			assert.Equal(t, pinned, svc.DesiredRef)
		})
	}
}

func TestPhaseOneAgentGetsNoJobs(t *testing.T) {
	store := &memStore{agents: map[string]model.Agent{agenttoken.Hash("pua_old"): {ID: agentID, Name: "node-2"}}}
	svcRow := model.Service{ID: 1, Kind: protocol.KindCompose, ComposeProject: "slusk", ComposeService: "web",
		AgentID: ptr(agentID), ActualRef: slusk + ":1.4.2@sha256:old"}
	store.seedService(svcRow)
	m := startMaster(t, store)
	js := jobs.New(store, m.hub, m.broker, "gitea.shcizo.se/shcizo/package-updater:latest", slog.New(slog.NewTextHandler(io.Discard, nil)))

	// A raw client speaking the phase-1 protocol: its hello has no
	// protocol_version.
	ctx := context.Background()
	c, _, err := websocket.Dial(ctx, m.url, &websocket.DialOptions{
		HTTPHeader: http.Header{"Authorization": {"Bearer pua_old"}},
		HTTPClient: &http.Client{Transport: &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}},
	})
	require.NoError(t, err)
	defer c.CloseNow()
	hello, err := protocol.Encode(protocol.TypeHello, map[string]any{"agent_version": "1", "hostname": "node-2", "swarm_manager": false})
	require.NoError(t, err)
	require.NoError(t, wsjson.Write(ctx, c, hello))
	snap, err := protocol.Encode(protocol.TypeSnapshot, protocol.Snapshot{Services: []protocol.Service{web}})
	require.NoError(t, err)
	require.NoError(t, wsjson.Write(ctx, c, snap))

	require.Eventually(t, func() bool {
		_, err := m.hub.Route(svcRow, protocol.MinJobVersion)
		return errors.Is(err, model.ErrAgentTooOld)
	}, waitFor, pollEvery)

	j, err := js.Trigger(ctx, 1, "sub-1")
	require.NoError(t, err)
	assert.Equal(t, model.JobFailed, j.Status)
	assert.Equal(t, "agent too old for updates", j.Error)

	rctx, cancel := context.WithTimeout(ctx, 200*time.Millisecond)
	defer cancel()
	var env protocol.Envelope
	err = wsjson.Read(rctx, c, &env)
	assert.True(t, err != nil && strings.Contains(err.Error(), "deadline"), "no job.start expected, got %q (err %v)", env.Type, err)
}