package migration

import (
	"bytes"
	"context"
	"io"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

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

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

type recordingRunner struct {
	calls    [][]string
	behavior func(args []string) (exitCode int, stdout string, bytes int64)
}

func (r *recordingRunner) Run(ctx context.Context, args []string, out io.Writer) (int, int64) {
	r.calls = append(r.calls, args)
	exit, text, bytes := r.behavior(args)
	if text != "" {
		_, _ = out.Write([]byte(text))
	}
	return exit, bytes
}

func combinedActionOutput(w io.Writer) ActionOutput {
	return ActionOutput{Progress: w, Log: w}
}

func writeMigrationManifestFixture(t *testing.T, dir string) string {
	t.Helper()
	path := filepath.Join(dir, "tenants-manifest.json")
	content := `{
  "region": "prod",
  "namespace": "observability",
  "vmcluster": "vmcluster",
  "stepInterval": "1d",
  "tenants": [
    {
      "tenant": "1001:0",
      "sourceAddr": "http://src/select/1001:0/prometheus",
      "destinationAddr": "http://dst/insert/1001:0/prometheus",
      "shards": [
        {"shardID": "1001:0:1", "timeStart": "2026-06-16T00:00:00Z", "timeEnd": "2026-06-17T00:00:00Z"},
        {"shardID": "1001:0:2", "timeStart": "2026-06-17T00:00:00Z", "timeEnd": "2026-06-18T00:00:00Z"}
      ]
    }
  ]
}`
	require.NoError(t, os.WriteFile(path, []byte(content), 0o644))
	return path
}

func TestParseVmctlBytesLine(t *testing.T) {
	f := func(line string, want int64) {
		t.Helper()
		assert.Equal(t, want, parseVmctlBytesLine(line), "line=%q", line)
	}
	f("total bytes: 1.5 GB", 1500*1000*1000)
	f("total bytes: 1024 B", 1024)
	f("total bytes: 1.23 MB", 1230000)
	f("total bytes: 5 KB", 5000)
	f("total bytes: 0 B", 0)
	f("no bytes here", 0)
	f("Total bytes: 1 GB", 1000*1000*1000)
}

func TestByteSnifferWriterParsesBytesLine(t *testing.T) {
	var sniffer byteSnifferWriter
	_, _ = sniffer.Write([]byte("vmctl progress\n total bytes: 2 GB \nmore output\n"))
	assert.Equal(t, int64(2*1000*1000*1000), sniffer.bytes())
}

func TestVmctlCommandTimedWithFlags(t *testing.T) {
	tenant := schema.ManifestTenant{SourceAddr: "http://s", DestinationAddr: "http://d"}
	shard := schema.ManifestShard{
		ShardID:   "x:1",
		TimeStart: "2026-06-16T00:00:00Z",
		TimeEnd:   "2026-06-17T00:00:00Z",
	}
	rate := 1000
	conc := 4
	cmd := vmctlCommand(tenant, shard, "day", &rate, &conc)
	joined := strings.Join(cmd, " ")
	for _, want := range []string{
		"--vm-native-filter-match {__name__!=\"\"}",
		"--vm-native-filter-time-start 2026-06-16T00:00:00Z",
		"--vm-native-filter-time-end 2026-06-17T00:00:00Z",
		"--vm-native-step-interval day",
		"--vm-rate-limit 1000",
		"--vm-concurrency 4",
		"--disable-progress-bar",
		"-s",
	} {
		assert.Contains(t, joined, want)
	}
}

func TestRunMigrateHappyPath(t *testing.T) {
	dir := t.TempDir()
	manifestPath := writeMigrationManifestFixture(t, dir)
	statusPath := status.PathForManifest(manifestPath)

	manifest, err := status.LoadManifest(manifestPath)
	require.NoError(t, err)
	doc := status.BuildInitialStatus(manifest, manifestPath)
	require.NoError(t, status.WriteJSONAtomic(statusPath, doc))

	rr := &recordingRunner{
		behavior: func(args []string) (int, string, int64) {
			start := ""
			for i, a := range args {
				if a == "--vm-native-filter-time-start" && i+1 < len(args) {
					start = args[i+1]
				}
			}
			if start == "2026-06-16T00:00:00Z" {
				return 0, "vmctl done\ntotal bytes: 2.5 GB\n", 2500 * 1000 * 1000
			}
			return 0, "vmctl done\ntotal bytes: 1 GB\n", 1000 * 1000 * 1000
		},
	}

	args := VmctlOptions{
		Manifest:          manifestPath,
		VmctlStepInterval: "day",
		VmConcurrency:     new(24),
	}
	var stdout bytes.Buffer
	require.NoError(t, RunVmctl(context.Background(), args, rr, combinedActionOutput(&stdout)))
	out := stdout.String()
	assert.Contains(t, out, "migration start:")
	assert.Contains(t, out, "done migration")
	assert.Contains(t, out, "migration done:")

	loaded, err := status.LoadDocument(statusPath)
	require.NoError(t, err)
	for _, s := range loaded.Shards {
		assert.Equal(t, schema.PhaseImported, s.Phase)
		require.NotNil(t, s.Bytes)
		assert.Greater(t, *s.Bytes, int64(0))
		assert.Equal(t, 1, s.Attempts)
	}

	args2 := VmctlOptions{Manifest: manifestPath, VmctlStepInterval: "day"}
	var stdout2 bytes.Buffer
	require.NoError(t, RunVmctl(context.Background(), args2, rr, combinedActionOutput(&stdout2)))
	assert.Contains(t, stdout2.String(), "selected=0")
}

func TestRunVmctlSeparatesProgressAndActionLog(t *testing.T) {
	dir := t.TempDir()
	manifestPath := writeMigrationManifestFixture(t, dir)
	runner := &recordingRunner{behavior: func([]string) (int, string, int64) {
		return 0, "", 0
	}}
	var progress, actionLog bytes.Buffer

	err := RunVmctl(context.Background(), VmctlOptions{
		Manifest:          manifestPath,
		VmctlStepInterval: "day",
	}, runner, ActionOutput{Progress: &progress, Log: &actionLog})
	require.NoError(t, err)

	assert.Contains(t, progress.String(), "migration start:")
	assert.NotContains(t, progress.String(), "vmctl command:")
	assert.Contains(t, actionLog.String(), "vmctl command:")
	assert.NotContains(t, actionLog.String(), "migration start:")
}

func TestRunMigrateShardFailure(t *testing.T) {
	dir := t.TempDir()
	manifestPath := writeMigrationManifestFixture(t, dir)
	statusPath := status.PathForManifest(manifestPath)
	manifest, err := status.LoadManifest(manifestPath)
	require.NoError(t, err)
	doc := status.BuildInitialStatus(manifest, manifestPath)
	require.NoError(t, status.WriteJSONAtomic(statusPath, doc))

	rr := &recordingRunner{
		behavior: func(args []string) (int, string, int64) {
			return 1, "vmctl error: connection refused\n", 0
		},
	}

	args := VmctlOptions{Manifest: manifestPath, VmctlStepInterval: "day"}
	var stdout bytes.Buffer
	err = RunVmctl(context.Background(), args, rr, combinedActionOutput(&stdout))
	assert.Error(t, err)
	var ee *exitcode.Error
	require.ErrorAs(t, err, &ee)
	assert.Equal(t, exitcode.Partial, ee.Code)
	loaded, err := status.LoadDocument(statusPath)
	require.NoError(t, err)
	for _, s := range loaded.Shards {
		assert.Equal(t, schema.PhaseFailed, s.Phase)
		require.NotNil(t, s.LastError)
		assert.Contains(t, *s.LastError, "code 1")
	}
}

func TestRunMigrateShardCancelledMidRun(t *testing.T) {
	dir := t.TempDir()
	manifestPath := writeMigrationManifestFixture(t, dir)
	statusPath := status.PathForManifest(manifestPath)
	manifest, err := status.LoadManifest(manifestPath)
	require.NoError(t, err)
	doc := status.BuildInitialStatus(manifest, manifestPath)
	require.NoError(t, status.WriteJSONAtomic(statusPath, doc))

	ctx, cancel := context.WithCancel(context.Background())
	started := make(chan struct{})
	rr := &recordingRunner{behavior: func(args []string) (int, string, int64) {
		close(started)  // shard is Running, vmctl is mid-call
		<-ctx.Done()    // simulate vmctl blocking until ctx kills it
		return 1, "", 0 // killed process exits non-zero
	}}

	done := make(chan error, 1)
	go func() {
		done <- RunVmctl(ctx, VmctlOptions{Manifest: manifestPath, VmctlStepInterval: "day"}, rr, combinedActionOutput(&bytes.Buffer{}))
	}()
	<-started
	cancel()

	select {
	case err := <-done:
		assert.Error(t, err, "cancelled run: expected non-nil error, got nil")
	case <-time.After(5 * time.Second):
		t.Fatal("Run did not return after cancel")
	}

	loaded, err := status.LoadDocument(statusPath)
	require.NoError(t, err)
	for _, s := range loaded.Shards {
		assert.NotEqual(t, schema.PhaseImported, s.Phase, "shard %s: phase imported after cancel; a killed migration must not be marked complete", s.ShardID)
	}
}

func TestRunMigrateReportsFailureWhenRetriedShardFailsAgain(t *testing.T) {
	dir := t.TempDir()
	manifestPath := writeMigrationManifestFixture(t, dir)
	statusPath := status.PathForManifest(manifestPath)
	manifest, err := status.LoadManifest(manifestPath)
	require.NoError(t, err)
	doc := status.BuildInitialStatus(manifest, manifestPath)
	doc.Shards[0].Phase = schema.PhaseFailed
	doc.Shards[0].Attempts = 1
	doc.Shards[1].Phase = schema.PhaseImported
	require.NoError(t, status.WriteJSONAtomic(statusPath, doc))
	runner := &recordingRunner{behavior: func([]string) (int, string, int64) {
		return 1, "vmctl failed again\n", 0
	}}
	err = RunVmctl(context.Background(), VmctlOptions{
		Manifest:          manifestPath,
		VmctlStepInterval: "day",
	}, runner, combinedActionOutput(&bytes.Buffer{}))
	var exitErr *exitcode.Error
	assert.Error(t, err)
	require.ErrorAs(t, err, &exitErr)
	assert.Equal(t, exitcode.Partial, exitErr.Code)
}

func TestRunMigrateResumesMidShard(t *testing.T) {
	dir := t.TempDir()
	manifestPath := writeMigrationManifestFixture(t, dir)
	statusPath := status.PathForManifest(manifestPath)
	manifest, err := status.LoadManifest(manifestPath)
	require.NoError(t, err)
	doc := status.BuildInitialStatus(manifest, manifestPath)
	doc.Shards[0].Phase = schema.PhaseImported
	doc.Shards[0].Bytes = new(int64(500 * 1000 * 1000))
	doc.Shards[0].Attempts = 1
	require.NoError(t, status.WriteJSONAtomic(statusPath, doc))

	var recordedShard string
	rr := &recordingRunner{
		behavior: func(args []string) (int, string, int64) {
			for i, a := range args {
				if a == "--vm-native-filter-time-start" && i+1 < len(args) {
					recordedShard = args[i+1]
				}
			}
			return 0, "total bytes: 800 MB\n", 800 * 1000 * 1000
		},
	}
	args := VmctlOptions{Manifest: manifestPath, VmctlStepInterval: "day"}
	var stdout bytes.Buffer
	require.NoError(t, RunVmctl(context.Background(), args, rr, combinedActionOutput(&stdout)))
	assert.Contains(t, stdout.String(), "selected=1")
	assert.Equal(t, "2026-06-17T00:00:00Z", recordedShard)
}

func TestMutateShardFinalLeavesNoStrayFieldsOnIllegalTransition(t *testing.T) {
	doc := status.BuildInitialStatus(&schema.Manifest{Tenants: []schema.ManifestTenant{{
		Tenant: "1001:0",
		Shards: []schema.ManifestShard{{ShardID: "1001:0:1"}},
	}}}, "m.json")
	// pending → imported is illegal (must pass through running); the
	// mutation must fail without leaving a stray FinishedAt behind.
	_, err := mutateShardFinal(doc, "1001:0:1", 0, 42)
	assert.Error(t, err)
	s := status.FindShard(doc, "1001:0:1")
	require.NotNil(t, s)
	assert.Nil(t, s.FinishedAt, "illegal transition must not mutate the in-memory shard")
	assert.Equal(t, schema.PhasePending, s.Phase)

	require.NoError(t, mutateShardRunning(doc, "1001:0:1", 1))
	label, err := mutateShardFinal(doc, "1001:0:1", 0, 42)
	require.NoError(t, err)
	s = status.FindShard(doc, "1001:0:1")
	require.NotNil(t, s)
	require.NotNil(t, s.FinishedAt)
	require.NotNil(t, s.Bytes)
	assert.Equal(t, schema.PhaseImported, s.Phase)
	assert.Equal(t, int64(42), *s.Bytes)
	assert.NotEmpty(t, label)
}

func TestFailShardNotFound(t *testing.T) {
	doc := &schema.Document{Shards: []schema.Shard{
		{ShardID: "1001:0:1", Phase: schema.PhaseImported},
	}}

	require.NoError(t, failShardNotFound(doc, "1001:0:1"))
	s := status.FindShard(doc, "1001:0:1")
	require.NotNil(t, s)
	assert.Equal(t, schema.PhaseFailed, s.Phase)
	if assert.NotNil(t, s.LastError) {
		assert.Equal(t, "shardID not found in manifest", *s.LastError)
	}
	assert.NotNil(t, s.FinishedAt)

	err := failShardNotFound(doc, "absent")
	assert.ErrorContains(t, err, "shard absent: not present in status document")
}

func TestDefaultVmctlRunnerRun(t *testing.T) {
	t.Run("SuccessParsesBytes", func(t *testing.T) {
		script := filepath.Join(t.TempDir(), "fake-vmctl")
		body := "#!/bin/sh\n" +
			"echo '2026/08/17 12:00:00 INF copying shards'\n" +
			"echo '2026/08/17 12:00:01 INF shard copied; total bytes: 1.54 GB; speed: 12.34 MB/s'\n"
		require.NoError(t, os.WriteFile(script, []byte(body), 0o755))

		var out bytes.Buffer
		code, got := NewDefaultVmctlRunner().Run(context.Background(), []string{script, "vm-native", "--flag"}, &out)
		assert.Zero(t, code)
		assert.Equal(t, int64(1540000000), got)
		assert.Contains(t, out.String(), "total bytes: 1.54 GB")
	})

	t.Run("NonZeroExit", func(t *testing.T) {
		script := filepath.Join(t.TempDir(), "fake-vmctl")
		require.NoError(t, os.WriteFile(script, []byte("#!/bin/sh\necho boom 1>&2; exit 3\n"), 0o755))
		var out bytes.Buffer
		code, got := NewDefaultVmctlRunner().Run(context.Background(), []string{script}, &out)
		assert.Equal(t, 3, code)
		assert.Zero(t, got)
		assert.Contains(t, out.String(), "boom")
	})

	t.Run("BinMissing", func(t *testing.T) {
		var out bytes.Buffer
		code, got := NewDefaultVmctlRunner().Run(context.Background(), []string{filepath.Join(t.TempDir(), "absent-vmctl")}, &out)
		assert.Equal(t, 127, code)
		assert.Zero(t, got)
	})
}