package mill

import (
	"context"
	"io"
	"log/slog"
	"os"
	"path/filepath"
	"testing"
	"time"

	"tangled.org/core/notifier"
	"tangled.org/core/spindle/db"
	"tangled.org/core/spindle/models"

	millproto "tangled.org/core/spindle/mill/proto"
	millv1 "tangled.org/core/spindle/mill/proto/gen"
)

func discardLogger() *slog.Logger {
	return slog.New(slog.NewTextHandler(io.Discard, nil))
}

func nopEncoder() scriptedEncoder {
	return scriptedEncoder(func(*millproto.Message) error { return nil })
}

// HashToken is the storage/lookup transform for executor credentials. It must be
// deterministic (so a token minted once keeps authenticating), collision-free
// across distinct tokens, and it must never echo the raw token (so a DB leak of
// the hash never yields a usable credential).
func TestHashToken(t *testing.T) {
	const raw = "super-secret-executor-token"

	if HashToken(raw) != HashToken(raw) {
		t.Fatal("HashToken is not deterministic; the same token would stop authenticating")
	}
	if HashToken("token-a") == HashToken("token-b") {
		t.Fatal("HashToken collided two distinct tokens")
	}
	if HashToken(raw) == raw {
		t.Fatal("HashToken returned the raw token; a hash leak would expose a usable credential")
	}
}

// GenerateToken must hand out distinct, non-empty tokens on every call; a
// generator that repeats or emits empties would let one executor's credential
// collide with another's.
func TestGenerateTokenDistinct(t *testing.T) {
	const n = 100
	seen := make(map[string]struct{}, n)
	for i := range n {
		tok, err := GenerateToken()
		if err != nil {
			t.Fatalf("GenerateToken: %v", err)
		}
		if tok == "" {
			t.Fatalf("GenerateToken returned an empty token on call %d", i)
		}
		if _, dup := seen[tok]; dup {
			t.Fatalf("GenerateToken repeated a token after %d calls: %q", i, tok)
		}
		seen[tok] = struct{}{}
	}
}

// attachSession is the anti-hijack gate: a valid token must not be able to
// displace a live executor that already holds that node identity. Only once the
// incumbent is disconnected (in its reconnect grace window) may a reconnect
// adopt the identity.
func TestAttachSessionRejectsSecondLiveSession(t *testing.T) {
	l := discardLogger()
	m := New(l, Config{ReconnectGrace: time.Minute})

	sessionOf := func(node string) *millSession {
		m.mu.Lock()
		defer m.mu.Unlock()
		return m.sessions[node]
	}

	sess1 := newSession("node-1", nopEncoder(), l)
	if _, ok := m.attachSession(sess1); !ok {
		t.Fatal("first attach of a node was rejected; want accept")
	}
	if sessionOf("node-1") != sess1 {
		t.Fatal("first session was not registered as the live session")
	}

	// a second attach for the same node while sess1 is still live must be
	// rejected and must not evict the incumbent.
	sess2 := newSession("node-1", nopEncoder(), l)
	if _, ok := m.attachSession(sess2); ok {
		t.Fatal("second live attach for an already-live node was accepted; a valid token hijacked the executor")
	}
	if sessionOf("node-1") != sess1 {
		t.Fatal("rejected newcomer evicted the incumbent session")
	}

	// once the incumbent is disconnected (reconnect grace), a reconnect adopts
	// the identity.
	m.detachSession(sess1)
	sess3 := newSession("node-1", nopEncoder(), l)
	if _, ok := m.attachSession(sess3); !ok {
		t.Fatal("attach during the incumbent's reconnect grace was rejected; want adopt")
	}
	if sessionOf("node-1") != sess3 {
		t.Fatal("adopted session was not installed as the live session")
	}
}

// onAttemptResult must only wake the RunStep of a lease the reporting session
// actually owns. A terminal relayed for another node's lease, or for a lease
// that does not exist, must not be delivered.
func TestOnAttemptResultIgnoresForeignLease(t *testing.T) {
	l := discardLogger()
	m := New(l, Config{ReconnectGrace: time.Minute})

	foreign := newLease("lease-foreign", "node-a", "dummy")
	m.mu.Lock()
	m.leases[foreign.id] = foreign
	m.mu.Unlock()

	// node-b reports a terminal for a lease owned by node-a.
	sessB := newSession("node-b", nopEncoder(), l)
	m.attachSession(sessB)
	m.onAttemptResult(sessB, &millv1.AttemptResult{
		Offset:         1,
		LeaseId:        foreign.id,
		TerminalStatus: string(models.StatusKindSuccess),
	})

	if _, ok := pollTerminal(foreign); ok {
		t.Fatal("attempt-result on a foreign lease delivered a terminal; an executor forged another node's job result")
	}
	if foreign.getState() == leaseDone {
		t.Fatal("attempt-result on a foreign lease sealed the lease")
	}
}

func TestOnAttemptResultIgnoresAbsentLease(t *testing.T) {
	l := discardLogger()
	m := New(l, Config{ReconnectGrace: time.Minute})

	// a bystander lease the reporting node does own, to prove an absent-lease
	// relay does not spill onto some other lease.
	bystander := newLease("lease-bystander", "node-b", "dummy")
	m.mu.Lock()
	m.leases[bystander.id] = bystander
	m.mu.Unlock()

	sessB := newSession("node-b", nopEncoder(), l)
	m.attachSession(sessB)

	// no lease with this id exists: must be a no-op (and must not panic).
	m.onAttemptResult(sessB, &millv1.AttemptResult{
		Offset:         1,
		LeaseId:        "lease-nonexistent",
		TerminalStatus: string(models.StatusKindSuccess),
	})

	if _, ok := pollTerminal(bystander); ok {
		t.Fatal("attempt-result for an absent lease delivered a terminal to a bystander lease")
	}
	if bystander.getState() == leaseDone {
		t.Fatal("attempt-result for an absent lease sealed a bystander lease")
	}
}

// The owned-lease path proves the ignore tests above are not passing merely
// because delivery is broken: a correctly owned terminal IS delivered.
func TestOnAttemptResultDeliversOwnedLease(t *testing.T) {
	l := discardLogger()
	m := New(l, Config{ReconnectGrace: time.Minute})

	owned := newLease("lease-owned", "node-b", "dummy")
	m.mu.Lock()
	m.leases[owned.id] = owned
	m.mu.Unlock()

	sessB := newSession("node-b", nopEncoder(), l)
	m.attachSession(sessB)
	m.onAttemptResult(sessB, &millv1.AttemptResult{
		Offset:         1,
		LeaseId:        owned.id,
		TerminalStatus: string(models.StatusKindSuccess),
	})

	res, ok := pollTerminal(owned)
	if !ok {
		t.Fatal("attempt-result on an owned lease was not delivered")
	}
	if got := res.GetTerminalStatus(); got != string(models.StatusKindSuccess) {
		t.Fatalf("delivered terminal status = %q, want %q", got, string(models.StatusKindSuccess))
	}
	if owned.getState() != leaseDone {
		t.Fatal("owned lease was not sealed after its terminal was delivered")
	}
}

// onLogRelay must only write a log line for a lease the reporting session owns.
// A line relayed for another node's lease is dropped; a line for an owned lease
// is written verbatim to that lease's log file.
func TestOnLogRelayRelayOwnership(t *testing.T) {
	l := discardLogger()
	dir := t.TempDir()
	m := New(l, Config{LogDir: dir, ReconnectGrace: time.Minute})

	foreign := newLease("lease-foreign", "node-x", "dummy")
	foreign.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "foreign"}, Name: "build"}
	owned := newLease("lease-owned", "node-z", "dummy")
	owned.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "owned"}, Name: "build"}
	m.mu.Lock()
	m.leases[foreign.id] = foreign
	m.leases[owned.id] = owned
	m.mu.Unlock()

	sessY := newSession("node-y", nopEncoder(), l)
	m.attachSession(sessY)
	sessZ := newSession("node-z", nopEncoder(), l)
	m.attachSession(sessZ)

	// node-y relays a log line for a lease owned by node-x: dropped, no file.
	m.onLogRelay(sessY, &millv1.LogLine{
		Offset:  1,
		LeaseId: foreign.id,
		RawJson: []byte(`{"line":"forged"}`),
	})
	if _, err := os.Stat(models.LogFilePath(dir, foreign.wid)); !os.IsNotExist(err) {
		t.Fatalf("log relay for a foreign lease wrote a file (stat err=%v); an executor forged another node's logs", err)
	}

	// node-z relays a log line for its own lease: written verbatim.
	line := []byte(`{"line":"legit"}`)
	m.onLogRelay(sessZ, &millv1.LogLine{
		Offset:  1,
		LeaseId: owned.id,
		RawJson: line,
	})
	got, err := os.ReadFile(models.LogFilePath(dir, owned.wid))
	if err != nil {
		t.Fatalf("owned log relay did not write the log file: %v", err)
	}
	if want := string(line) + "\n"; string(got) != want {
		t.Fatalf("owned log line = %q, want %q", string(got), want)
	}
}

// onStatusRelay must only author a status row for a lease the reporting session
// owns. A status relayed for another node's lease is ignored; a status for an
// owned lease is re-authored under the lease's (authoritative) workflow id.
func TestOnStatusRelayRelayOwnership(t *testing.T) {
	ctx := context.Background()
	l := discardLogger()

	bdb, err := db.Make(ctx, 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(l, Config{ReconnectGrace: time.Minute})
	m.Attach(bdb, &n)

	foreign := newLease("lease-foreign", "node-x", "dummy")
	foreign.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "foreign"}, Name: "build"}
	owned := newLease("lease-owned", "node-z", "dummy")
	owned.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "owned"}, Name: "build"}
	m.mu.Lock()
	m.leases[foreign.id] = foreign
	m.leases[owned.id] = owned
	m.mu.Unlock()

	sessY := newSession("node-y", nopEncoder(), l)
	m.attachSession(sessY)
	sessZ := newSession("node-z", nopEncoder(), l)
	m.attachSession(sessZ)

	// node-y relays a status for a lease owned by node-x: no row authored.
	m.onStatusRelay(sessY, &millv1.StatusEvent{
		Offset:  1,
		LeaseId: foreign.id,
		Status:  "running",
	})
	if _, err := bdb.GetStatus(foreign.wid); err == nil {
		t.Fatal("status relay for a foreign lease authored a status row; an executor forged another pipeline's status")
	}

	// node-z relays a status for its own lease: authored under the lease's wid.
	m.onStatusRelay(sessZ, &millv1.StatusEvent{
		Offset:  1,
		LeaseId: owned.id,
		Status:  "running",
	})
	st, err := bdb.GetStatus(owned.wid)
	if err != nil {
		t.Fatalf("owned status relay did not author a status row: %v", err)
	}
	if st.Status != "running" {
		t.Fatalf("owned status = %q, want %q", st.Status, "running")
	}
}
