package server

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/observability/vm-migration/internal/migration"
	"github.com/tidbcloud/observability/vm-migration/internal/schema"
)

func TestRegistryGetList(t *testing.T) {
	migrationsDir := filepath.Join(t.TempDir(), "migrations")
	writePreparedMigration(t, migrationsDir, "a", []string{"t1"})
	writePreparedMigration(t, migrationsDir, "b", []string{"t1"})

	r := NewRegistry(context.Background(), migrationsDir)
	r.LoadFromDisk()

	_, ok := r.Get("a")
	assert.True(t, ok, "Get(a) not found")
	_, ok = r.Get("missing")
	assert.False(t, ok, "Get(missing) should return false")
	assert.Len(t, r.List(), 2, "List: got %d want 2", len(r.List()))
}

func writePreparedMigration(t *testing.T, migrationsDir, name string, tenants []string) {
	t.Helper()
	dir := filepath.Join(migrationsDir, name)
	assert.NoError(t, os.MkdirAll(dir, 0o755))
	plan := `{"tenants":[`
	for i, tn := range tenants {
		if i > 0 {
			plan += ","
		}
		plan += `{"name":"` + tn + `","tenantsFile":"` + tn + `.json"}`
	}
	plan += `]}`
	assert.NoError(t, os.WriteFile(filepath.Join(dir, "migration-loop.json"), []byte(plan), 0o644))
}

func TestRegistryLoadFromDisk(t *testing.T) {
	t.Run("EmptyOrMissingDir", func(t *testing.T) {
		cases := map[string]func(t *testing.T, migrationsDir string){
			"missing": func(t *testing.T, migrationsDir string) {},
			"empty": func(t *testing.T, migrationsDir string) {
				assert.NoError(t, os.MkdirAll(migrationsDir, 0o755))
			},
		}
		for name, setup := range cases {
			t.Run(name, func(t *testing.T) {
				migrationsDir := filepath.Join(t.TempDir(), "migrations")
				setup(t, migrationsDir)
				r := NewRegistry(context.Background(), migrationsDir)
				r.LoadFromDisk()
				assert.Empty(t, r.List(), "List: got %d, want 0", len(r.List()))
			})
		}
	})

	t.Run("LoadsPrepared", func(t *testing.T) {
		migrationsDir := filepath.Join(t.TempDir(), "migrations")
		writePreparedMigration(t, migrationsDir, "prod-2026-08", []string{"1001:0"})

		r := NewRegistry(context.Background(), migrationsDir)
		r.LoadFromDisk()

		m, ok := r.Get("prod-2026-08")
		assert.True(t, ok, "Get(prod-2026-08) not found")
		assert.Equal(t, "prod-2026-08", m.Name())
	})

	t.Run("SkipsIncomplete", func(t *testing.T) {
		migrationsDir := filepath.Join(t.TempDir(), "migrations")
		writePreparedMigration(t, migrationsDir, "prepared", []string{"1001:0"})
		staleDir := filepath.Join(migrationsDir, "stale")
		assert.NoError(t, os.MkdirAll(staleDir, 0o755))

		r := NewRegistry(context.Background(), migrationsDir)
		r.LoadFromDisk()

		_, ok := r.Get("prepared")
		assert.True(t, ok, "prepared should be loaded")
		_, ok = r.Get("stale")
		assert.False(t, ok, "stale should be skipped")
	})

	t.Run("Multiple", func(t *testing.T) {
		migrationsDir := filepath.Join(t.TempDir(), "migrations")
		writePreparedMigration(t, migrationsDir, "m1", []string{"1001:0"})
		writePreparedMigration(t, migrationsDir, "m2", []string{"1002:0"})
		writePreparedMigration(t, migrationsDir, "m3", []string{"1003:0"})

		r := NewRegistry(context.Background(), migrationsDir)
		r.LoadFromDisk()

		assert.Len(t, r.List(), 3, "List: got %d want 3", len(r.List()))
	})

	t.Run("Idempotent", func(t *testing.T) {
		migrationsDir := filepath.Join(t.TempDir(), "migrations")
		writePreparedMigration(t, migrationsDir, "m1", []string{"1001:0"})

		r := NewRegistry(context.Background(), migrationsDir)
		r.LoadFromDisk()
		r.LoadFromDisk()

		assert.Len(t, r.List(), 1, "List after double load: got %d want 1", len(r.List()))
	})

	t.Run("NameFromDirName", func(t *testing.T) {
		migrationsDir := filepath.Join(t.TempDir(), "migrations")
		writePreparedMigration(t, migrationsDir, "custom-name", []string{"1001:0"})

		r := NewRegistry(context.Background(), migrationsDir)
		r.LoadFromDisk()

		m, ok := r.Get("custom-name")
		assert.True(t, ok, "Get(custom-name) not found")
		assert.Equal(t, "custom-name", m.Name())
		assert.Equal(t, filepath.Join(migrationsDir, "custom-name"), m.Dir())
	})

	t.Run("SkipsNonDirectory", func(t *testing.T) {
		migrationsDir := filepath.Join(t.TempDir(), "migrations")
		writePreparedMigration(t, migrationsDir, "m1", []string{"1001:0"})
		assert.NoError(t, os.WriteFile(filepath.Join(migrationsDir, "stray.txt"), []byte("x"), 0o644))

		r := NewRegistry(context.Background(), migrationsDir)
		r.LoadFromDisk()

		_, ok := r.Get("m1")
		assert.True(t, ok, "m1 should still load")
	})
}

func writeInflightMarker(t *testing.T, dir string, rec jobRecord) {
	t.Helper()
	data, err := json.Marshal(rec)
	assert.NoError(t, err)
	assert.NoError(t, os.WriteFile(filepath.Join(dir, inflightMarkerFile), data, 0o644))
}

func newResumableRegistry(t *testing.T, migrationsDir string) *Registry {
	t.Helper()
	r := NewRegistry(context.Background(), migrationsDir)
	r.loopOptions = []migration.Option{migration.WithActionExecutor(fakeRegistryAction("echo resumed"))}
	r.LoadFromDisk()
	return r
}

func TestRegistryResumeInterrupted(t *testing.T) {
	t.Run("OrphanLoopRunMarkerResumes", func(t *testing.T) {
		migrationsDir := filepath.Join(t.TempDir(), "migrations")
		writeRegistryMigration(t, migrationsDir, "m1")
		writeInflightMarker(t, filepath.Join(migrationsDir, "m1"),
			jobRecord{ID: "migrate-run-1", Kind: jobMigrateRun, StartedAt: "2026-08-15T00:00:00Z"})

		r := newResumableRegistry(t, migrationsDir)
		resumed := r.ResumeInterrupted()
		assert.Equal(t, []string{"m1"}, resumed)

		m, ok := r.Get("m1")
		assert.True(t, ok)
		st := m.Status()
		assert.True(t, st.Running, "resumed migrate-run should be in flight")
		assert.Equal(t, jobMigrateRun, st.Kind)

		old, ok := m.GetJob("migrate-run-1")
		assert.True(t, ok, "interrupted attempt must land in history")
		assert.Equal(t, "interrupted; resumed as "+st.JobId, old.Result)
		assert.NotEmpty(t, old.FinishedAt)

		t.Cleanup(func() {
			m.Cancel(m.Status().JobId)
			waitRegistryIdle(t, m, 5*time.Second)
		})
		waitRegistryIdle(t, m, 5*time.Second)
		assert.Empty(t, r.ResumeInterrupted(), "second call after finish must be a no-op")
	})

	t.Run("ResumedRunUsesDefaultRetryBudget", func(t *testing.T) {
		migrationsDir := filepath.Join(t.TempDir(), "migrations")
		writeRegistryMigration(t, migrationsDir, "m1")
		writeInflightMarker(t, filepath.Join(migrationsDir, "m1"),
			jobRecord{ID: "migrate-run-1", Kind: jobMigrateRun, StartedAt: "2026-08-15T00:00:00Z"})

		var calls atomic.Int32
		r := NewRegistry(context.Background(), migrationsDir)
		r.loopOptions = []migration.Option{migration.WithActionExecutor(
			flakyExecutor(&calls, 1, errors.New("flaky step failure")))}
		r.retryInterval = time.Millisecond
		r.LoadFromDisk()

		require.Equal(t, []string{"m1"}, r.ResumeInterrupted())

		m, ok := r.Get("m1")
		require.True(t, ok)
		waitRegistryIdle(t, m, 5*time.Second)

		jobs := m.Jobs()
		require.NotEmpty(t, jobs)
		assert.Empty(t, jobs[len(jobs)-1].Result, "resumed run must recover under the default budget")
		assert.Contains(t, string(m.logs.Get().snapshot()),
			fmt.Sprintf("retry 1/%d", DefaultRunRetries))
	})

	t.Run("StaleMarkerCleanedWithoutResume", func(t *testing.T) {
		migrationsDir := filepath.Join(t.TempDir(), "migrations")
		writeRegistryMigration(t, migrationsDir, "m1")
		finished := jobRecord{
			ID:         "migrate-run-1",
			Kind:       jobMigrateRun,
			StartedAt:  "2026-08-15T00:00:00Z",
			FinishedAt: "2026-08-15T00:01:00Z",
		}
		writeInflightMarker(t, filepath.Join(migrationsDir, "m1"), finished)
		assert.NoError(t, appendJobRecord(filepath.Join(migrationsDir, "m1", "jobs.jsonl"), &finished))

		r := newResumableRegistry(t, migrationsDir)
		assert.Empty(t, r.ResumeInterrupted())

		m, _ := r.Get("m1")
		assert.False(t, m.Status().Running, "stale marker must not resume")
		_, err := os.Stat(filepath.Join(migrationsDir, "m1", inflightMarkerFile))
		assert.True(t, os.IsNotExist(err), "stale marker must be removed")
	})

	t.Run("ProbeMarkerCleanedWithoutResume", func(t *testing.T) {
		migrationsDir := filepath.Join(t.TempDir(), "migrations")
		writeRegistryMigration(t, migrationsDir, "m1")
		writeInflightMarker(t, filepath.Join(migrationsDir, "m1"),
			jobRecord{ID: "export-probe-1", Kind: jobExportProbe, StartedAt: "2026-08-15T00:00:00Z"})

		r := newResumableRegistry(t, migrationsDir)
		assert.Empty(t, r.ResumeInterrupted())

		m, _ := r.Get("m1")
		assert.False(t, m.Status().Running, "probe marker must not resume")
		rec, ok := m.GetJob("export-probe-1")
		assert.True(t, ok, "dropped probe attempt must land in history")
		assert.Equal(t, "interrupted; not resumed (read-only job)", rec.Result)
		assert.NotEmpty(t, rec.FinishedAt)
		_, err := os.Stat(filepath.Join(migrationsDir, "m1", inflightMarkerFile))
		assert.True(t, os.IsNotExist(err), "probe marker must be removed")
	})

	t.Run("CorruptMarkerKeptForInspection", func(t *testing.T) {
		migrationsDir := filepath.Join(t.TempDir(), "migrations")
		writeRegistryMigration(t, migrationsDir, "m1")
		assert.NoError(t, os.WriteFile(
			filepath.Join(migrationsDir, "m1", inflightMarkerFile), []byte("not json"), 0o644))

		r := newResumableRegistry(t, migrationsDir)
		assert.Empty(t, r.ResumeInterrupted())

		m, _ := r.Get("m1")
		assert.False(t, m.Status().Running)
		_, err := os.Stat(filepath.Join(migrationsDir, "m1", inflightMarkerFile))
		assert.NoError(t, err, "corrupt marker must be kept for inspection")
	})

	t.Run("NoMarkersNoop", func(t *testing.T) {
		migrationsDir := filepath.Join(t.TempDir(), "migrations")
		writeRegistryMigration(t, migrationsDir, "m1")

		r := newResumableRegistry(t, migrationsDir)
		assert.Empty(t, r.ResumeInterrupted())
	})
}

func TestResumeSkipsWhenJobsHistoryUnreadable(t *testing.T) {
	migrationsDir := filepath.Join(t.TempDir(), "migrations")
	writeRegistryMigration(t, migrationsDir, "m1")
	dir := filepath.Join(migrationsDir, "m1")
	finished := jobRecord{
		ID:         "migrate-run-1",
		Kind:       jobMigrateRun,
		StartedAt:  "2026-08-15T00:00:00Z",
		FinishedAt: "2026-08-15T00:01:00Z",
	}
	writeInflightMarker(t, dir, finished)
	// A directory where jobs.jsonl should live makes the read fail with
	// something other than NotExist.
	assert.NoError(t, os.MkdirAll(filepath.Join(dir, "jobs.jsonl"), 0o755))

	r := newResumableRegistry(t, migrationsDir)
	assert.Empty(t, r.ResumeInterrupted(), "unreadable history must skip, not re-dispatch")

	m, _ := r.Get("m1")
	assert.False(t, m.Status().Running)
	_, err := os.Stat(filepath.Join(dir, inflightMarkerFile))
	assert.NoError(t, err, "marker must be kept when history is unreadable")
}

func TestResumeErrorSurfacesInStatus(t *testing.T) {
	t.Run("DispatchFailureRecorded", func(t *testing.T) {
		// A loop-status claiming a completed verify step with no shard
		// status file: reconcile fails closed, the resume dispatch fails,
		// and the failure must be visible in Status.
		migrationsDir := filepath.Join(t.TempDir(), "migrations")
		dir := filepath.Join(migrationsDir, "m1")
		plan := schema.MigrationLoopPlan{
			Runtime: &schema.MigrationLoopRuntime{
				Vminsert: "http://i/{tenant}",
				Window:   schema.MigrationLoopWindow{Start: "2026-06-16T00:00:00Z", End: "2026-06-17T00:00:00Z", StepInterval: "1d"},
			},
			Vmselect: &schema.MigrationLoopVmselect{Origin: "http://o/{tenant}"},
			Tenants: []schema.MigrationLoopTenant{{
				Name:        "t1",
				TenantsFile: "t1.json",
				StatusFile:  "t1-status.json",
				Cmds: []schema.MigrationLoopCmd{
					{Step: "manifest", Action: &schema.MigrationLoopAction{Kind: schema.ActionManifest, TenantsFile: "t1.json", Manifest: "t1-manifest.json"}},
					{Step: "migrate", Action: &schema.MigrationLoopAction{Kind: schema.ActionMigrate, Manifest: "t1-manifest.json", StepInterval: "day"}},
					{Step: "verify", Action: &schema.MigrationLoopAction{Kind: schema.ActionVerify, Manifest: "t1-manifest.json"}},
				},
			}},
		}
		writeRegistryPlan(t, migrationsDir, "m1", plan)
		loopStatus := schema.MigrationLoopStatus{
			UpdatedAt: "2026-08-15T00:00:00Z",
			Tenants: []schema.MigrationLoopTenantStatus{{
				Name:        "t1",
				Completed:   true,
				CurrentStep: "done",
				Steps: []schema.MigrationLoopStepStatus{
					{Step: "migrate", Status: "completed"},
					{Step: "verify", Status: "completed"},
				},
			}},
		}
		seed, _ := json.MarshalIndent(loopStatus, "", "  ")
		assert.NoError(t, os.WriteFile(filepath.Join(dir, "migration-loop-status.json"), append(seed, '\n'), 0o644))
		writeInflightMarker(t, dir, jobRecord{ID: "migrate-run-1", Kind: jobMigrateRun, StartedAt: "2026-08-15T00:00:00Z"})

		r := newResumableRegistry(t, migrationsDir)
		assert.Empty(t, r.ResumeInterrupted())

		m, _ := r.Get("m1")
		st := m.Status()
		assert.False(t, st.Running)
		assert.Contains(t, st.ResumeError, "claims completed", "ResumeError = %q", st.ResumeError)
		assert.NotEmpty(t, st.ResumeErrorAt)
	})

	t.Run("SuccessfulResumeClearsRecord", func(t *testing.T) {
		migrationsDir := filepath.Join(t.TempDir(), "migrations")
		writeRegistryMigration(t, migrationsDir, "m1")
		writeInflightMarker(t, filepath.Join(migrationsDir, "m1"),
			jobRecord{ID: "migrate-run-1", Kind: jobMigrateRun, StartedAt: "2026-08-15T00:00:00Z"})

		r := newResumableRegistry(t, migrationsDir)
		assert.Equal(t, []string{"m1"}, r.ResumeInterrupted())

		m, _ := r.Get("m1")
		assert.Empty(t, m.Status().ResumeError, "a dispatched resume must clear any earlier failure record")
		t.Cleanup(func() {
			m.Cancel(m.Status().JobId)
			waitRegistryIdle(t, m, 5*time.Second)
		})
	})

	t.Run("CorruptMarkerRecorded", func(t *testing.T) {
		migrationsDir := filepath.Join(t.TempDir(), "migrations")
		dir := filepath.Join(migrationsDir, "m1")
		writeRegistryMigration(t, migrationsDir, "m1")
		assert.NoError(t, os.WriteFile(filepath.Join(dir, inflightMarkerFile), []byte("not json"), 0o644))

		r := newResumableRegistry(t, migrationsDir)
		assert.Empty(t, r.ResumeInterrupted())

		m, _ := r.Get("m1")
		assert.Contains(t, m.Status().ResumeError, "read in-flight marker")
	})
}