package k8score

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"

	corev1 "k8s.io/api/core/v1"
	policyv1 "k8s.io/api/policy/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/kubernetes/fake"
	k8stesting "k8s.io/client-go/testing"
)

func drainTestPod(name string, mutate ...func(*corev1.Pod)) corev1.Pod {
	yes := true
	pod := corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: "shop",
			Labels:    map[string]string{"app": "web"},
			OwnerReferences: []metav1.OwnerReference{{
				Kind: "ReplicaSet", Name: "web-abc", Controller: &yes,
			}},
		},
		Spec:   corev1.PodSpec{NodeName: "worker-1"},
		Status: corev1.PodStatus{Phase: corev1.PodRunning},
	}
	for _, m := range mutate {
		m(&pod)
	}
	return pod
}

func withOwner(kind string, controller bool) func(*corev1.Pod) {
	return func(p *corev1.Pod) {
		c := controller
		p.OwnerReferences = []metav1.OwnerReference{{Kind: kind, Name: "owner", Controller: &c}}
	}
}

func withoutOwner(p *corev1.Pod) { p.OwnerReferences = nil }
func withEmptyDir(p *corev1.Pod) {
	p.Spec.Volumes = []corev1.Volume{{Name: "tmp", VolumeSource: corev1.VolumeSource{EmptyDir: &corev1.EmptyDirVolumeSource{}}}}
}
func withMirror(p *corev1.Pod) { p.Annotations = map[string]string{corev1.MirrorPodAnnotationKey: "x"} }
func withDeletionTimestamp(p *corev1.Pod) {
	now := metav1.Now()
	p.DeletionTimestamp = &now
}
func withPhase(ph corev1.PodPhase) func(*corev1.Pod) {
	return func(p *corev1.Pod) { p.Status.Phase = ph }
}

func withReady(ready bool) func(*corev1.Pod) {
	return func(p *corev1.Pod) {
		st := corev1.ConditionFalse
		if ready {
			st = corev1.ConditionTrue
		}
		p.Status.Conditions = []corev1.PodCondition{{Type: corev1.PodReady, Status: st}}
	}
}

func pdb(ns, name string, selector *metav1.LabelSelector, allowed int32) policyv1.PodDisruptionBudget {
	return policyv1.PodDisruptionBudget{
		ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: ns, Generation: 1},
		Spec:       policyv1.PodDisruptionBudgetSpec{Selector: selector},
		Status:     policyv1.PodDisruptionBudgetStatus{DisruptionsAllowed: allowed, ObservedGeneration: 1},
	}
}

func withPolicy(policy policyv1.UnhealthyPodEvictionPolicyType) func(*policyv1.PodDisruptionBudget) {
	return func(p *policyv1.PodDisruptionBudget) { p.Spec.UnhealthyPodEvictionPolicy = &policy }
}

func withHealth(current, desired int32) func(*policyv1.PodDisruptionBudget) {
	return func(p *policyv1.PodDisruptionBudget) {
		p.Status.CurrentHealthy, p.Status.DesiredHealthy = current, desired
	}
}

func stale(p *policyv1.PodDisruptionBudget) { p.Generation = 2 }

func pdbWith(base policyv1.PodDisruptionBudget, mutate ...func(*policyv1.PodDisruptionBudget)) policyv1.PodDisruptionBudget {
	for _, m := range mutate {
		m(&base)
	}
	return base
}

func TestClassifyPodForDrain(t *testing.T) {
	webSelector := &metav1.LabelSelector{MatchLabels: map[string]string{"app": "web"}}
	defaults := DrainOptions{IgnoreDaemonSets: true}

	tests := []struct {
		name string
		pod  corev1.Pod
		opts DrainOptions
		pdbs []policyv1.PodDisruptionBudget
		want DrainOutcome
	}{
		{name: "managed running pod is evicted", pod: drainTestPod("a"), opts: defaults, want: DrainOutcomeEvict},
		{name: "succeeded pod is skipped", pod: drainTestPod("a", withPhase(corev1.PodSucceeded)), opts: defaults, want: DrainOutcomeSkip},
		{name: "failed pod is skipped", pod: drainTestPod("a", withPhase(corev1.PodFailed)), opts: defaults, want: DrainOutcomeSkip},
		{name: "mirror pod is skipped even with force", pod: drainTestPod("a", withMirror), opts: DrainOptions{IgnoreDaemonSets: true, Force: true}, want: DrainOutcomeSkip},
		{name: "DaemonSet pod is skipped when ignored", pod: drainTestPod("a", withOwner("DaemonSet", true)), opts: defaults, want: DrainOutcomeSkip},
		{name: "DaemonSet pod is never evicted, even when not ignored", pod: drainTestPod("a", withOwner("DaemonSet", true)), opts: DrainOptions{}, want: DrainOutcomeSkip},
		{name: "terminating pod is evicted without consulting PDBs", pod: drainTestPod("a", withDeletionTimestamp), opts: defaults, pdbs: []policyv1.PodDisruptionBudget{pdb("shop", "web", webSelector, 0)}, want: DrainOutcomeEvict},
		{name: "terminating DaemonSet pod is still skipped", pod: drainTestPod("a", withOwner("DaemonSet", true), withDeletionTimestamp), opts: defaults, want: DrainOutcomeSkip},
		{name: "non-controller DaemonSet reference does not make a DaemonSet pod", pod: drainTestPod("a", withOwner("DaemonSet", false)), opts: defaults, want: DrainOutcomeSkip}, // unmanaged, no force
		{name: "unmanaged pod is skipped without force", pod: drainTestPod("a", withoutOwner), opts: defaults, want: DrainOutcomeSkip},
		{name: "unmanaged pod is evicted with force", pod: drainTestPod("a", withoutOwner), opts: DrainOptions{IgnoreDaemonSets: true, Force: true}, want: DrainOutcomeEvict},
		{name: "owner reference without controller flag counts as unmanaged", pod: drainTestPod("a", withOwner("ReplicaSet", false)), opts: defaults, want: DrainOutcomeSkip},
		{name: "emptyDir pod is skipped unless allowed", pod: drainTestPod("a", withEmptyDir), opts: defaults, want: DrainOutcomeSkip},
		{name: "emptyDir pod is evicted when allowed", pod: drainTestPod("a", withEmptyDir), opts: DrainOptions{IgnoreDaemonSets: true, DeleteEmptyDirData: true}, want: DrainOutcomeEvict},
		{name: "exhausted PDB may block", pod: drainTestPod("a"), opts: defaults, pdbs: []policyv1.PodDisruptionBudget{pdb("shop", "web", webSelector, 0)}, want: DrainOutcomeMayBlock},
		{name: "PDB with budget left does not block", pod: drainTestPod("a"), opts: defaults, pdbs: []policyv1.PodDisruptionBudget{pdb("shop", "web", webSelector, 1)}, want: DrainOutcomeEvict},
		{name: "PDB in another namespace is ignored", pod: drainTestPod("a"), opts: defaults, pdbs: []policyv1.PodDisruptionBudget{pdb("other", "web", webSelector, 0)}, want: DrainOutcomeEvict},
		{name: "PDB with non-matching selector is ignored", pod: drainTestPod("a"), opts: defaults, pdbs: []policyv1.PodDisruptionBudget{pdb("shop", "db", &metav1.LabelSelector{MatchLabels: map[string]string{"app": "db"}}, 0)}, want: DrainOutcomeEvict},
		{name: "PDB with nil selector matches nothing", pod: drainTestPod("a"), opts: defaults, pdbs: []policyv1.PodDisruptionBudget{pdb("shop", "none", nil, 0)}, want: DrainOutcomeEvict},
		{name: "PDB with empty selector matches every pod in the namespace", pod: drainTestPod("a"), opts: defaults, pdbs: []policyv1.PodDisruptionBudget{pdb("shop", "all", &metav1.LabelSelector{}, 0)}, want: DrainOutcomeMayBlock},
		{name: "skip decisions win over PDB", pod: drainTestPod("a", withEmptyDir), opts: defaults, pdbs: []policyv1.PodDisruptionBudget{pdb("shop", "web", webSelector, 0)}, want: DrainOutcomeSkip},
		{name: "pending pod bypasses PDBs", pod: drainTestPod("a", withPhase(corev1.PodPending)), opts: defaults, pdbs: []policyv1.PodDisruptionBudget{pdb("shop", "web", webSelector, 0)}, want: DrainOutcomeEvict},
		{name: "two PDBs covering the pod block the eviction even with budget left", pod: drainTestPod("a"), opts: defaults, pdbs: []policyv1.PodDisruptionBudget{pdb("shop", "web", webSelector, 1), pdb("shop", "all", &metav1.LabelSelector{}, 1)}, want: DrainOutcomeMayBlock},
		{name: "unreconciled PDB status blocks", pod: drainTestPod("a"), opts: defaults, pdbs: []policyv1.PodDisruptionBudget{pdbWith(pdb("shop", "web", webSelector, 1), stale)}, want: DrainOutcomeMayBlock},
		{name: "unready pod under AlwaysAllow is evicted despite an exhausted budget", pod: drainTestPod("a", withReady(false)), opts: defaults, pdbs: []policyv1.PodDisruptionBudget{pdbWith(pdb("shop", "web", webSelector, 0), withPolicy(policyv1.AlwaysAllow))}, want: DrainOutcomeEvict},
		{name: "unready pod under IfHealthyBudget with a healthy budget is evicted", pod: drainTestPod("a", withReady(false)), opts: defaults, pdbs: []policyv1.PodDisruptionBudget{pdbWith(pdb("shop", "web", webSelector, 0), withHealth(2, 2))}, want: DrainOutcomeEvict},
		{name: "unready pod under IfHealthyBudget with an unhealthy budget may block", pod: drainTestPod("a", withReady(false)), opts: defaults, pdbs: []policyv1.PodDisruptionBudget{pdbWith(pdb("shop", "web", webSelector, 0), withHealth(1, 2))}, want: DrainOutcomeMayBlock},
		{name: "ready pod under AlwaysAllow with an exhausted budget still may block", pod: drainTestPod("a", withReady(true)), opts: defaults, pdbs: []policyv1.PodDisruptionBudget{pdbWith(pdb("shop", "web", webSelector, 0), withPolicy(policyv1.AlwaysAllow))}, want: DrainOutcomeMayBlock},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := ClassifyPodForDrain(tt.pod, tt.opts, tt.pdbs)
			if got.Outcome != tt.want {
				t.Fatalf("outcome = %q (%s), want %q", got.Outcome, got.Reason, tt.want)
			}
			if got.Reason == "" {
				t.Fatalf("every decision must carry an operator-facing reason")
			}
			if got.Namespace != tt.pod.Namespace || got.Name != tt.pod.Name {
				t.Fatalf("decision must identify the pod: %+v", got)
			}
			if tt.want == DrainOutcomeMayBlock && got.PDB == "" {
				t.Fatalf("may-block decision must name the PDB: %+v", got)
			}
			if got.EmptyDir != hasLocalStorage(tt.pod) {
				t.Fatalf("emptyDir flag must reflect the pod's volumes regardless of outcome: %+v", got)
			}
			if got.Terminating != (tt.pod.DeletionTimestamp != nil) {
				t.Fatalf("terminating flag must reflect the pod's deletion timestamp: %+v", got)
			}
		})
	}
}

func TestPlanNodeDrainIsReadOnlyAndCountsOutcomes(t *testing.T) {
	client := fake.NewSimpleClientset(
		&corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "worker-1"}},
		func() *corev1.Pod { p := drainTestPod("web-1"); return &p }(),
		func() *corev1.Pod { p := drainTestPod("web-2", withEmptyDir); return &p }(),
		func() *corev1.Pod { p := drainTestPod("agent", withOwner("DaemonSet", true)); return &p }(),
		func() *corev1.Pod { p := drainTestPod("elsewhere"); p.Spec.NodeName = "worker-2"; return &p }(),
		&policyv1.PodDisruptionBudget{
			ObjectMeta: metav1.ObjectMeta{Name: "web", Namespace: "shop"},
			Spec:       policyv1.PodDisruptionBudgetSpec{Selector: &metav1.LabelSelector{MatchLabels: map[string]string{"app": "web"}}},
			Status:     policyv1.PodDisruptionBudgetStatus{DisruptionsAllowed: 0},
		},
	)
	// the fake clientset does not implement field selectors; filter like the API server would
	client.PrependReactor("list", "pods", func(action k8stesting.Action) (bool, runtime.Object, error) {
		list := &corev1.PodList{}
		all, _ := client.Tracker().List(schema.GroupVersionResource{Version: "v1", Resource: "pods"}, schema.GroupVersionKind{Version: "v1", Kind: "Pod"}, "")
		for _, p := range all.(*corev1.PodList).Items {
			if p.Spec.NodeName == "worker-1" {
				list.Items = append(list.Items, p)
			}
		}
		return true, list, nil
	})

	plan, err := PlanNodeDrain(context.Background(), client, "worker-1", DrainOptions{IgnoreDaemonSets: true})
	if err != nil {
		t.Fatalf("PlanNodeDrain: %v", err)
	}
	for _, a := range client.Actions() {
		switch a.GetVerb() {
		case "get", "list":
		default:
			t.Fatalf("plan must not mutate the cluster, saw %s %s", a.GetVerb(), a.GetResource().Resource)
		}
	}
	if !plan.Estimate || plan.Node != "worker-1" || plan.GeneratedAt.IsZero() {
		t.Fatalf("plan header is incomplete: %+v", plan)
	}
	if !plan.PDBsEvaluated {
		t.Fatalf("PDBs were listable and must be reported as evaluated")
	}
	for _, d := range plan.Pods {
		if (d.Outcome != DrainOutcomeSkip) != d.PDBChecked {
			t.Fatalf("pdbChecked must be true exactly for decisions that reached the budget check: %+v", d)
		}
	}
	if got := plan.Summary; got.Evict != 0 || got.Skip != 2 || got.MayBlock != 1 {
		t.Fatalf("summary = %+v, want evict=0 skip=2 mayBlock=1", got)
	}
	if len(plan.Pods) != 3 {
		t.Fatalf("plan must list only pods on the node, got %d", len(plan.Pods))
	}
}

func TestPlanNodeDrainReportsUnevaluatedPDBs(t *testing.T) {
	client := fake.NewSimpleClientset(
		&corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "worker-1"}},
		func() *corev1.Pod { p := drainTestPod("web-1"); return &p }(),
	)
	client.PrependReactor("list", "poddisruptionbudgets", func(k8stesting.Action) (bool, runtime.Object, error) {
		return true, nil, apierrors.NewForbidden(schema.GroupResource{Group: "policy", Resource: "poddisruptionbudgets"}, "", nil)
	})

	plan, err := PlanNodeDrain(context.Background(), client, "worker-1", DrainOptions{IgnoreDaemonSets: true})
	if err != nil {
		t.Fatalf("a forbidden PDB list must not fail the plan: %v", err)
	}
	if plan.PDBsEvaluated || plan.PDBError == "" {
		t.Fatalf("plan must say PDBs were not evaluated instead of reporting zero: %+v", plan)
	}
	if plan.Summary.MayBlock != 0 || plan.Pods[0].Outcome != DrainOutcomeEvict || plan.Pods[0].PDBChecked {
		t.Fatalf("pod must be classified without PDB knowledge and flagged: %+v", plan.Pods[0])
	}
}

func TestPlanNodeDrainUnknownNode(t *testing.T) {
	client := fake.NewSimpleClientset()
	_, err := PlanNodeDrain(context.Background(), client, "ghost", DrainOptions{})
	if !apierrors.IsNotFound(err) {
		t.Fatalf("want NotFound for an unknown node, got %v", err)
	}
}

func TestDrainNodeReportsSkippedPodsWithoutPDBClaims(t *testing.T) {
	client := fake.NewSimpleClientset(
		&corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "worker-1"}},
		func() *corev1.Pod { p := drainTestPod("bare", withoutOwner); return &p }(),
	)
	res, err := DrainNode(context.Background(), client, "worker-1", DrainOptions{IgnoreDaemonSets: true, Timeout: time.Second})
	if err != nil {
		t.Fatalf("DrainNode: %v", err)
	}
	if len(res.SkippedPods) != 1 || res.SkippedPods[0].Outcome != DrainOutcomeSkip || res.SkippedPods[0].PDBChecked {
		t.Fatalf("skipped pods must be reported with pdbChecked=false (PDBs are not consulted during execution): %+v", res.SkippedPods)
	}
}

func TestPlanNodeDrainListsPDBsOnlyWhereADecisionNeedsThem(t *testing.T) {
	client := fake.NewSimpleClientset(
		&corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "worker-1"}},
		func() *corev1.Pod {
			p := drainTestPod("agent", withOwner("DaemonSet", true))
			p.Namespace = "infra"
			return &p
		}(),
		func() *corev1.Pod { p := drainTestPod("web-1"); return &p }(),
	)
	client.PrependReactor("list", "poddisruptionbudgets", func(action k8stesting.Action) (bool, runtime.Object, error) {
		if action.GetNamespace() == "infra" {
			t.Fatalf("PDBs must not be listed for a namespace whose pods are all skipped")
		}
		return false, nil, nil
	})
	plan, err := PlanNodeDrain(context.Background(), client, "worker-1", DrainOptions{IgnoreDaemonSets: true})
	if err != nil {
		t.Fatalf("PlanNodeDrain: %v", err)
	}
	if !plan.PDBsEvaluated {
		t.Fatalf("plan must count as evaluated when every needed namespace was listed: %+v", plan)
	}
	for _, d := range plan.Pods {
		wantChecked := d.Name == "web-1"
		if d.PDBChecked != wantChecked {
			t.Fatalf("pdbChecked for %s = %v, want %v", d.Name, d.PDBChecked, wantChecked)
		}
	}
}

// Six PDB-blocked pods sleep in their 429 retry loops; with only five concurrency
// slots, a pod with no budget at all must still get its eviction attempted promptly.
// The free pod's reactor is gated on five distinct blocked pods having been attempted
// first, so this test cannot pass by racing ahead of the blocked retry loops.
func TestDrainNodeEvictsFreePodWhilePDBBlockedPodsRetry(t *testing.T) {
	objs := []runtime.Object{&corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: "worker-1"}}}
	for i := 0; i < 6; i++ {
		p := drainTestPod(fmt.Sprintf("blocked-%d", i))
		objs = append(objs, &p)
	}
	free := drainTestPod("free")
	objs = append(objs, &free)
	client := fake.NewSimpleClientset(objs...)

	// Reactors run under the fake clientset's lock, so the gate must never block:
	// the free pod is refused with the same 429 until five distinct blocked pods
	// have been attempted, then its next retry is allowed through.
	var mu sync.Mutex
	blockedSeen := map[string]bool{}
	gateOpen := false
	var freeEvictedAfter time.Duration
	start := time.Now()
	pdbRefused := apierrors.NewTooManyRequests("Cannot evict pod as it would violate the pod's disruption budget.", 0)

	client.PrependReactor("create", "pods", func(action k8stesting.Action) (bool, runtime.Object, error) {
		if action.GetSubresource() != "eviction" {
			return false, nil, nil
		}
		ev, ok := action.(k8stesting.CreateAction).GetObject().(*policyv1.Eviction)
		if !ok {
			t.Errorf("eviction create carried unexpected object")
			return true, nil, nil
		}
		mu.Lock()
		defer mu.Unlock()
		if ev.Name == "free" {
			if !gateOpen {
				return true, nil, pdbRefused
			}
			if freeEvictedAfter == 0 {
				freeEvictedAfter = time.Since(start)
			}
			return true, nil, nil
		}
		blockedSeen[ev.Name] = true
		if len(blockedSeen) >= 5 {
			gateOpen = true
		}
		return true, nil, pdbRefused
	})

	const timeout = 3 * time.Second
	res, err := DrainNode(context.Background(), client, "worker-1", DrainOptions{IgnoreDaemonSets: true, Timeout: timeout})
	if err != nil {
		t.Fatalf("DrainNode: %v", err)
	}

	if len(res.EvictedPods) != 1 || res.EvictedPods[0] != "shop/free" {
		t.Fatalf("free pod must be evicted, got evicted=%v errors=%v", res.EvictedPods, res.Errors)
	}
	// Half the deadline separates "evicted while the blocked pods retried" from
	// "evicted only after the deadline freed their slots".
	if freeEvictedAfter <= 0 || freeEvictedAfter > timeout/2 {
		t.Fatalf("free pod was evicted after %v; slots are being held across PDB backoff sleeps", freeEvictedAfter)
	}
	if len(res.Errors) != 6 {
		t.Fatalf("all six blocked pods must report errors, got %v", res.Errors)
	}
	for _, e := range res.Errors {
		if !strings.Contains(e, "timed out waiting for PDB to allow eviction") {
			t.Fatalf("blocked pod error must name the PDB timeout, got %q", e)
		}
	}
}