package executor

import (
	"context"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/bluesky-social/indigo/atproto/syntax"
	"github.com/gorilla/websocket"

	"tangled.org/core/api/tangled"
	"tangled.org/core/notifier"
	"tangled.org/core/spindle/config"
	"tangled.org/core/spindle/db"
	"tangled.org/core/spindle/engine"
	"tangled.org/core/spindle/mill/proto"
	millv1 "tangled.org/core/spindle/mill/proto/gen"
	"tangled.org/core/spindle/models"
	"tangled.org/core/spindle/secrets"
)

type captureEncoder struct {
	messages chan *millproto.Message
}

func newCaptureEncoder() *captureEncoder {
	return &captureEncoder{messages: make(chan *millproto.Message, 4)}
}

func (e *captureEncoder) Encode(msg *millproto.Message) error {
	e.messages <- msg
	return nil
}

type fakeSlot struct{ released int }

func (s *fakeSlot) Release() { s.released++ }

// fakeEngine records which methods got called and never acquires a slot of its
// own; the wrapper must supply the slot.
type fakeEngine struct {
	setupCalled   bool
	runCalled     bool
	destroyCalled bool
	acquireCalled bool
	secrets       chan []secrets.UnlockedSecret
	done          chan struct{}
}

func (e *fakeEngine) InitWorkflow(twf tangled.Pipeline_Workflow, tpl tangled.Pipeline) (*models.Workflow, error) {
	return &models.Workflow{Name: twf.Name}, nil
}
func (e *fakeEngine) SetupWorkflow(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, l models.WorkflowLogger) error {
	e.setupCalled = true
	return nil
}
func (e *fakeEngine) WorkflowTimeout() time.Duration { return 7 * time.Minute }
func (e *fakeEngine) DestroyWorkflow(ctx context.Context, wid models.WorkflowId) error {
	e.destroyCalled = true
	if e.done != nil {
		close(e.done)
	}
	return nil
}
func (e *fakeEngine) RunStep(ctx context.Context, wid models.WorkflowId, w *models.Workflow, idx int, s []secrets.UnlockedSecret, l models.WorkflowLogger) error {
	e.runCalled = true
	if e.secrets != nil {
		e.secrets <- s
	}
	return nil
}

func (e *fakeEngine) AcquireWorkflowSlot(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, mode engine.AcquireMode) (engine.WorkflowSlot, error) {
	e.acquireCalled = true
	return engine.NoopSlot{}, nil
}

type fakeStep struct{}

func (fakeStep) Name() string          { return "test" }
func (fakeStep) Command() string       { return "true" }
func (fakeStep) Kind() models.StepKind { return models.StepKindUser }
func TestReservedEngineHandsBackHeldSlotOnce(t *testing.T) {
	inner := &fakeEngine{}
	slot := &fakeSlot{}
	re := newReservedEngine(inner, slot)

	got, err := re.(engine.WorkflowSlotter).AcquireWorkflowSlot(context.Background(), models.WorkflowId{}, nil, engine.Wait)
	if err != nil {
		t.Fatalf("AcquireWorkflowSlot() error = %v", err)
	}
	if got != engine.WorkflowSlot(slot) {
		t.Fatal("AcquireWorkflowSlot() did not return the held slot")
	}
	if inner.acquireCalled {
		t.Fatal("wrapper must not call the inner engine's AcquireWorkflowSlot")
	}

	if _, err := re.(engine.WorkflowSlotter).AcquireWorkflowSlot(context.Background(), models.WorkflowId{}, nil, engine.Wait); err == nil {
		t.Fatal("second AcquireWorkflowSlot() should error")
	}
}

func TestReservedEngineForwardsLifecycle(t *testing.T) {
	inner := &fakeEngine{}
	re := newReservedEngine(inner, &fakeSlot{})

	if d := re.WorkflowTimeout(); d != 7*time.Minute {
		t.Fatalf("WorkflowTimeout() = %v, want 7m", d)
	}
	if err := re.SetupWorkflow(context.Background(), models.WorkflowId{}, &models.Workflow{}, models.NullLogger{}); err != nil {
		t.Fatal(err)
	}
	if err := re.RunStep(context.Background(), models.WorkflowId{}, &models.Workflow{Steps: []models.Step{}}, 0, nil, models.NullLogger{}); err != nil {
		t.Fatal(err)
	}
	if err := re.DestroyWorkflow(context.Background(), models.WorkflowId{}); err != nil {
		t.Fatal(err)
	}
	if !inner.setupCalled || !inner.runCalled || !inner.destroyCalled {
		t.Fatalf("lifecycle not forwarded: %+v", inner)
	}
}

func TestHandleCommitIsIdempotent(t *testing.T) {
	enc := newCaptureEncoder()

	e := &Executor{
		l:      slog.New(slog.NewTextHandler(io.Discard, nil)),
		enc:    enc,
		active: map[string]*reservation{"lease-1": {leaseID: "lease-1", committed: true}},
	}

	e.handleCommit(context.Background(), &millv1.CommitLease{LeaseId: "lease-1"})
	msg := <-enc.messages
	if got := msg.GetCommitted().GetLeaseId(); got != "lease-1" {
		t.Fatalf("Committed lease = %q, want lease-1", got)
	}
}

func TestHandleCommitRejectsMissingReservation(t *testing.T) {
	enc := newCaptureEncoder()
	e := &Executor{
		l:      slog.New(slog.NewTextHandler(io.Discard, nil)),
		enc:    enc,
		active: make(map[string]*reservation),
	}

	e.handleCommit(context.Background(), &millv1.CommitLease{LeaseId: "expired"})
	result := (<-enc.messages).GetReserveResult()
	if result == nil {
		t.Fatal("missing reservation commit did not receive a ReserveResult")
	}
	if result.GetLeaseId() != "expired" || result.GetAccepted() {
		t.Fatalf("ReserveResult = %+v, want correlated rejection", result)
	}
}

func TestHandleCommitPreservesPreauthorizedSecrets(t *testing.T) {
	d, err := db.Make(context.Background(), filepath.Join(t.TempDir(), "spindle.db"))
	if err != nil {
		t.Fatalf("db.Make() error = %v", err)
	}
	t.Cleanup(func() { d.Close() })
	n := notifier.New()
	enc := newCaptureEncoder()
	inner := &fakeEngine{
		secrets: make(chan []secrets.UnlockedSecret, 1),
		done:    make(chan struct{}),
	}
	repoDid, err := syntax.ParseDID("did:plc:abcdefghijklmnopqrstuvwxyz")
	if err != nil {
		t.Fatalf("ParseDID() error = %v", err)
	}
	res := &reservation{
		leaseID:    "lease-1",
		wid:        models.WorkflowId{PipelineId: models.PipelineId{Knot: "did:plc:test", Rkey: "pipeline"}, Name: "build"},
		realEngine: inner,
		slot:       engine.NoopSlot{},
		wf:         &models.Workflow{Name: "build", Steps: []models.Step{fakeStep{}}},
		repoDid:    repoDid,
	}
	e := &Executor{
		cfg:    &config.Config{Server: config.Server{LogDir: t.TempDir()}},
		db:     d,
		n:      &n,
		l:      slog.New(slog.NewTextHandler(io.Discard, nil)),
		enc:    enc,
		relay:  newRelayLog(),
		active: map[string]*reservation{res.leaseID: res},
	}

	e.handleCommit(context.Background(), &millv1.CommitLease{
		LeaseId: res.leaseID,
		Secrets: []*millv1.Secret{{Key: "TOKEN", Value: "secret-value"}},
	})
	if got := (<-enc.messages).GetCommitted().GetLeaseId(); got != res.leaseID {
		t.Fatalf("Committed lease = %q, want %q", got, res.leaseID)
	}
	select {
	case got := <-inner.secrets:
		if len(got) != 1 || got[0].Key != "TOKEN" || got[0].Value != "secret-value" {
			t.Fatalf("RunStep secrets = %+v", got)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("RunStep did not receive CommitLease secrets")
	}
	select {
	case <-inner.done:
	case <-time.After(2 * time.Second):
		t.Fatal("workflow did not finish")
	}
	if res.stopTail != nil {
		res.stopTail()
	}
}

func TestRunSessionCancellationClosesStalledWebsocket(t *testing.T) {
	connected := make(chan struct{})
	release := make(chan struct{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := websocket.Upgrade(w, r, nil, 1024, 1024)
		if err != nil {
			return
		}
		defer conn.Close()
		close(connected)
		<-release
	}))
	t.Cleanup(func() {
		close(release)
		srv.Close()
	})

	e := testSessionExecutor("ws" + strings.TrimPrefix(srv.URL, "http"))
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- e.runSession(ctx) }()
	<-connected
	cancel()

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("runSession did not return after context cancellation")
	}
}

func TestRunSessionConsumesAcksWhileReplaying(t *testing.T) {
	e := testSessionExecutor("")
	e.relay.appendLog("lease", []byte("first"))
	e.relay.appendLog("lease", make([]byte, 7<<20))

	serverErr := make(chan error, 1)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := websocket.Upgrade(w, r, nil, 1024, 1024)
		if err != nil {
			serverErr <- err
			return
		}
		defer conn.Close()
		if tcp, ok := conn.UnderlyingConn().(interface{ SetReadBuffer(int) error }); ok {
			_ = tcp.SetReadBuffer(1024)
		}
		stream := millproto.NewWSStream(conn)
		enc := millproto.NewEncoder(stream)
		dec := millproto.NewDecoder(stream)
		if _, err := dec.Decode(); err != nil {
			serverErr <- err
			return
		}
		if err := enc.Encode(&millproto.Message{Resume: &millv1.Resume{}}); err != nil {
			serverErr <- err
			return
		}
		first, err := dec.Decode()
		if err != nil {
			serverErr <- err
			return
		}
		if got := offsetOf(first); got != 1 {
			serverErr <- fmt.Errorf("first replay offset = %d, want 1", got)
			return
		}
		if err := enc.Encode(&millproto.Message{Ack: &millv1.Ack{UpToOffset: 1}}); err != nil {
			serverErr <- err
			return
		}

		deadline := time.Now().Add(10 * time.Second)
		for {
			e.relay.mu.Lock()
			acked := e.relay.acked
			e.relay.mu.Unlock()
			if acked >= 1 {
				break
			}
			if time.Now().After(deadline) {
				serverErr <- context.DeadlineExceeded
				return
			}
			time.Sleep(time.Millisecond)
		}

		serverErr <- nil
	}))
	t.Cleanup(srv.Close)
	e.millURL = "ws" + strings.TrimPrefix(srv.URL, "http")

	done := make(chan error, 1)
	go func() { done <- e.runSession(context.Background()) }()
	select {
	case err := <-serverErr:
		if err != nil {
			t.Fatalf("server: %v", err)
		}
	case <-time.After(20 * time.Second):
		t.Fatal("replay stalled while the peer waited for its acknowledgement to be consumed")
	}
	select {
	case <-done:
	case <-time.After(5 * time.Second):
		t.Fatal("runSession did not stop after peer closed")
	}
}

func testSessionExecutor(url string) *Executor {
	return &Executor{
		millURL: url,
		seats:   1,
		engines: make(map[string]models.Engine),
		cfg:     &config.Config{Server: config.Server{Dev: true}},
		l:       slog.New(slog.NewTextHandler(io.Discard, nil)),
		relay:   newRelayLog(),
		active:  make(map[string]*reservation),
	}
}

func TestFinishJobReportsCancelledReservationAsCancelled(t *testing.T) {
	res := &reservation{
		leaseID:   "lease-1",
		wid:       models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"},
		cancelled: true,
	}
	e := &Executor{
		relay:  newRelayLog(),
		active: map[string]*reservation{res.leaseID: res},
	}

	e.finishJob(res, &tangled.PipelineStatus{
		Pipeline: string(res.wid.PipelineId.AtUri()),
		Workflow: res.wid.Name,
		Status:   string(models.StatusKindFailed),
	})

	if len(e.relay.entries) != 1 {
		t.Fatalf("relay entries = %d, want 1", len(e.relay.entries))
	}
	got := e.relay.entries[0].GetAttemptResult().GetTerminalStatus()
	if got != string(models.StatusKindCancelled) {
		t.Fatalf("terminal status = %q, want cancelled", got)
	}
}
