package server

import (
	"context"
	"encoding/json"
	"errors"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"testing"
	"time"

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

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

type echoJob struct {
	kind  jobKind
	delay time.Duration
	err   error
}

func (j *echoJob) Kind() jobKind { return j.kind }
func (j *echoJob) Run(ctx context.Context) error {
	if j.delay > 0 {
		select {
		case <-time.After(j.delay):
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	return j.err
}

type artifactEchoJob struct {
	echoJob
	artifact string
}

type failingArtifactJob struct{ echoJob }

func (*failingArtifactJob) WriteArtifact(string, *jobRecord) error {
	return errors.New("artifact storage failed")
}

type warningEchoJob struct {
	echoJob
	warnings []string
}

func (j *warningEchoJob) Warnings() []string { return append([]string(nil), j.warnings...) }

func (j *artifactEchoJob) WriteArtifact(path string, _ *jobRecord) error {
	return os.WriteFile(path, []byte(j.artifact), 0o644)
}

func newTestMigration(t *testing.T) *Migration {
	t.Helper()
	dir := t.TempDir()
	m := &Migration{
		name:    "test",
		dir:     dir,
		plan:    nil,
		rootCtx: context.Background(),
		logs:    newLogHub(),
	}
	return m
}

func waitMigrationIdle(t *testing.T, m *Migration, d time.Duration) {
	t.Helper()
	deadline := time.Now().Add(d)
	for time.Now().Before(deadline) {
		if m.Status().Running {
			time.Sleep(10 * time.Millisecond)
			continue
		}
		st := m.Status()
		if st.LastJobId == "" {
			return // no jobs yet, idle is genuine
		}
		if r, ok := m.GetJob(st.LastJobId); ok && r.FinishedAt != "" {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	assert.Fail(t, "migration still running or jobs.jsonl stale after %v", d)
}

// startJob starts job and drops the returned id; callers that need it read
// m.Status().JobId.
func startJob(t *testing.T, m *Migration, job job) {
	t.Helper()
	_, err := m.start(job)
	assert.NoError(t, err)
}

// runJobToIdle starts job on a fresh migration and waits for it to finish.
func runJobToIdle(t *testing.T, job job) *Migration {
	t.Helper()
	m := newTestMigration(t)
	startJob(t, m, job)
	waitMigrationIdle(t, m, 2*time.Second)
	return m
}

func TestMigration(t *testing.T) {
	t.Run("StartAssignsJobId", func(t *testing.T) {
		m := newTestMigration(t)
		startJob(t, m, &echoJob{kind: jobMigrateRun})
		st := m.Status()
		assert.True(t, st.Running, "Running = false, want true")
		assert.NotEmpty(t, st.JobId, "JobId empty after start")
		assert.True(t, strings.HasPrefix(st.JobId, "migrate-run-"), "JobId = %q, want prefix migrate-run-", st.JobId)
		waitMigrationIdle(t, m, 2*time.Second)
	})

	t.Run("JobFinishedAppendedToJSONL", func(t *testing.T) {
		m := runJobToIdle(t, &echoJob{kind: jobMigrateRun})
		got, err := readJobRecords(m.jobsJSONLPath())
		assert.NoError(t, err)
		assert.Len(t, got, 1)
		r := got[0]
		assert.True(t, strings.HasPrefix(r.ID, "migrate-run-"), "ID = %q", r.ID)
		assert.Equal(t, jobMigrateRun, r.Kind)
		assert.NotEmpty(t, r.FinishedAt, "FinishedAt empty")
		assert.False(t, r.Cancelled, "Cancelled = true, want false")
	})

	t.Run("JobWarningsPersisted", func(t *testing.T) {
		m := runJobToIdle(t, &warningEchoJob{
			echoJob:  echoJob{kind: jobMigrateRun},
			warnings: []string{"tenant=t1 action log unavailable"},
		})
		got, err := readJobRecords(m.jobsJSONLPath())
		assert.NoError(t, err)
		assert.Len(t, got, 1)
		assert.Equal(t, []string{"tenant=t1 action log unavailable"}, got[0].Warnings)
	})

	t.Run("ArtifactFailureWarningPersisted", func(t *testing.T) {
		m := runJobToIdle(t, &failingArtifactJob{echoJob{kind: jobExportProbe}})
		got, err := readJobRecords(m.jobsJSONLPath())
		assert.NoError(t, err)
		assert.Len(t, got, 1)
		assert.Len(t, got[0].Warnings, 1)
		assert.Contains(t, got[0].Warnings[0], "artifact unavailable")
	})

	t.Run("LastResultFromJSONL", func(t *testing.T) {
		m := newTestMigration(t)
		startJob(t, m, &echoJob{kind: jobMigrateRun})
		waitMigrationIdle(t, m, 2*time.Second)
		startJob(t, m, &echoJob{kind: jobExportProbe, err: errors.New("boom")})
		waitMigrationIdle(t, m, 2*time.Second)

		st := m.Status()
		assert.Equal(t, jobExportProbe, st.LastKind, "LastKind = %q, want %q", st.LastKind, jobExportProbe)
		assert.Equal(t, "boom", st.LastResult, "LastResult = %q, want %q", st.LastResult, "boom")
		got := m.Jobs()
		assert.Len(t, got, 2, "len(Jobs()) = %d, want 2", len(got))
	})

	t.Run("StatusReflectsRunning", func(t *testing.T) {
		m := newTestMigration(t)
		startJob(t, m, &echoJob{kind: jobMigrateRun, delay: 200 * time.Millisecond})
		st := m.Status()
		assert.True(t, st.Running, "Running = false during in-flight job")
		assert.NotEmpty(t, st.JobId, "JobId empty during in-flight job")
		waitMigrationIdle(t, m, 2*time.Second)
	})

	t.Run("JobArtifactLoopKindNotFound", func(t *testing.T) {
		m := runJobToIdle(t, &echoJob{kind: jobMigrateRun})
		st := m.Status()
		_, err := m.JobArtifact(st.LastJobId)
		assert.Error(t, err, "JobArtifact on loop job: want error, got nil")
		assert.True(t, os.IsNotExist(err), "err = %v, want os.IsNotExist", err)
	})

	t.Run("JobArtifactWritesWhenWriter", func(t *testing.T) {
		m := runJobToIdle(t, &artifactEchoJob{
			echoJob:  echoJob{kind: jobExportProbe},
			artifact: `{"k":"v"}`,
		})
		st := m.Status()
		got, err := m.JobArtifact(st.LastJobId)
		assert.NoError(t, err)
		var doc map[string]any
		assert.NoError(t, json.Unmarshal(got, &doc), "artifact is not JSON")
		assert.Equal(t, "v", doc["k"], "artifact = %s, want k=v", string(got))
		path := filepath.Join(m.dir, "artifacts", st.LastJobId+".json")
		_, err = os.Stat(path)
		assert.NoError(t, err, "artifact file missing")
	})

	t.Run("FailureDoesNotWriteArtifact", func(t *testing.T) {
		job := &artifactEchoJob{
			echoJob:  echoJob{kind: jobExportProbe, err: errors.New("broken config")},
			artifact: `{"mustNotExist":true}`,
		}
		m := runJobToIdle(t, job)
		_, err := m.JobArtifact(m.Status().LastJobId)
		assert.True(t, os.IsNotExist(err), "failure artifact error = %v, want os.ErrNotExist", err)
	})

	t.Run("InflightMarkerLifecycle", func(t *testing.T) {
		m := newTestMigration(t)
		startJob(t, m, &echoJob{kind: jobMigrateRun, delay: 200 * time.Millisecond})
		data, err := os.ReadFile(m.inflightMarkerPath())
		assert.NoError(t, err, "in-flight marker missing while job runs")
		var marker jobRecord
		assert.NoError(t, json.Unmarshal(data, &marker))
		assert.Equal(t, m.Status().JobId, marker.ID, "marker ID must match running job")
		waitMigrationIdle(t, m, 2*time.Second)

		_, err = os.Stat(m.inflightMarkerPath())
		assert.True(t, os.IsNotExist(err), "marker must be removed after normal finish")
	})

	t.Run("JobsMergeInflightMarker", func(t *testing.T) {
		m := newTestMigration(t)
		startJob(t, m, &echoJob{kind: jobMigrateRun, delay: 200 * time.Millisecond})
		got := m.Jobs()
		assert.Len(t, got, 1, "in-flight job must be visible in Jobs()")
		assert.Equal(t, m.Status().JobId, got[0].ID)
		assert.Empty(t, got[0].FinishedAt, "in-flight record must carry no finishedAt")
		waitMigrationIdle(t, m, 2*time.Second)

		got = m.Jobs()
		assert.Len(t, got, 1)
		assert.NotEmpty(t, got[0].FinishedAt, "finished record must replace the marker view")
	})

	t.Run("JobsShowInterruptedMarker", func(t *testing.T) {
		m := runJobToIdle(t, &echoJob{kind: jobMigrateRun})
		// Submitted after the finished job, so it sorts last.
		later := status.FormatTime(status.UTCNow().Add(time.Hour))
		interrupted := &jobRecord{ID: "migrate-run-999", Kind: jobMigrateRun, StartedAt: later}
		assert.NoError(t, status.DumpJSONAtomic(m.inflightMarkerPath(), interrupted))
		got := m.Jobs()
		require.Len(t, got, 2, "interrupted marker must surface in Jobs()")
		last := got[len(got)-1]
		assert.Equal(t, "migrate-run-999", last.ID, "later-submitted marker record sorts last")
		assert.Empty(t, last.FinishedAt)
	})

	t.Run("JobsDedupStaleMarker", func(t *testing.T) {
		m := runJobToIdle(t, &echoJob{kind: jobMigrateRun})
		// Crash between the jobs.jsonl append and the marker removal: the
		// marker's ID is already in history.
		done := m.Jobs()[0]
		assert.NoError(t, status.DumpJSONAtomic(m.inflightMarkerPath(), &done))
		got := m.Jobs()
		assert.Len(t, got, 1, "stale marker must not duplicate its record")
	})

	t.Run("JobsSortedByParsedStartedAt", func(t *testing.T) {
		m := newTestMigration(t)
		// Same second, different fraction width: lexical comparison would
		// order the whole-second record after the fractional one.
		assert.NoError(t, appendJobRecord(m.jobsJSONLPath(), &jobRecord{
			ID:         "migrate-run-2",
			Kind:       jobMigrateRun,
			StartedAt:  "2026-08-20T10:00:00.5Z",
			FinishedAt: "2026-08-20T10:00:01Z",
		}))
		assert.NoError(t, appendJobRecord(m.jobsJSONLPath(), &jobRecord{
			ID:         "migrate-run-1",
			Kind:       jobMigrateRun,
			StartedAt:  "2026-08-20T10:00:00Z",
			FinishedAt: "2026-08-20T10:00:00.9Z",
		}))
		got := m.Jobs()
		require.Len(t, got, 2)
		assert.Equal(t, "migrate-run-1", got[0].ID, "whole-second StartedAt must sort before the later fractional one")
	})

	t.Run("ManualCancelRemovesMarkerAndRecordsCancelled", func(t *testing.T) {
		m := newTestMigration(t)
		startJob(t, m, &echoJob{kind: jobMigrateRun, delay: 30 * time.Second})
		assert.NoError(t, m.Cancel(m.Status().JobId))
		waitMigrationIdle(t, m, 2*time.Second)

		_, err := os.Stat(m.inflightMarkerPath())
		assert.True(t, os.IsNotExist(err), "marker must be removed after manual cancel")
		st := m.Status()
		r, ok := m.GetJob(st.LastJobId)
		assert.True(t, ok, "cancelled job must be recorded")
		assert.True(t, r.Cancelled, "record.Cancelled = false, want true")
	})

	t.Run("RootShutdownKeepsMarkerAndSkipsRecord", func(t *testing.T) {
		rootCtx, cancelRoot := context.WithCancel(context.Background())
		m := newTestMigration(t)
		m.rootCtx = rootCtx
		startJob(t, m, &echoJob{kind: jobMigrateRun, delay: 30 * time.Second})
		cancelRoot()
		waitMigrationIdle(t, m, 2*time.Second)

		data, err := os.ReadFile(m.inflightMarkerPath())
		assert.NoError(t, err, "marker must survive process shutdown")
		var marker jobRecord
		assert.NoError(t, json.Unmarshal(data, &marker))
		raw, err := readJobRecords(m.jobsJSONLPath())
		assert.NoError(t, err)
		assert.False(t, slices.ContainsFunc(raw, func(r jobRecord) bool { return r.ID == marker.ID }),
			"interrupted job must not appear in jobs.jsonl")
		_, visible := m.GetJob(marker.ID)
		assert.True(t, visible, "interrupted job must stay visible via the marker merge")
	})
}

func TestManualCancelThenRootShutdownFinalizesAsManualStop(t *testing.T) {
	rootCtx, cancelRoot := context.WithCancel(context.Background())
	m := newTestMigration(t)
	m.rootCtx = rootCtx
	startJob(t, m, &echoJob{kind: jobMigrateRun, delay: 30 * time.Second})

	// Operator cancels first; shutdown lands while the job is unwinding.
	assert.NoError(t, m.Cancel(m.Status().JobId))
	cancelRoot()
	waitMigrationIdle(t, m, 5*time.Second)

	_, err := os.Stat(m.inflightMarkerPath())
	assert.True(t, os.IsNotExist(err), "manual cancel must finalize (remove marker) even under shutdown")
	st := m.Status()
	r, ok := m.GetJob(st.LastJobId)
	assert.True(t, ok, "manually cancelled job must be recorded")
	assert.True(t, r.Cancelled, "record.Cancelled = false, want true")
}

func TestRootShutdownAloneStillResumes(t *testing.T) {
	rootCtx, cancelRoot := context.WithCancel(context.Background())
	m := newTestMigration(t)
	m.rootCtx = rootCtx
	startJob(t, m, &echoJob{kind: jobMigrateRun, delay: 30 * time.Second})
	cancelRoot()
	waitMigrationIdle(t, m, 5*time.Second)

	_, err := os.Stat(m.inflightMarkerPath())
	assert.NoError(t, err, "shutdown without manual cancel keeps the marker for resume")
}