package mill

import (
	"context"
	"path/filepath"
	"testing"
	"time"

	"tangled.org/core/notifier"
	"tangled.org/core/spindle/db"
	millv1 "tangled.org/core/spindle/mill/proto/gen"
	"tangled.org/core/spindle/models"
)

func restoreTestMill(t *testing.T, cfg Config) (*Mill, *db.DB) {
	t.Helper()
	bdb, err := db.Make(context.Background(), filepath.Join(t.TempDir(), "mill.db"))
	if err != nil {
		t.Fatalf("db.Make: %v", err)
	}
	t.Cleanup(func() { bdb.Close() })
	n := notifier.New()
	m := New(discardLogger(), cfg)
	m.Attach(bdb, &n)
	return m, bdb
}

func restoredMill(t *testing.T, bdb *db.DB, cfg Config) *Mill {
	t.Helper()
	n := notifier.New()
	m := New(discardLogger(), cfg)
	m.Attach(bdb, &n)
	if err := m.RestoreState(); err != nil {
		t.Fatalf("RestoreState: %v", err)
	}
	return m
}

func TestRestoreStateRebuildsLeasesAndCursors(t *testing.T) {
	_, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})

	if err := bdb.SaveMillLease(db.MillLease{
		LeaseID: "lease-1", NodeID: "node-1", Engine: "dummy",
		Knot: "knot.example", Rkey: "rkey1", Workflow: "build", State: leaseRowRunning,
	}); err != nil {
		t.Fatalf("SaveMillLease: %v", err)
	}
	if err := bdb.SetExecutorCursor("node-1", 7); err != nil {
		t.Fatalf("SetExecutorCursor: %v", err)
	}

	m2 := restoredMill(t, bdb, Config{ReconnectGrace: time.Minute})

	m2.mu.Lock()
	lease := m2.leases["lease-1"]
	offset := m2.nodeOffset["node-1"]
	m2.mu.Unlock()

	if lease == nil {
		t.Fatal("restored mill has no lease-1")
	}
	if !lease.orphaned {
		t.Fatal("restored lease is not orphaned; a terminal would be delivered to a waiter that does not exist")
	}
	if lease.getState() != leaseRunning {
		t.Fatalf("restored lease state = %v, want leaseRunning", lease.getState())
	}
	wantWid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.example", Rkey: "rkey1"}, Name: "build"}
	if lease.wid != wantWid {
		t.Fatalf("restored lease wid = %+v, want %+v", lease.wid, wantWid)
	}
	if offset != 7 {
		t.Fatalf("restored cursor = %d, want 7", offset)
	}
}

func TestOrphanTerminalAuthorsStatusRow(t *testing.T) {
	_, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	if err := bdb.SaveMillLease(db.MillLease{
		LeaseID: "lease-1", NodeID: "node-1", Engine: "dummy",
		Knot: "knot.example", Rkey: "rkey1", Workflow: "build", State: leaseRowRunning,
	}); err != nil {
		t.Fatalf("SaveMillLease: %v", err)
	}
	if err := bdb.SetExecutorCursor("node-1", 3); err != nil {
		t.Fatalf("SetExecutorCursor: %v", err)
	}

	m := restoredMill(t, bdb, Config{ReconnectGrace: time.Minute})

	sess := newSession("node-1", nopEncoder(), discardLogger())
	resume, ok := m.attachSession(sess)
	if !ok {
		t.Fatal("attachSession rejected the reconnecting executor")
	}
	if resume != 3 {
		t.Fatalf("attachSession resume offset = %d, want restored cursor 3", resume)
	}

	m.onAttemptResult(sess, &millv1.AttemptResult{
		Offset:         4,
		LeaseId:        "lease-1",
		TerminalStatus: string(models.StatusKindSuccess),
	})

	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.example", Rkey: "rkey1"}, Name: "build"}
	st, err := bdb.GetStatus(wid)
	if err != nil {
		t.Fatalf("GetStatus after orphan terminal: %v", err)
	}
	if st.Status != string(models.StatusKindSuccess) {
		t.Fatalf("orphan terminal authored status %q, want success", st.Status)
	}

	m.mu.Lock()
	_, still := m.leases["lease-1"]
	m.mu.Unlock()
	if still {
		t.Fatal("finished orphan still in the lease map")
	}
	if rows, _ := bdb.ListMillLeases(); len(rows) != 0 {
		t.Fatalf("finished orphan still persisted: %+v", rows)
	}
}

func TestSnapshotReconciliationFailsDroppedOrphans(t *testing.T) {
	_, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	for _, l := range []db.MillLease{
		{LeaseID: "lease-kept", NodeID: "node-1", Engine: "dummy", Knot: "k", Rkey: "r1", Workflow: "w", State: leaseRowRunning},
		{LeaseID: "lease-gone", NodeID: "node-1", Engine: "dummy", Knot: "k", Rkey: "r2", Workflow: "w", State: leaseRowRunning},
	} {
		if err := bdb.SaveMillLease(l); err != nil {
			t.Fatalf("SaveMillLease(%s): %v", l.LeaseID, err)
		}
	}

	m := restoredMill(t, bdb, Config{ReconnectGrace: time.Minute})
	sess := newSession("node-1", nopEncoder(), discardLogger())
	m.attachSession(sess)

	m.onSnapshot(sess, &millv1.NodeSnapshot{
		NodeId:         "node-1",
		ActiveLeaseIds: []string{"lease-kept"},
	})

	m.mu.Lock()
	_, kept := m.leases["lease-kept"]
	_, gone := m.leases["lease-gone"]
	m.mu.Unlock()
	if !kept {
		t.Fatal("reconciliation dropped a lease the executor still holds")
	}
	if gone {
		t.Fatal("reconciliation kept a lease the executor no longer holds")
	}

	st, err := bdb.GetStatus(models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r2"}, Name: "w"})
	if err != nil {
		t.Fatalf("GetStatus for dropped orphan: %v", err)
	}
	if st.Status != string(models.StatusKindFailed) {
		t.Fatalf("dropped orphan authored status %q, want failed", st.Status)
	}
}

func TestSweepFailsOrphansOfAbsentExecutors(t *testing.T) {
	_, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	if err := bdb.SaveMillLease(db.MillLease{
		LeaseID: "lease-1", NodeID: "node-absent", Engine: "dummy",
		Knot: "k", Rkey: "r1", Workflow: "w", State: leaseRowReserved,
	}); err != nil {
		t.Fatalf("SaveMillLease: %v", err)
	}

	m := restoredMill(t, bdb, Config{ReconnectGrace: time.Minute})
	m.sweepUnclaimedOrphans()

	m.mu.Lock()
	_, still := m.leases["lease-1"]
	m.mu.Unlock()
	if still {
		t.Fatal("sweep kept an orphan whose executor never reconnected")
	}
	st, err := bdb.GetStatus(models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r1"}, Name: "w"})
	if err != nil {
		t.Fatalf("GetStatus after sweep: %v", err)
	}
	if st.Status != string(models.StatusKindFailed) {
		t.Fatalf("sweep authored status %q, want failed", st.Status)
	}
	if rows, _ := bdb.ListMillLeases(); len(rows) != 0 {
		t.Fatalf("swept orphan still persisted: %+v", rows)
	}
}

func TestAckOffsetPersistsCursor(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	sess := newSession("node-1", nopEncoder(), discardLogger())
	m.attachSession(sess)

	m.ackOffset(sess, 12)

	cursors, err := bdb.ListExecutorCursors()
	if err != nil {
		t.Fatalf("ListExecutorCursors: %v", err)
	}
	if cursors["node-1"] != 12 {
		t.Fatalf("persisted cursor = %d, want 12", cursors["node-1"])
	}
}

func TestOrphanTerminalFailureKeepsLeaseAndOffsetRetryable(t *testing.T) {
	_, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	if err := bdb.SaveMillLease(db.MillLease{
		LeaseID: "lease-1", NodeID: "node-1", Engine: "dummy",
		Knot: "knot.example", Rkey: "rkey1", Workflow: "build", State: leaseRowRunning,
	}); err != nil {
		t.Fatalf("SaveMillLease: %v", err)
	}
	if err := bdb.SetExecutorCursor("node-1", 3); err != nil {
		t.Fatalf("SetExecutorCursor: %v", err)
	}
	if _, err := bdb.Exec(`
		create trigger reject_orphan_lease_delete
		before delete on mill_leases
		begin
			select raise(abort, 'forced delete failure');
		end
	`); err != nil {
		t.Fatalf("create failure trigger: %v", err)
	}

	m := restoredMill(t, bdb, Config{ReconnectGrace: time.Minute})
	sess := newSession("node-1", nopEncoder(), discardLogger())
	if _, ok := m.attachSession(sess); !ok {
		t.Fatal("attachSession rejected reconnect")
	}
	result := &millv1.AttemptResult{
		Offset:         4,
		LeaseId:        "lease-1",
		TerminalStatus: string(models.StatusKindSuccess),
	}
	if err := m.onAttemptResult(sess, result); err == nil {
		t.Fatal("orphan terminal relay succeeded despite forced transaction failure")
	}

	m.mu.Lock()
	lease := m.leases["lease-1"]
	offset := m.nodeOffset["node-1"]
	m.mu.Unlock()
	if lease == nil || lease.getState() == leaseDone {
		t.Fatal("failed orphan completion made the in-memory lease unretryable")
	}
	if offset != 3 {
		t.Fatalf("in-memory relay offset = %d, want 3", offset)
	}
	if rows, err := bdb.ListMillLeases(); err != nil || len(rows) != 1 {
		t.Fatalf("durable leases after transaction rollback = %+v, err = %v; want retained lease", rows, err)
	}
	var events int
	if err := bdb.QueryRow(`select count(*) from events`).Scan(&events); err != nil {
		t.Fatalf("count events: %v", err)
	}
	if events != 0 {
		t.Fatalf("terminal events after transaction rollback = %d, want 0", events)
	}

	if _, err := bdb.Exec(`drop trigger reject_orphan_lease_delete`); err != nil {
		t.Fatalf("drop failure trigger: %v", err)
	}
	if err := m.onAttemptResult(sess, result); err != nil {
		t.Fatalf("retry orphan terminal: %v", err)
	}
	m.mu.Lock()
	_, still := m.leases["lease-1"]
	offset = m.nodeOffset["node-1"]
	m.mu.Unlock()
	if still {
		t.Fatal("successful orphan completion retained in-memory lease")
	}
	if offset != 4 {
		t.Fatalf("in-memory relay offset after retry = %d, want 4", offset)
	}
	if rows, err := bdb.ListMillLeases(); err != nil || len(rows) != 0 {
		t.Fatalf("durable leases after successful retry = %+v, err = %v; want none", rows, err)
	}
}
