package storage

import (
	"strconv"
	"testing"
	"time"
)

func TestKvBasedExecStatusStg(t *testing.T) {
	kvs := NewMemKvStg()
	stg := NewKvBasedExecStatusStg(kvs)
	if err := stg.OnNewTask("foo", "bar"); err != nil {
		t.Fatal(err)
	}
	done, err := stg.IsTaskFinished("foo")
	if err != nil {
		t.Fatal(err)
	}
	t.Log(done)
}

func TestCleanFinishedTasksRespectsRetainDuration(t *testing.T) {
	kvs := NewMemKvStg()
	stg := NewKvBasedExecStatusStg(kvs)

	// Create and finish two tasks.
	for _, id := range []string{"task-1", "task-2"} {
		if err := stg.OnNewTask(id, "data-"+id); err != nil {
			t.Fatal(err)
		}
		if err := stg.SetTaskFinished(id); err != nil {
			t.Fatal(err)
		}
	}

	// With a long retention, nothing should be cleaned.
	cleaned, err := stg.CleanFinishedTasks(1*time.Hour, 0)
	if err != nil {
		t.Fatal(err)
	}
	if cleaned != 0 {
		t.Fatalf("expected 0 cleaned with long retention, got %d", cleaned)
	}

	// Verify task records still exist.
	for _, id := range []string{"task-1", "task-2"} {
		_, existed, err := kvs.Get(stg.taskKey(id))
		if err != nil {
			t.Fatal(err)
		}
		if !existed {
			t.Fatalf("task %s should still exist after cleanup with long retention", id)
		}
	}

	// With zero retention, all finished tasks should be cleaned.
	cleaned, err = stg.CleanFinishedTasks(0, 0)
	if err != nil {
		t.Fatal(err)
	}
	if cleaned != 2 {
		t.Fatalf("expected 2 cleaned with zero retention, got %d", cleaned)
	}
}

func TestCleanFinishedTasksDoesNotDeleteActive(t *testing.T) {
	kvs := NewMemKvStg()
	stg := NewKvBasedExecStatusStg(kvs)

	if err := stg.OnNewTask("active-1", "data"); err != nil {
		t.Fatal(err)
	}
	// Do NOT finish this task.

	if err := stg.OnNewTask("done-1", "data"); err != nil {
		t.Fatal(err)
	}
	if err := stg.SetTaskFinished("done-1"); err != nil {
		t.Fatal(err)
	}

	// Even with zero retention, active tasks must not be deleted.
	cleaned, err := stg.CleanFinishedTasks(0, 0)
	if err != nil {
		t.Fatal(err)
	}
	if cleaned != 1 {
		t.Fatalf("expected 1 cleaned (only finished task), got %d", cleaned)
	}

	// Active task record should still exist.
	_, existed, err := kvs.Get(stg.taskKey("active-1"))
	if err != nil {
		t.Fatal(err)
	}
	if !existed {
		t.Fatal("active task should not be deleted")
	}
}

func TestSetTaskFinishedRecordsTimestamp(t *testing.T) {
	kvs := NewMemKvStg()
	stg := NewKvBasedExecStatusStg(kvs)

	if err := stg.OnNewTask("ts-1", "data"); err != nil {
		t.Fatal(err)
	}

	before := time.Now()
	if err := stg.SetTaskFinished("ts-1"); err != nil {
		t.Fatal(err)
	}
	after := time.Now()

	val, existed, err := kvs.Get(stg.finishedAtKey("ts-1"))
	if err != nil {
		t.Fatal(err)
	}
	if !existed {
		t.Fatal("expected finished_at key to exist after SetTaskFinished")
	}
	if val == "" {
		t.Fatal("expected non-empty finished_at value")
	}

	// Parse the stored timestamp and verify it is within the expected range.
	finishedAtNano, parseErr := strconv.ParseInt(val, 10, 64)
	if parseErr != nil {
		t.Fatalf("failed to parse finished_at value %q: %v", val, parseErr)
	}
	finishedAt := time.Unix(0, finishedAtNano)
	if finishedAt.Before(before) || finishedAt.After(after) {
		t.Fatalf("finished_at %v is not between before %v and after %v", finishedAt, before, after)
	}

	// Verify the task is also marked finished.
	done, err := stg.IsTaskFinished("ts-1")
	if err != nil {
		t.Fatal(err)
	}
	if !done {
		t.Fatal("expected task to be finished")
	}
}

func TestCleanFinishedTasksBackfillsLegacyTask(t *testing.T) {
	kvs := NewMemKvStg()
	stg := NewKvBasedExecStatusStg(kvs)

	// Simulate a legacy finished task: create it, then manually remove
	// it from the active list WITHOUT calling SetTaskFinished (which would
	// write finished_at). This mimics the pre-annotation state.
	if err := stg.OnNewTask("legacy-1", "data"); err != nil {
		t.Fatal(err)
	}
	// Remove from active list to mark as "finished" without writing finished_at.
	if err := kvs.Delete(stg.activeTaskKey("legacy-1")); err != nil {
		t.Fatal(err)
	}

	// Verify there is no finished_at marker.
	// MemKvStg.Get auto-creates nodes, so check value is empty instead.
	val, _, _ := kvs.Get(stg.finishedAtKey("legacy-1"))
	if val != "" {
		t.Fatalf("expected no finished_at marker for legacy task, got %q", val)
	}

	// First cleanup cycle with retention > 0: should backfill finished_at
	// and NOT delete the task.
	cleaned, err := stg.CleanFinishedTasks(1*time.Hour, 0)
	if err != nil {
		t.Fatal(err)
	}
	if cleaned != 0 {
		t.Fatalf("expected 0 cleaned on first cycle (backfill only), got %d", cleaned)
	}

	// Verify the finished_at marker was backfilled.
	val, existed, err := kvs.Get(stg.finishedAtKey("legacy-1"))
	if err != nil {
		t.Fatal(err)
	}
	if !existed || val == "" {
		t.Fatal("expected finished_at to be backfilled for legacy task")
	}
	finishedAtNano, parseErr := strconv.ParseInt(val, 10, 64)
	if parseErr != nil {
		t.Fatalf("failed to parse backfilled finished_at: %v", parseErr)
	}
	finishedAt := time.Unix(0, finishedAtNano)

	// The backfilled time should be recent (within the last minute).
	if time.Since(finishedAt) > 1*time.Minute {
		t.Fatalf("backfilled finished_at %v is too old", finishedAt)
	}

	// Task record should still exist.
	taskVal, _, err := kvs.Get(stg.taskKey("legacy-1"))
	if err != nil {
		t.Fatal(err)
	}
	if taskVal == "" {
		t.Fatal("task should still exist after backfill cycle")
	}

	// Second cleanup cycle with zero retention: now the task should be deleted.
	cleaned, err = stg.CleanFinishedTasks(0, 0)
	if err != nil {
		t.Fatal(err)
	}
	if cleaned != 1 {
		t.Fatalf("expected 1 cleaned on second cycle, got %d", cleaned)
	}
}

func TestCleanFinishedTasksEnforcesMaxCount(t *testing.T) {
	kvs := NewMemKvStg()
	stg := NewKvBasedExecStatusStg(kvs)

	// Create and finish 5 tasks with staggered finish times.
	ids := []string{"t-1", "t-2", "t-3", "t-4", "t-5"}
	for i, id := range ids {
		if err := stg.OnNewTask(id, "data-"+id); err != nil {
			t.Fatal(err)
		}
		// Manually set finished_at with increasing timestamps so we can control ordering.
		ts := strconv.FormatInt(time.Now().Add(time.Duration(i)*time.Minute).UnixNano(), 10)
		if err := kvs.Set(stg.finishedAtKey(id), ts); err != nil {
			t.Fatal(err)
		}
		if err := kvs.Delete(stg.activeTaskKey(id)); err != nil {
			t.Fatal(err)
		}
	}

	// Allow max 3 finished tasks → should delete 2 oldest (t-1, t-2).
	cleaned, err := stg.CleanFinishedTasks(1*time.Hour, 3)
	if err != nil {
		t.Fatal(err)
	}
	if cleaned != 2 {
		t.Fatalf("expected 2 excess tasks cleaned, got %d", cleaned)
	}

	// Verify t-1 and t-2 are deleted.
	for _, id := range []string{"t-1", "t-2"} {
		val, _, _ := kvs.Get(stg.taskKey(id))
		if val != "" {
			t.Fatalf("task %s should have been deleted", id)
		}
	}
	// Verify t-3, t-4, t-5 still exist.
	for _, id := range []string{"t-3", "t-4", "t-5"} {
		val, _, _ := kvs.Get(stg.taskKey(id))
		if val == "" {
			t.Fatalf("task %s should still exist", id)
		}
	}
}

func TestCleanFinishedTasksMaxCountNoopWhenUnderLimit(t *testing.T) {
	kvs := NewMemKvStg()
	stg := NewKvBasedExecStatusStg(kvs)

	// Create and finish 2 tasks.
	for _, id := range []string{"a-1", "a-2"} {
		if err := stg.OnNewTask(id, "data"); err != nil {
			t.Fatal(err)
		}
		if err := stg.SetTaskFinished(id); err != nil {
			t.Fatal(err)
		}
	}

	// Max 5 → nothing to clean.
	cleaned, err := stg.CleanFinishedTasks(1*time.Hour, 5)
	if err != nil {
		t.Fatal(err)
	}
	if cleaned != 0 {
		t.Fatalf("expected 0 cleaned when under limit, got %d", cleaned)
	}
}

func TestCleanFinishedTasksMaxCountSkipsActiveTasks(t *testing.T) {
	kvs := NewMemKvStg()
	stg := NewKvBasedExecStatusStg(kvs)

	// Create 3 tasks: 2 finished, 1 active.
	for _, id := range []string{"f-1", "f-2", "active-1"} {
		if err := stg.OnNewTask(id, "data"); err != nil {
			t.Fatal(err)
		}
	}
	if err := stg.SetTaskFinished("f-1"); err != nil {
		t.Fatal(err)
	}
	if err := stg.SetTaskFinished("f-2"); err != nil {
		t.Fatal(err)
	}
	// active-1 remains active.

	// Max 1 finished task → should delete 1 oldest finished, not touch active.
	cleaned, err := stg.CleanFinishedTasks(1*time.Hour, 1)
	if err != nil {
		t.Fatal(err)
	}
	if cleaned != 1 {
		t.Fatalf("expected 1 cleaned, got %d", cleaned)
	}

	// Active task should still be present in active list.
	done, err := stg.IsTaskFinished("active-1")
	if err != nil {
		t.Fatal(err)
	}
	if done {
		t.Fatal("active-1 should not be finished")
	}
}

func TestCleanFinishedTasksLegacyParticipatesInCountLimit(t *testing.T) {
	kvs := NewMemKvStg()
	stg := NewKvBasedExecStatusStg(kvs)

	// Create 3 legacy finished tasks (no finished_at marker).
	for _, id := range []string{"leg-1", "leg-2", "leg-3"} {
		if err := stg.OnNewTask(id, "data"); err != nil {
			t.Fatal(err)
		}
		if err := kvs.Delete(stg.activeTaskKey(id)); err != nil {
			t.Fatal(err)
		}
	}

	// With retainDuration=1h and maxFinishedTasks=2:
	// All 3 are backfilled with time.Now() (which is after cutoff = now-1h),
	// so all 3 survive time-based cleanup. Then count-based cleanup should
	// delete 1 (3-2=1 excess).
	cleaned, err := stg.CleanFinishedTasks(1*time.Hour, 2)
	if err != nil {
		t.Fatal(err)
	}
	if cleaned != 1 {
		t.Fatalf("expected 1 excess legacy task cleaned, got %d", cleaned)
	}
}

func TestCleanFinishedTasksLegacyDeletedImmediatelyWhenZeroRetain(t *testing.T) {
	kvs := NewMemKvStg()
	stg := NewKvBasedExecStatusStg(kvs)

	// Create a legacy finished task (no finished_at marker).
	if err := stg.OnNewTask("leg-imm", "data"); err != nil {
		t.Fatal(err)
	}
	if err := kvs.Delete(stg.activeTaskKey("leg-imm")); err != nil {
		t.Fatal(err)
	}

	// With retainDuration=0, backfilledAt = cutoff = now, so the task should
	// be deleted in the same cycle (not deferred to next cycle).
	cleaned, err := stg.CleanFinishedTasks(0, 0)
	if err != nil {
		t.Fatal(err)
	}
	if cleaned != 1 {
		t.Fatalf("expected 1 cleaned immediately with zero retention, got %d", cleaned)
	}

	// Task record should be gone.
	val, _, _ := kvs.Get(stg.taskKey("leg-imm"))
	if val != "" {
		t.Fatal("task should have been deleted")
	}
}