package cli

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

	"git.cloudinit.dev/coreci/orca/internal/certpaths"
	"git.cloudinit.dev/coreci/orca/internal/model"
	"git.cloudinit.dev/coreci/orca/internal/store"
)

// mockDrainExec is a record-and-replay execer for the drain commands
// (same pattern as internal/stepca mockExec). It matches each incoming
// command against a list of (substring, output, exitCode) responses;
// the first match wins. An entry with an empty substring matches any
// command. The exit code is 0 (success) unless explicitly set; a
// non-zero code is returned as an *sshExitErr so stopAlloc can treat
// code 5 ("unit not loaded") as idempotent success.
type mockDrainExec struct {
	mu        sync.Mutex
	responses []mockDrainResp
	calls     []mockDrainCall
}

type mockDrainResp struct {
	match string
	out   string
	exit  int
}

type mockDrainCall struct {
	peer string
	cmd  string
}

func (m *mockDrainExec) Exec(_ context.Context, peer, cmd string) ([]byte, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.calls = append(m.calls, mockDrainCall{peer: peer, cmd: cmd})
	for _, r := range m.responses {
		if r.match == "" || strings.Contains(cmd, r.match) {
			if r.exit != 0 {
				return []byte(r.out), &sshExitErr{code: r.exit}
			}
			return []byte(r.out), nil
		}
	}
	return nil, nil
}

func (m *mockDrainExec) callsFor(match string) []mockDrainCall {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out []mockDrainCall
	for _, c := range m.calls {
		if strings.Contains(c.cmd, match) {
			out = append(out, c)
		}
	}
	return out
}

func (m *mockDrainExec) countCalls(match string) int {
	return len(m.callsFor(match))
}

// drainTestEnv wires a mockDrainExec into drainExecOverride and returns
// the mock + a cleanup func. Tests MUST defer the cleanup.
func drainTestEnv(t *testing.T) *mockDrainExec {
	t.Helper()
	prev := drainExecOverride
	mx := &mockDrainExec{}
	drainExecOverride = mx
	t.Cleanup(func() { drainExecOverride = prev })
	return mx
}

// drainNodeForTest inserts a node with a fixed id+name and returns it,
// so drain commands can target it by name. Uses the test ORCA_HOME db.
func drainNodeForTest(t *testing.T, name, addr string) *model.Node {
	t.Helper()
	db, err := store.Open(certpaths.DBPath())
	if err != nil {
		t.Fatalf("open db: %v", err)
	}
	defer db.Close()
	n := &model.Node{
		ID:       "node-" + name,
		Name:     name,
		Address:  addr,
		State:    model.NodeStateReady,
		JoinedAt: time.Now().UTC(),
		LastSeen: time.Now().UTC(),
		Kind:     string(model.NodeKindLinux),
	}
	if err := store.NewNodeRepo(db).Insert(context.Background(), n); err != nil {
		t.Fatalf("insert node: %v", err)
	}
	return n
}

func TestNodeDrain_StopsAllocsAndMarksDrained(t *testing.T) {
	_, cleanup := initTestEnv(t)
	defer cleanup()
	resetRootFlags(t)

	node := drainNodeForTest(t, "drainee", "drainee:8443")
	drainTestEnv(t) // sets drainExecOverride (reset in cleanup)
	// Scripted exec: first list-units returns 2 running allocs; the
	// post-stop poll returns empty so waitAllocsStopped completes.
	mx := &scriptedDrainExec{}
	mx.queue("list-units", "orca-alloc-web-0.service loaded active running\norca-alloc-web-1.service loaded active running\n", 0)
	mx.queue("systemctl stop orca-alloc-web-0", "", 0)
	mx.queue("systemctl stop orca-alloc-web-1", "", 0)
	mx.queue("list-units", "", 0)
	drainExecOverride = mx

	var buf bytes.Buffer
	rootCmd.SetOut(&buf)
	rootCmd.SetErr(&buf)
	nodeDrainCmd.SetOut(&buf)
	nodeDrainCmd.SetErr(&buf)
	rootCmd.SetArgs([]string{"node", "drain", node.Name, "--timeout", "5s"})
	if err := rootCmd.Execute(); err != nil {
		t.Fatalf("node drain: %v", err)
	}
	out := buf.String()
	if !strings.Contains(out, "drained") {
		t.Errorf("expected drained message, got: %s", out)
	}
	if mx.countCalls("systemctl stop orca-alloc-web-0") == 0 || mx.countCalls("systemctl stop orca-alloc-web-1") == 0 {
		t.Errorf("expected stop commands for both allocs, calls: %+v", mx.calls)
	}

	// Node state should be drained in the DB.
	db, _ := store.Open(certpaths.DBPath())
	defer db.Close()
	got, err := store.NewNodeRepo(db).Get(context.Background(), node.ID)
	if err != nil {
		t.Fatalf("get node: %v", err)
	}
	if got.State != model.NodeStateDrained {
		t.Errorf("node state = %q, want drained", got.State)
	}
}

func TestNodeDrain_NoAllocs_MarksDrainedImmediately(t *testing.T) {
	_, cleanup := initTestEnv(t)
	defer cleanup()
	resetRootFlags(t)

	node := drainNodeForTest(t, "emptynode", "emptynode:8443")
	mx := &scriptedDrainExec{}
	mx.queue("list-units", "", 0) // no allocs
	drainExecOverride = mx

	var buf bytes.Buffer
	rootCmd.SetOut(&buf)
	rootCmd.SetErr(&buf)
	nodeDrainCmd.SetOut(&buf)
	nodeDrainCmd.SetErr(&buf)
	rootCmd.SetArgs([]string{"node", "drain", node.Name, "--timeout", "5s"})
	if err := rootCmd.Execute(); err != nil {
		t.Fatalf("node drain: %v", err)
	}
	out := buf.String()
	if !strings.Contains(out, "drained") {
		t.Errorf("expected drained message, got: %s", out)
	}
	// Should not have issued any stop commands.
	if mx.countCalls("systemctl stop") != 0 {
		t.Errorf("expected no stop commands, got: %+v", mx.calls)
	}

	db, _ := store.Open(certpaths.DBPath())
	defer db.Close()
	got, _ := store.NewNodeRepo(db).Get(context.Background(), node.ID)
	if got.State != model.NodeStateDrained {
		t.Errorf("node state = %q, want drained", got.State)
	}
}

func TestNodeDrain_RespectsTimeout(t *testing.T) {
	_, cleanup := initTestEnv(t)
	defer cleanup()
	resetRootFlags(t)

	node := drainNodeForTest(t, "stucknode", "stucknode:8443")
	// The alloc never leaves the running list → waitAllocsStopped hits
	// the timeout. The drain reports partial and leaves the node in
	// "draining".
	mx := &scriptedDrainExec{}
	mx.queueAlways("list-units", "orca-alloc-stuck-0.service loaded active running\n", 0)
	mx.queueAlways("systemctl stop", "", 0)
	drainExecOverride = mx

	// Poll interval is 500ms; use a 1s timeout so the test is fast but
	// still exercises the timeout path.
	var buf bytes.Buffer
	rootCmd.SetOut(&buf)
	rootCmd.SetErr(&buf)
	nodeDrainCmd.SetOut(&buf)
	nodeDrainCmd.SetErr(&buf)
	start := time.Now()
	rootCmd.SetArgs([]string{"node", "drain", node.Name, "--timeout", "1s"})
	err := rootCmd.Execute()
	elapsed := time.Since(start)
	if err == nil {
		t.Fatalf("expected drain to error on timeout, got nil")
	}
	if elapsed > 5*time.Second {
		t.Errorf("drain took too long (%v); timeout not respected", elapsed)
	}

	// Node should remain in "draining" (not drained) since an alloc is
	// still running.
	db, _ := store.Open(certpaths.DBPath())
	defer db.Close()
	got, _ := store.NewNodeRepo(db).Get(context.Background(), node.ID)
	if got.State != model.NodeStateDraining {
		t.Errorf("node state = %q, want draining (timeout)", got.State)
	}
}

func TestNodeDrain_NodeNotFound(t *testing.T) {
	_, cleanup := initTestEnv(t)
	defer cleanup()
	resetRootFlags(t)
	drainTestEnv(t)

	var buf bytes.Buffer
	rootCmd.SetOut(&buf)
	rootCmd.SetErr(&buf)
	nodeDrainCmd.SetOut(&buf)
	nodeDrainCmd.SetErr(&buf)
	rootCmd.SetArgs([]string{"node", "drain", "no.such.node", "--timeout", "1s"})
	err := rootCmd.Execute()
	if err == nil {
		t.Fatal("expected error for unknown node, got nil")
	}
	if !strings.Contains(err.Error(), "not found") {
		t.Errorf("error should mention not found, got: %v", err)
	}
}

func TestDaemonDrainAndStop_StopsEachPeer(t *testing.T) {
	_, cleanup := initTestEnv(t)
	defer cleanup()
	resetRootFlags(t)

	drainNodeForTest(t, "peer-a", "peer-a:8443")
	drainNodeForTest(t, "peer-b", "peer-b:8443")
	drainNodeForTest(t, "peer-c", "peer-c:8443")

	mx := &scriptedDrainExec{}
	mx.queueAlways("systemctl stop orca-daemon.service", "", 0)
	drainExecOverride = mx

	var buf bytes.Buffer
	rootCmd.SetOut(&buf)
	rootCmd.SetErr(&buf)
	daemonCmd.SetOut(&buf)
	daemonCmd.SetErr(&buf)
	daemonDrainAndStopCmd.SetOut(&buf)
	daemonDrainAndStopCmd.SetErr(&buf)
	rootCmd.SetArgs([]string{"daemon", "drain-and-stop"})
	if err := rootCmd.Execute(); err != nil {
		t.Fatalf("daemon drain-and-stop: %v", err)
	}
	out := buf.String()
	if !strings.Contains(out, "complete") {
		t.Errorf("expected complete message, got: %s", out)
	}
	// One stop call per peer.
	stops := mx.countCalls("systemctl stop orca-daemon.service")
	if stops != 3 {
		t.Errorf("expected 3 daemon stop calls, got %d (calls: %+v)", stops, mx.calls)
	}
}

func TestDaemonDrainAndStop_AlreadyStoppedIsIdempotent(t *testing.T) {
	_, cleanup := initTestEnv(t)
	defer cleanup()
	resetRootFlags(t)

	drainNodeForTest(t, "migrated", "migrated:8443")
	mx := &scriptedDrainExec{}
	// systemctl stop returns exit 5 ("unit not loaded") → already
	// migrated, idempotent.
	mx.queueAlways("systemctl stop orca-daemon.service", "", 5)
	drainExecOverride = mx

	var buf bytes.Buffer
	rootCmd.SetOut(&buf)
	rootCmd.SetErr(&buf)
	daemonCmd.SetOut(&buf)
	daemonCmd.SetErr(&buf)
	daemonDrainAndStopCmd.SetOut(&buf)
	daemonDrainAndStopCmd.SetErr(&buf)
	rootCmd.SetArgs([]string{"daemon", "drain-and-stop"})
	if err := rootCmd.Execute(); err != nil {
		t.Fatalf("daemon drain-and-stop should be idempotent: %v", err)
	}
	out := buf.String()
	if !strings.Contains(out, "already stopped") && !strings.Contains(out, "already-stopped") {
		t.Errorf("expected already-stopped in output, got: %s", out)
	}
}

func TestJobMigrate_StopsOldAndStartsNewOnTarget(t *testing.T) {
	_, cleanup := initTestEnv(t)
	defer cleanup()
	resetRootFlags(t)

	src := drainNodeForTest(t, "src", "src:8443")
	_ = drainNodeForTest(t, "tgt", "tgt:8443")
	_ = src

	mx := &scriptedDrainExec{}
	// src node lists the job's alloc running.
	mx.queue("list-units",
		"orca-alloc-web-0.service loaded active running\n", 0)
	// tgt node has no allocs (so migrate starts a new one).
	mx.queue("list-units", "", 0)
	// stop on src succeeds.
	mx.queue("systemctl stop orca-alloc-web-0", "", 0)
	// start on tgt succeeds.
	mx.queue("systemctl start", "", 0)
	drainExecOverride = mx

	var buf bytes.Buffer
	rootCmd.SetOut(&buf)
	rootCmd.SetErr(&buf)
	jobCmd.SetOut(&buf)
	jobCmd.SetErr(&buf)
	jobMigrateCmd.SetOut(&buf)
	jobMigrateCmd.SetErr(&buf)
	rootCmd.SetArgs([]string{"job", "migrate", "web", "--to", "tgt"})
	if err := rootCmd.Execute(); err != nil {
		t.Fatalf("job migrate: %v", err)
	}
	out := buf.String()
	if !strings.Contains(out, "Migrated") {
		t.Errorf("expected Migrated message, got: %s", out)
	}
	if mx.countCalls("systemctl stop orca-alloc-web-0") == 0 {
		t.Errorf("expected stop for web-0 on src, calls: %+v", mx.calls)
	}
	if mx.countCalls("systemctl start orca-alloc-web-migrated") == 0 {
		t.Errorf("expected start of new alloc on tgt, calls: %+v", mx.calls)
	}
}

func TestJobMigrate_IdempotentWhenAlreadyOnTarget(t *testing.T) {
	_, cleanup := initTestEnv(t)
	defer cleanup()
	resetRootFlags(t)

	_ = drainNodeForTest(t, "src", "src:8443")
	_ = drainNodeForTest(t, "tgt", "tgt:8443")

	mx := &scriptedDrainExec{}
	// The job is running ONLY on tgt (the target). No allocs on src.
	mx.queue("list-units", "orca-alloc-web-0.service loaded active running\n", 0) // src: empty would be ideal; see scripted note
	// Use scripted: src empty, tgt has web-0.
	mx = &scriptedDrainExec{}
	mx.queue("list-units", "", 0)                                                 // first list-units call (src) → empty
	mx.queue("list-units", "orca-alloc-web-0.service loaded active running\n", 0) // second (tgt) → has web-0
	drainExecOverride = mx

	var buf bytes.Buffer
	rootCmd.SetOut(&buf)
	rootCmd.SetErr(&buf)
	jobCmd.SetOut(&buf)
	jobCmd.SetErr(&buf)
	jobMigrateCmd.SetOut(&buf)
	jobMigrateCmd.SetErr(&buf)
	rootCmd.SetArgs([]string{"job", "migrate", "web", "--to", "tgt"})
	if err := rootCmd.Execute(); err != nil {
		t.Fatalf("job migrate (idempotent): %v", err)
	}
	out := buf.String()
	if !strings.Contains(out, "already running on") {
		t.Errorf("expected already-running idempotent message, got: %s", out)
	}
	// No stop or start commands should have been issued.
	if mx.countCalls("systemctl stop") != 0 || mx.countCalls("systemctl start") != 0 {
		t.Errorf("idempotent migrate must not stop/start, calls: %+v", mx.calls)
	}
}

func TestJobMigrate_RequiresToFlag(t *testing.T) {
	_, cleanup := initTestEnv(t)
	defer cleanup()
	resetRootFlags(t)
	drainTestEnv(t)

	var buf bytes.Buffer
	rootCmd.SetOut(&buf)
	rootCmd.SetErr(&buf)
	jobCmd.SetOut(&buf)
	jobCmd.SetErr(&buf)
	jobMigrateCmd.SetOut(&buf)
	jobMigrateCmd.SetErr(&buf)
	rootCmd.SetArgs([]string{"job", "migrate", "web"})
	err := rootCmd.Execute()
	if err == nil {
		t.Fatal("expected error for missing --to, got nil")
	}
	if !strings.Contains(err.Error(), "--to is required") {
		t.Errorf("error should mention --to required, got: %v", err)
	}
}

func TestSetNodeStateUpdatesDB(t *testing.T) {
	_, cleanup := initTestEnv(t)
	defer cleanup()
	// Sanity test for the store-level SetNodeState used by drain.
	db, err := store.Open(certpaths.DBPath())
	if err != nil {
		t.Fatalf("open db: %v", err)
	}
	defer db.Close()
	repo := store.NewNodeRepo(db)
	n := &model.Node{
		ID:       "node-setstate",
		Name:     "setstate",
		Address:  "setstate:8443",
		State:    model.NodeStateReady,
		JoinedAt: time.Now().UTC(),
		LastSeen: time.Now().UTC(),
	}
	ctx := context.Background()
	if err := repo.Insert(ctx, n); err != nil {
		t.Fatalf("insert: %v", err)
	}
	if err := repo.SetNodeState(ctx, n.ID, string(model.NodeStateDraining)); err != nil {
		t.Fatalf("SetNodeState: %v", err)
	}
	got, err := repo.Get(ctx, n.ID)
	if err != nil {
		t.Fatalf("get: %v", err)
	}
	if got.State != model.NodeStateDraining {
		t.Errorf("state = %q, want draining", got.State)
	}

	// Missing node returns ErrNotFound.
	err = repo.SetNodeState(ctx, "ghost", "draining")
	if !errors.Is(err, store.ErrNotFound) {
		t.Errorf("SetNodeState(ghost) = %v, want ErrNotFound", err)
	}
}

// scriptedDrainExec is a record-and-replay execer that returns queued
// responses in order. queueAlways installs a sticky response that
// matches every subsequent call containing the substring. This is
// more convenient than mockDrainExec when the same command (e.g.
// list-units) must return different outputs across calls.
type scriptedDrainExec struct {
	mu     sync.Mutex
	queued []scriptedResp
	sticky []scriptedResp
	calls  []mockDrainCall
}

type scriptedResp struct {
	match string
	out   string
	exit  int
}

func (s *scriptedDrainExec) queue(match, out string, exit int) {
	s.queued = append(s.queued, scriptedResp{match: match, out: out, exit: exit})
}

func (s *scriptedDrainExec) queueAlways(match, out string, exit int) {
	s.sticky = append(s.sticky, scriptedResp{match: match, out: out, exit: exit})
}

func (s *scriptedDrainExec) Exec(_ context.Context, peer, cmd string) ([]byte, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.calls = append(s.calls, mockDrainCall{peer: peer, cmd: cmd})
	// Sticky responses win if they match.
	for _, r := range s.sticky {
		if r.match == "" || strings.Contains(cmd, r.match) {
			if r.exit != 0 {
				return []byte(r.out), &sshExitErr{code: r.exit}
			}
			return []byte(r.out), nil
		}
	}
	// Queued responses: pop the first matching entry.
	for i, r := range s.queued {
		if r.match == "" || strings.Contains(cmd, r.match) {
			s.queued = append(s.queued[:i], s.queued[i+1:]...)
			if r.exit != 0 {
				return []byte(r.out), &sshExitErr{code: r.exit}
			}
			return []byte(r.out), nil
		}
	}
	return nil, nil
}

func (s *scriptedDrainExec) callsFor(match string) []mockDrainCall {
	s.mu.Lock()
	defer s.mu.Unlock()
	var out []mockDrainCall
	for _, c := range s.calls {
		if strings.Contains(c.cmd, match) {
			out = append(out, c)
		}
	}
	return out
}

func (s *scriptedDrainExec) countCalls(match string) int {
	return len(s.callsFor(match))
}