package status

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

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

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

func TestOpenLoadsExistingDocument(t *testing.T) {
	path := filepath.Join(t.TempDir(), "status.json")
	doc := &schema.Document{
		ManifestRef: "tenants-1-manifest.json",
		Shards: []schema.Shard{
			{ShardID: "x:1", Tenant: "x", Phase: schema.PhaseImported},
		},
	}
	require.NoError(t, WriteJSONAtomic(path, doc))

	store, err := Open(path, nil)
	require.NoError(t, err)
	d := store.Document()
	assert.Equal(t, "tenants-1-manifest.json", d.ManifestRef)
	assert.Len(t, d.Shards, 1)
	assert.Equal(t, schema.PhaseImported, d.Shards[0].Phase)
}

func TestOpenInitWritesFreshFileWhenAbsent(t *testing.T) {
	path := filepath.Join(t.TempDir(), "status.json")
	_, err := os.Stat(path)
	assert.ErrorIs(t, err, os.ErrNotExist)

	doc := &schema.Document{
		ManifestRef: "fresh.json",
		Shards: []schema.Shard{
			{ShardID: "x:1", Tenant: "x", Phase: schema.PhasePending},
			{ShardID: "x:2", Tenant: "x", Phase: schema.PhasePending},
		},
	}
	store, err := Open(path, func() (*schema.Document, error) { return doc, nil })
	require.NoError(t, err)
	assert.Equal(t, 2, store.Document().Summary.Total, "RefreshSummary should have run")

	_, err = os.Stat(path)
	assert.NoError(t, err, "init file not written")
	reloaded, err := LoadDocument(path)
	require.NoError(t, err)
	assert.Equal(t, "fresh.json", reloaded.ManifestRef)
	assert.Equal(t, 2, reloaded.Summary.Total)
}

func TestOpenAbsentFileNoInitErrors(t *testing.T) {
	_, err := Open(filepath.Join(t.TempDir(), "missing.json"), nil)
	assert.Error(t, err)
}

func TestMutateRefreshesSummaryAndPersists(t *testing.T) {
	path := filepath.Join(t.TempDir(), "status.json")
	store, err := Open(path, func() (*schema.Document, error) {
		return &schema.Document{
			Shards: []schema.Shard{{ShardID: "x:1", Tenant: "x", Phase: schema.PhasePending}},
		}, nil
	})
	require.NoError(t, err)

	require.NoError(t, store.Mutate(context.Background(), func(doc *schema.Document) error {
		doc.Shards[0].Phase = schema.PhaseImported
		return nil
	}))

	assert.Equal(t, 1, store.Document().Summary.Imported)

	reloaded, err := LoadDocument(path)
	require.NoError(t, err)
	assert.Equal(t, 1, reloaded.Summary.Imported)
}

func TestMutateFnErrorSkipsWrite(t *testing.T) {
	path := filepath.Join(t.TempDir(), "status.json")
	store, err := Open(path, func() (*schema.Document, error) {
		return &schema.Document{Shards: []schema.Shard{{ShardID: "x:1", Phase: schema.PhasePending}}}, nil
	})
	require.NoError(t, err)

	require.NoError(t, store.Mutate(context.Background(), func(d *schema.Document) error {
		d.Shards[0].Attempts = 1
		return nil
	}))

	sentinel := errors.New("boom")
	err = store.Mutate(context.Background(), func(d *schema.Document) error {
		d.Shards[0].Attempts = 99 // partial mutation
		return sentinel
	})
	assert.ErrorIs(t, err, sentinel)

	reloaded, err := LoadDocument(path)
	require.NoError(t, err)
	assert.Equal(t, 1, reloaded.Shards[0].Attempts, "writeFn should not have run on fn error")
}

func TestMutateWriteFailureLeavesDocUnchanged(t *testing.T) {
	doc := &schema.Document{
		Shards: []schema.Shard{{ShardID: "x:1", Tenant: "x", Phase: schema.PhasePending}},
	}
	RefreshSummary(doc)
	writeErr := errors.New("disk full")
	store := &Store{
		writeFn: func(*schema.Document) error { return writeErr },
		doc:     doc,
	}

	err := store.Mutate(context.Background(), func(d *schema.Document) error {
		d.Shards[0].Phase = schema.PhaseImported
		return nil
	})
	assert.ErrorIs(t, err, writeErr)

	store.View(func(d *schema.Document) {
		assert.Equal(t, schema.PhasePending, d.Shards[0].Phase, "in-memory doc must keep the last persisted state")
		assert.Equal(t, 1, d.Summary.Pending)
		assert.Zero(t, d.Summary.Imported)
	})
}

func TestMutateCtxCancelledShortCircuits(t *testing.T) {
	path := filepath.Join(t.TempDir(), "status.json")
	store, err := Open(path, func() (*schema.Document, error) {
		return &schema.Document{Shards: []schema.Shard{{ShardID: "x:1"}}}, nil
	})
	require.NoError(t, err)

	ctx, cancel := context.WithCancel(context.Background())
	cancel() // cancel BEFORE calling Mutate
	var ran bool
	err = store.Mutate(ctx, func(d *schema.Document) error {
		ran = true
		return nil
	})
	assert.False(t, ran, "fn should not run when ctx is cancelled")
	assert.ErrorIs(t, err, context.Canceled)
}

func TestMutateConcurrentIsSafe(t *testing.T) {
	const N = 50
	path := filepath.Join(t.TempDir(), "status.json")
	store, err := Open(path, func() (*schema.Document, error) {
		shards := make([]schema.Shard, N)
		for i := range shards {
			shards[i] = schema.Shard{
				ShardID: fmt.Sprintf("x:%d", i+1),
				Tenant:  "x",
				// PhaseRunning: Pending → Imported is not a legal transition.
				Phase: schema.PhaseRunning,
			}
		}
		return &schema.Document{Shards: shards}, nil
	})
	require.NoError(t, err)

	var wg sync.WaitGroup
	var failures atomic.Int32
	for i := 1; i <= N; i++ {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			id := fmt.Sprintf("x:%d", idx)
			if err := store.Mutate(context.Background(), func(d *schema.Document) error {
				for j := range d.Shards {
					if d.Shards[j].ShardID != id {
						continue
					}
					if err := d.Shards[j].Phase.Transition(schema.PhaseImported); err != nil {
						return fmt.Errorf("shard %s: %w", id, err)
					}
					d.Shards[j].Phase = schema.PhaseImported
					return nil
				}
				return fmt.Errorf("shard %s missing", id)
			}); err != nil {
				t.Logf("goroutine %d: %v", idx, err)
				failures.Add(1)
			}
		}(i)
	}
	wg.Wait()

	assert.Equal(t, int32(0), failures.Load(), "some goroutines failed")

	store.View(func(d *schema.Document) {
		assert.Equal(t, N, d.Summary.Imported)
		assert.Zero(t, d.Summary.Pending)
	})

	reloaded, err := LoadDocument(path)
	assert.NoError(t, err, "reload corrupted")
	assert.Len(t, reloaded.Shards, N, "slice corrupted?")
}

func TestNewInMemoryDoesNotPersist(t *testing.T) {
	path := filepath.Join(t.TempDir(), "ignored.json")
	doc := &schema.Document{Shards: []schema.Shard{{ShardID: "x:1", Phase: schema.PhasePending}}}
	store := NewInMemory(doc)

	require.NoError(t, store.Mutate(context.Background(), func(d *schema.Document) error {
		d.Shards[0].Phase = schema.PhaseImported
		return nil
	}))

	_, err := os.Stat(path)
	assert.ErrorIs(t, err, os.ErrNotExist, "NewInMemory should not write to disk")
	assert.Equal(t, 1, store.Document().Summary.Imported)
}

func TestMutateHoldsLockAcrossFnAndWrite(t *testing.T) {
	store := NewInMemory(&schema.Document{Shards: []schema.Shard{{ShardID: "x:1"}}})
	writeStarted := make(chan struct{})
	releaseWrite := make(chan struct{})
	var writes atomic.Int32
	store.writeFn = func(*schema.Document) error {
		if writes.Add(1) == 1 {
			close(writeStarted)
			<-releaseWrite
		}
		return nil
	}
	ctx := t.Context()
	firstDone := make(chan error, 1)
	go func() {
		firstDone <- store.Mutate(ctx, func(d *schema.Document) error {
			d.Shards[0].Tenant = "mutated"
			return nil
		})
	}()

	<-writeStarted
	lockAvailable := store.mu.TryLock()
	if lockAvailable {
		store.mu.Unlock()
	}
	secondEntered := make(chan struct{})
	secondDone := make(chan error, 1)
	go func() {
		secondDone <- store.Mutate(ctx, func(*schema.Document) error {
			close(secondEntered)
			return nil
		})
	}()

	close(releaseWrite)
	require.NoError(t, <-firstDone)
	<-secondEntered
	require.NoError(t, <-secondDone)
	assert.False(t, lockAvailable, "writeFn must run while the Store lock is held")
}