package scheduler

import (
	"strings"
	"testing"

	"git.cloudinit.dev/coreci/orca/internal/jobspec"
)

// threeLinuxNodes returns a small cluster of three Linux nodes with
// distinct free capacities so best-fit ordering is unambiguous.
func threeLinuxNodes() []NodeInfo {
	return []NodeInfo{
		{Hostname: "node-a", Runtimes: []string{"process"}, Tags: nil, CPU: 4, Memory: 4096, FreeCPU: 4, FreeMem: 4096, Kind: "linux"},
		{Hostname: "node-b", Runtimes: []string{"process"}, Tags: nil, CPU: 8, Memory: 8192, FreeCPU: 8, FreeMem: 8192, Kind: "linux"},
		{Hostname: "node-c", Runtimes: []string{"process"}, Tags: nil, CPU: 2, Memory: 2048, FreeCPU: 2, FreeMem: 2048, Kind: "linux"},
	}
}

func jobSpec(name, oneOf string, constraints []string) *jobspec.WorkloadSpec {
	return &jobspec.WorkloadSpec{
		Kind:        "Job",
		Name:        name,
		Count:       1,
		Runtime:     &jobspec.RuntimeBlock{OneOf: oneOf},
		Constraints: constraints,
	}
}

func serviceSpec(name, oneOf string, count int, constraints []string) *jobspec.WorkloadSpec {
	return &jobspec.WorkloadSpec{
		Kind:        "Service",
		Name:        name,
		Count:       count,
		Runtime:     &jobspec.RuntimeBlock{OneOf: oneOf},
		Constraints: constraints,
	}
}

func daemonSetSpec(name, oneOf string, constraints []string) *jobspec.WorkloadSpec {
	return &jobspec.WorkloadSpec{
		Kind:        "DaemonSet",
		Name:        name,
		Count:       1,
		Runtime:     &jobspec.RuntimeBlock{OneOf: oneOf},
		Constraints: constraints,
	}
}

// ---------------------------------------------------------------------------
// Job
// ---------------------------------------------------------------------------

func TestScheduleJob_BestFit(t *testing.T) {
	nodes := threeLinuxNodes()
	req := WorkloadRequest{Spec: jobSpec("batch", "process", nil), Namespace: "ns"}
	got, err := Schedule(nodes, req)
	if err != nil {
		t.Fatalf("Schedule: %v", err)
	}
	if len(got) != 1 {
		t.Fatalf("placements = %d, want 1", len(got))
	}
	if got[0].Node != "node-b" {
		t.Errorf("Node = %q, want node-b (most free capacity)", got[0].Node)
	}
	if !strings.HasPrefix(got[0].AllocID, "ns/batch-") {
		t.Errorf("AllocID = %q, want ns/batch-*", got[0].AllocID)
	}
	if got[0].Score <= 0 {
		t.Errorf("Score = %d, want > 0", got[0].Score)
	}
}

func TestScheduleJob_NoFittingNode(t *testing.T) {
	nodes := threeLinuxNodes()
	// wasm runtime not advertised by any node.
	req := WorkloadRequest{Spec: jobSpec("wasmjob", "wasm", nil), Namespace: "ns"}
	if _, err := Schedule(nodes, req); err == nil {
		t.Fatal("Schedule: expected error for no-fitting node, got nil")
	}
}

// ---------------------------------------------------------------------------
// Service
// ---------------------------------------------------------------------------

func TestScheduleService_SpreadAcrossNodes(t *testing.T) {
	nodes := threeLinuxNodes()
	req := WorkloadRequest{Spec: serviceSpec("web", "process", 3, nil), Namespace: "ns"}
	got, err := Schedule(nodes, req)
	if err != nil {
		t.Fatalf("Schedule: %v", err)
	}
	if len(got) != 3 {
		t.Fatalf("placements = %d, want 3", len(got))
	}
	seen := map[string]int{}
	for _, p := range got {
		seen[p.Node]++
	}
	if len(seen) != 3 {
		t.Errorf("anti-affinity spread: distinct nodes = %d, want 3; %v", len(seen), seen)
	}
}

func TestScheduleService_ColocationWhenFewerNodes(t *testing.T) {
	nodes := threeLinuxNodes()
	req := WorkloadRequest{Spec: serviceSpec("web", "process", 5, nil), Namespace: "ns"}
	got, err := Schedule(nodes, req)
	if err != nil {
		t.Fatalf("Schedule: %v", err)
	}
	if len(got) != 5 {
		t.Fatalf("placements = %d, want 5", len(got))
	}
	seen := map[string]int{}
	for _, p := range got {
		seen[p.Node]++
	}
	if len(seen) != 3 {
		t.Errorf("colocation: distinct nodes = %d, want 3 (all used)", len(seen))
	}
	// No node should host more than 2 (3 nodes, 5 replicas: 2+2+1).
	for n, c := range seen {
		if c > 2 {
			t.Errorf("node %s has %d replicas, want <= 2", n, c)
		}
	}
}

func TestScheduleService_NoFittingNode(t *testing.T) {
	nodes := threeLinuxNodes()
	req := WorkloadRequest{Spec: serviceSpec("wasm-svc", "wasm", 3, nil), Namespace: "ns"}
	if _, err := Schedule(nodes, req); err == nil {
		t.Fatal("Schedule: expected error for service with no fitting node")
	}
}

// ---------------------------------------------------------------------------
// DaemonSet
// ---------------------------------------------------------------------------

func TestScheduleDaemonSet_AllMatching(t *testing.T) {
	nodes := threeLinuxNodes()
	req := WorkloadRequest{Spec: daemonSetSpec("logrotate", "process", nil), Namespace: "ns"}
	got, err := Schedule(nodes, req)
	if err != nil {
		t.Fatalf("Schedule: %v", err)
	}
	if len(got) != 3 {
		t.Errorf("placements = %d, want 3 (one per node)", len(got))
	}
	seen := map[string]bool{}
	for _, p := range got {
		seen[p.Node] = true
	}
	if len(seen) != 3 {
		t.Errorf("DaemonSet distinct nodes = %d, want 3", len(seen))
	}
}

func TestScheduleDaemonSet_SomeExcludedByConstraint(t *testing.T) {
	nodes := threeLinuxNodes()
	// Only nodes with cpus >= 4 qualify: node-a (4) and node-b (8).
	req := WorkloadRequest{Spec: daemonSetSpec("heavy", "process", []string{"node.cpus >= 4"}), Namespace: "ns"}
	got, err := Schedule(nodes, req)
	if err != nil {
		t.Fatalf("Schedule: %v", err)
	}
	if len(got) != 2 {
		t.Errorf("placements = %d, want 2 (cpus>=4)", len(got))
	}
}

// ---------------------------------------------------------------------------
// Runtime compatibility
// ---------------------------------------------------------------------------

func TestSchedule_RuntimeCompatibilityWasm(t *testing.T) {
	nodes := []NodeInfo{
		{Hostname: "no-wasm", Runtimes: []string{"process"}, Kind: "linux", CPU: 8, Memory: 8192, FreeCPU: 8, FreeMem: 8192},
		{Hostname: "has-wasm", Runtimes: []string{"process", "wasmtime"}, Kind: "linux", CPU: 4, Memory: 4096, FreeCPU: 4, FreeMem: 4096},
	}
	// Even though no-wasm has more free capacity, the wasm workload
	// must land on has-wasm.
	req := WorkloadRequest{Spec: jobSpec("wasmjob", "wasm", nil), Namespace: "ns"}
	got, err := Schedule(nodes, req)
	if err != nil {
		t.Fatalf("Schedule: %v", err)
	}
	if got[0].Node != "has-wasm" {
		t.Errorf("Node = %q, want has-wasm (runtime compatibility)", got[0].Node)
	}
}

func TestSchedule_RuntimeCompatibilityPveVM(t *testing.T) {
	nodes := []NodeInfo{
		{Hostname: "linux-1", Runtimes: []string{"process"}, Kind: "linux", CPU: 8, Memory: 8192, FreeCPU: 8, FreeMem: 8192},
		{Hostname: "pve-1", Runtimes: []string{"process"}, Kind: "proxmox", CPU: 8, Memory: 8192, FreeCPU: 8, FreeMem: 8192},
	}
	req := WorkloadRequest{Spec: jobSpec("vmjob", "pve-vm", nil), Namespace: "ns"}
	got, err := Schedule(nodes, req)
	if err != nil {
		t.Fatalf("Schedule: %v", err)
	}
	if got[0].Node != "pve-1" {
		t.Errorf("Node = %q, want pve-1 (pve-vm requires proxmox kind)", got[0].Node)
	}
}

// ---------------------------------------------------------------------------
// Constraints
// ---------------------------------------------------------------------------

func TestSchedule_ConstraintKindExcludesProxmox(t *testing.T) {
	nodes := []NodeInfo{
		{Hostname: "linux-1", Runtimes: []string{"process"}, Kind: "linux", CPU: 8, Memory: 8192, FreeCPU: 8, FreeMem: 8192},
		{Hostname: "pve-1", Runtimes: []string{"process"}, Kind: "proxmox", CPU: 8, Memory: 8192, FreeCPU: 8, FreeMem: 8192},
	}
	req := WorkloadRequest{Spec: jobSpec("linuxonly", "process", []string{`node.kind == "linux"`}), Namespace: "ns"}
	got, err := Schedule(nodes, req)
	if err != nil {
		t.Fatalf("Schedule: %v", err)
	}
	if got[0].Node != "linux-1" {
		t.Errorf("Node = %q, want linux-1 (kind==linux)", got[0].Node)
	}
}

func TestSchedule_ConstraintCPUsExcludesSmall(t *testing.T) {
	nodes := threeLinuxNodes() // node-c has cpus=2
	req := WorkloadRequest{Spec: jobSpec("big", "process", []string{"node.cpus >= 4"}), Namespace: "ns"}
	got, err := Schedule(nodes, req)
	if err != nil {
		t.Fatalf("Schedule: %v", err)
	}
	if got[0].Node == "node-c" {
		t.Errorf("Node = node-c, want node-a or node-b (cpus>=4)")
	}
}

func TestSchedule_ConstraintNotInTags(t *testing.T) {
	nodes := []NodeInfo{
		{Hostname: "tagged", Runtimes: []string{"process"}, Tags: []string{"log-shipper"}, Kind: "linux", CPU: 8, Memory: 8192, FreeCPU: 8, FreeMem: 8192},
		{Hostname: "clean", Runtimes: []string{"process"}, Tags: nil, Kind: "linux", CPU: 4, Memory: 4096, FreeCPU: 4, FreeMem: 4096},
	}
	req := WorkloadRequest{Spec: jobSpec("worker", "process", []string{`"log-shipper" not in node.tags`}), Namespace: "ns"}
	got, err := Schedule(nodes, req)
	if err != nil {
		t.Fatalf("Schedule: %v", err)
	}
	if got[0].Node != "clean" {
		t.Errorf("Node = %q, want clean (log-shipper not in tags)", got[0].Node)
	}
}

// ---------------------------------------------------------------------------
// Affinity
// ---------------------------------------------------------------------------

func TestSchedule_AffinityPrefersColocatedNode(t *testing.T) {
	// Place a redis service first, then a worker with affinity for
	// redis; the worker should prefer the node where redis already
	// runs even if another node has more free capacity.
	nodes := []NodeInfo{
		{Hostname: "big", Runtimes: []string{"process"}, Kind: "linux", CPU: 8, Memory: 8192, FreeCPU: 8, FreeMem: 8192},
		{Hostname: "small", Runtimes: []string{"process"}, Kind: "linux", CPU: 4, Memory: 4096, FreeCPU: 4, FreeMem: 4096},
	}
	redisReq := WorkloadRequest{Spec: serviceSpec("redis", "process", 1, nil), Namespace: "ns"}
	redisPlacements, err := Schedule(nodes, redisReq)
	if err != nil {
		t.Fatalf("redis Schedule: %v", err)
	}
	// Redis lands on "big" (most free capacity). Now schedule the
	// worker with affinity to redis; it should also land on "big".
	workerReq := WorkloadRequest{Spec: &jobspec.WorkloadSpec{
		Kind:    "Job",
		Name:    "worker",
		Count:   1,
		Runtime: &jobspec.RuntimeBlock{OneOf: "process"},
		Affinity: []jobspec.AffinityRule{
			{Target: "redis", Weight: 1000},
		},
	}, Namespace: "ns"}
	// The affinity is name-based; we need to seed the worker schedule
	// with the redis placement so affinityScore can see it. Schedule
	// does not take prior placements, so test affinityScore directly.
	got := affinityScore(nodes[0], workerReq, redisPlacements)
	if got <= 0 {
		t.Errorf("affinityScore(big) = %d, want > 0 (redis colocated)", got)
	}
	gotSmall := affinityScore(nodes[1], workerReq, redisPlacements)
	if gotSmall != 0 {
		t.Errorf("affinityScore(small) = %d, want 0 (redis not colocated)", gotSmall)
	}
}

func TestSchedule_AffinityCELExpression(t *testing.T) {
	// Affinity with a CEL target: prefer nodes tagged "ssd".
	nodes := []NodeInfo{
		{Hostname: "hdd", Runtimes: []string{"process"}, Tags: []string{"hdd"}, Kind: "linux", CPU: 8, Memory: 8192, FreeCPU: 8, FreeMem: 8192},
		{Hostname: "ssd", Runtimes: []string{"process"}, Tags: []string{"ssd"}, Kind: "linux", CPU: 4, Memory: 4096, FreeCPU: 4, FreeMem: 4096},
	}
	req := WorkloadRequest{Spec: &jobspec.WorkloadSpec{
		Kind:    "Job",
		Name:    "db",
		Count:   1,
		Runtime: &jobspec.RuntimeBlock{OneOf: "process"},
		Affinity: []jobspec.AffinityRule{
			{Target: `"ssd" in node.tags`, Weight: 10000},
		},
	}, Namespace: "ns"}
	got, err := Schedule(nodes, req)
	if err != nil {
		t.Fatalf("Schedule: %v", err)
	}
	if got[0].Node != "ssd" {
		t.Errorf("Node = %q, want ssd (affinity to ssd tag outweighs capacity)", got[0].Node)
	}
}

// ---------------------------------------------------------------------------
// Error paths
// ---------------------------------------------------------------------------

func TestSchedule_EmptyNodes(t *testing.T) {
	req := WorkloadRequest{Spec: jobSpec("x", "process", nil), Namespace: "ns"}
	if _, err := Schedule(nil, req); err == nil {
		t.Fatal("Schedule: expected error for empty nodes, got nil")
	}
}

func TestSchedule_NilSpec(t *testing.T) {
	if _, err := Schedule(threeLinuxNodes(), WorkloadRequest{}); err == nil {
		t.Fatal("Schedule: expected error for nil spec, got nil")
	}
}

func TestSchedule_UnknownKind(t *testing.T) {
	req := WorkloadRequest{Spec: &jobspec.WorkloadSpec{Kind: "Cron", Name: "x", Count: 1}, Namespace: "ns"}
	if _, err := Schedule(threeLinuxNodes(), req); err == nil {
		t.Fatal("Schedule: expected error for unknown kind")
	}
}

// ---------------------------------------------------------------------------
// Score unit tests
// ---------------------------------------------------------------------------

func TestScore_FitsAndDoesNotFit(t *testing.T) {
	node := NodeInfo{Hostname: "n", Runtimes: []string{"process"}, Kind: "linux", CPU: 4, Memory: 4096, FreeCPU: 4, FreeMem: 4096}
	req := WorkloadRequest{Spec: jobSpec("j", "process", nil), Namespace: "ns"}
	score, fits := Score(node, req)
	if !fits {
		t.Error("fits = false, want true")
	}
	if score <= 0 {
		t.Errorf("score = %d, want > 0", score)
	}
}

func TestScore_RuntimeMismatchDoesNotFit(t *testing.T) {
	node := NodeInfo{Hostname: "n", Runtimes: []string{"process"}, Kind: "linux", CPU: 4, Memory: 4096, FreeCPU: 4, FreeMem: 4096}
	req := WorkloadRequest{Spec: jobSpec("j", "wasm", nil), Namespace: "ns"}
	if _, fits := Score(node, req); fits {
		t.Error("fits = true for wasm on process-only node, want false")
	}
}

func TestScore_ConstraintFailsDoesNotFit(t *testing.T) {
	node := NodeInfo{Hostname: "n", Runtimes: []string{"process"}, Kind: "linux", CPU: 4, Memory: 4096, FreeCPU: 4, FreeMem: 4096}
	req := WorkloadRequest{Spec: jobSpec("j", "process", []string{`node.kind == "proxmox"`}), Namespace: "ns"}
	if _, fits := Score(node, req); fits {
		t.Error("fits = true for kind==proxmox on linux node, want false")
	}
}

// ---------------------------------------------------------------------------
// allocID / isAllocFor helpers
// ---------------------------------------------------------------------------

func TestAllocID(t *testing.T) {
	req := WorkloadRequest{Spec: &jobspec.WorkloadSpec{Name: "web"}, Namespace: "prod"}
	if got := allocID(req, 2); got != "prod/web-2" {
		t.Errorf("allocID = %q, want prod/web-2", got)
	}
	req.Namespace = ""
	if got := allocID(req, 0); got != "default/web-0" {
		t.Errorf("allocID = %q, want default/web-0", got)
	}
}

func TestIsAllocFor(t *testing.T) {
	cases := []struct {
		allocID  string
		workload string
		want     bool
	}{
		{"ns/redis-0", "redis", true},
		{"ns/redis-12", "redis", true},
		{"ns/worker-0", "redis", false},
		{"redis-0", "redis", true},
		{"ns/web-canary-3", "web-canary", true},
	}
	for _, c := range cases {
		if got := isAllocFor(c.allocID, c.workload); got != c.want {
			t.Errorf("isAllocFor(%q,%q) = %v, want %v", c.allocID, c.workload, got, c.want)
		}
	}
}

// ---------------------------------------------------------------------------
// normalizeRuntime / hasRuntime
// ---------------------------------------------------------------------------

func TestHasRuntimeAliases(t *testing.T) {
	cases := []struct {
		name string
		node NodeInfo
		want bool
	}{
		{"wasm on wasmtime node", NodeInfo{Runtimes: []string{"wasmtime"}, Kind: "linux"}, true},
		{"wasm on process node", NodeInfo{Runtimes: []string{"process"}, Kind: "linux"}, false},
		{"pve-vm on linux node", NodeInfo{Runtimes: []string{"pve-vm"}, Kind: "linux"}, false},
		{"pve-vm on proxmox node", NodeInfo{Runtimes: nil, Kind: "proxmox"}, true},
		{"process on process node", NodeInfo{Runtimes: []string{"process"}, Kind: "linux"}, true},
		{"empty runtime on any node", NodeInfo{Runtimes: []string{"process"}, Kind: "linux"}, true},
	}
	for _, c := range cases {
		if c.name == "empty runtime on any node" {
			// hasRuntime is only called when OneOf != "".
			continue
		}
		if got := hasRuntime(c.node, "wasm"); c.name == "wasm on wasmtime node" || c.name == "wasm on process node" {
			if got != c.want {
				t.Errorf("%s: hasRuntime(wasm) = %v, want %v", c.name, got, c.want)
			}
		}
	}
	// Explicit pve-vm and process checks.
	if !hasRuntime(NodeInfo{Runtimes: nil, Kind: "proxmox"}, "pve-vm") {
		t.Error("pve-vm on proxmox node should fit")
	}
	if hasRuntime(NodeInfo{Runtimes: nil, Kind: "linux"}, "pve-vm") {
		t.Error("pve-vm on linux node should not fit")
	}
	if !hasRuntime(NodeInfo{Runtimes: []string{"process"}, Kind: "linux"}, "process") {
		t.Error("process on process node should fit")
	}
}