package actor

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"google.golang.org/protobuf/proto"
)

type testState struct {
	Value int `json:"value"`
}

type testGrain struct {
	id       uint64
	accessed time.Time
}

func (g *testGrain) GetId() uint64                      { return g.id }
func (g *testGrain) GetLastAccess() time.Time           { return g.accessed }
func (g *testGrain) GetLastChange() time.Time           { return g.accessed }
func (g *testGrain) GetCurrentState() (*testState, error) {
	return &testState{Value: int(g.id)}, nil
}

type mockHost struct {
	name        string
	healthy     bool
	closed      atomic.Bool
	closeCalls  atomic.Int32
	actorIDs    []uint64
	negotiateFn func([]string) ([]string, error)
	mu          sync.Mutex
	ownership   []uint64
	expiry      []uint64
}

func newMockHost(name string) *mockHost {
	return &mockHost{name: name, healthy: true}
}

func (m *mockHost) AnnounceExpiry(ids []uint64, _ []int64) {
	m.mu.Lock()
	m.expiry = append(m.expiry, ids...)
	m.mu.Unlock()
}

func (m *mockHost) Negotiate(otherHosts []string) ([]string, error) {
	if m.negotiateFn != nil {
		return m.negotiateFn(otherHosts)
	}
	return otherHosts, nil
}

func (m *mockHost) Name() string { return m.name }

func (m *mockHost) Proxy(_ uint64, _ http.ResponseWriter, _ *http.Request, _ io.Reader) (bool, error) {
	return false, fmt.Errorf("not implemented")
}

func (m *mockHost) Apply(_ context.Context, id uint64, _ ...proto.Message) (*MutationResult[testState], error) {
	return &MutationResult[testState]{Result: testState{Value: int(id)}}, nil
}

func (m *mockHost) Get(_ context.Context, id uint64) (*testState, error) {
	return &testState{Value: int(id)}, nil
}

func (m *mockHost) GetActorIds() []uint64 { return m.actorIDs }

func (m *mockHost) Close() error {
	m.closeCalls.Add(1)
	m.closed.Store(true)
	return nil
}

func (m *mockHost) Ping() bool { return m.healthy }

func (m *mockHost) IsHealthy() bool { return m.healthy }

func (m *mockHost) AnnounceOwnership(_ string, ids []uint64, _ []int64) {
	m.mu.Lock()
	m.ownership = append(m.ownership, ids...)
	m.mu.Unlock()
}

func newTestPool(t *testing.T, spawnHost func(string) (Host[testState], error)) *SimpleGrainPool[testState] {
	t.Helper()
	reg := NewMutationRegistry()
	pool, err := NewSimpleGrainPool(GrainPoolConfig[testState]{
		Hostname: "local",
		Spawn: func(_ context.Context, id uint64) (Grain[testState], error) {
			return &testGrain{id: id, accessed: time.Now()}, nil
		},
		SpawnHost:        spawnHost,
		TTL:              time.Hour,
		PoolSize:         100,
		MutationRegistry: reg,
	})
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(pool.Close)
	return pool
}

func TestSimpleGrainPoolGetAndApply(t *testing.T) {
	pool := newTestPool(t, func(string) (Host[testState], error) {
		return nil, fmt.Errorf("no remotes")
	})

	got, err := pool.Get(context.Background(), 42)
	if err != nil {
		t.Fatal(err)
	}
	if got.Value != 42 {
		t.Fatalf("Get value = %d, want 42", got.Value)
	}

	ids := pool.GetLocalIds()
	if len(ids) != 1 || ids[0] != 42 {
		t.Fatalf("GetLocalIds = %v, want [42]", ids)
	}
}

func TestSimpleGrainPoolAddRemoteDedup(t *testing.T) {
	var spawnCalls atomic.Int32
	host := newMockHost("peer-a")

	pool := newTestPool(t, func(name string) (Host[testState], error) {
		spawnCalls.Add(1)
		if name != host.name {
			t.Fatalf("spawn host %q, want %q", name, host.name)
		}
		return host, nil
	})

	const workers = 8
	var wg sync.WaitGroup
	wg.Add(workers)
	for range workers {
		go func() {
			defer wg.Done()
			if _, err := pool.AddRemote(host.name); err != nil {
				t.Error(err)
			}
		}()
	}
	wg.Wait()

	if pool.RemoteCount() != 1 {
		t.Fatalf("RemoteCount = %d, want 1", pool.RemoteCount())
	}
	if spawnCalls.Load() != 1 {
		t.Fatalf("spawn calls = %d, want 1", spawnCalls.Load())
	}
	if host.closeCalls.Load() != 0 {
		t.Fatalf("duplicate remote closed = %d, want 0", host.closeCalls.Load())
	}
}

func TestSimpleGrainPoolAddRemoteRaceClosesDuplicate(t *testing.T) {
	start := make(chan struct{})
	var spawnCalls atomic.Int32
	var spawned []*mockHost

	pool := newTestPool(t, func(name string) (Host[testState], error) {
		spawnCalls.Add(1)
		h := newMockHost(name)
		spawned = append(spawned, h)
		if spawnCalls.Load() == 1 {
			<-start
		}
		return h, nil
	})

	go func() {
		if _, err := pool.AddRemote("peer-b"); err != nil {
			t.Error(err)
		}
	}()
	time.Sleep(20 * time.Millisecond)
	close(start)
	if _, err := pool.AddRemote("peer-b"); err != nil {
		t.Fatal(err)
	}

	if pool.RemoteCount() != 1 {
		t.Fatalf("RemoteCount = %d, want 1", pool.RemoteCount())
	}
	if spawnCalls.Load() != 2 {
		t.Fatalf("spawn calls = %d, want 2", spawnCalls.Load())
	}

	closed := 0
	deadline := time.After(2 * time.Second)
	for closed < 1 {
		closed = 0
		for _, h := range spawned {
			if h.closeCalls.Load() > 0 {
				closed++
			}
		}
		if closed >= 1 {
			break
		}
		select {
		case <-deadline:
			t.Fatal("timed out waiting for duplicate remote host Close")
		default:
			time.Sleep(10 * time.Millisecond)
		}
	}
}

func TestSimpleGrainPoolRemoveHost(t *testing.T) {
	host := newMockHost("peer-c")
	pool := newTestPool(t, func(string) (Host[testState], error) { return host, nil })

	if _, err := pool.AddRemote(host.name); err != nil {
		t.Fatal(err)
	}
	pool.remoteMu.Lock()
	pool.remoteOwners[99] = host
	pool.remoteMu.Unlock()

	pool.RemoveHost(host.name)

	if pool.RemoteCount() != 0 {
		t.Fatalf("RemoteCount = %d, want 0", pool.RemoteCount())
	}
	if owner, ok := pool.OwnerHost(99); ok || owner != nil {
		t.Fatalf("OwnerHost still mapped: %#v ok=%v", owner, ok)
	}
	if host.closeCalls.Load() != 1 {
		t.Fatalf("Close calls = %d, want 1", host.closeCalls.Load())
	}
}

func TestSimpleGrainPoolHandleOwnershipChange(t *testing.T) {
	host := newMockHost("peer-d")
	pool := newTestPool(t, func(string) (Host[testState], error) { return host, nil })

	// First-spawn-wins arbitration: a remote broadcast wins only when its
	// lastChange stamp is strictly earlier than the local grain's stamp
	// (or the local grain is absent). Capture the local spawn time once
	// and pass a stamp 1ns earlier to the remote so the test continues to
	// assert "external owner takes over given an earlier remote spawn"
	// under the new arbitration rule.
	spawnTime := time.Now()
	pool.localMu.Lock()
	pool.grains[7] = &testGrain{id: 7, accessed: spawnTime}
	pool.localMu.Unlock()

	remoteStamp := spawnTime.UnixNano() - 1
	if err := pool.HandleOwnershipChange(host.name, []uint64{7}, []int64{remoteStamp}); err != nil {
		t.Fatal(err)
	}

	pool.localMu.RLock()
	_, local := pool.grains[7]
	pool.localMu.RUnlock()
	if local {
		t.Fatal("local grain should be removed after ownership change")
	}

	owner, ok := pool.OwnerHost(7)
	if !ok || owner.Name() != host.name {
		t.Fatalf("OwnerHost = (%v, %v), want (%q, true)", owner, ok, host.name)
	}
}

// TestSimpleGrainPoolHandleOwnershipChangeArbitration verifies the
// first-spawn-wins arbitration rules for HandleOwnershipChange:
//
//   - Remote with earlier lastChange than local  → accept remote.
//   - Remote with later lastChange than local   → keep local.
//   - Equal lastChange                           → lower hostname wins.
//   - Remote stamp == -1 (legacy path)           → accept remote (pre-arbitration).
//   - Local grain absent                        → accept remote.
func TestSimpleGrainPoolHandleOwnershipChangeArbitration(t *testing.T) {
	host := newMockHost("peer-arbitrate")
	// Each table case constructs a fresh pool inside t.Run — we don't
	// share state across cases because per-case hostname differs
	// (hostname is the tie-break oracle on equal lastChange stamps).

	now := time.Now().UnixNano()

	cases := []struct {
		name          string
		name_local    string // pool hostname for ordering check
		id            uint64
		localStamp    int64  // UnixNano of the local grain; -1 = no local grain
		remoteStamp   int64  // UnixNano stamped on the broadcast
		expectLocal   bool   // true → local grain still resident
		expectRemote  bool   // true → remoteOwners[id] = host
	}{
		{
			name:        "remote older stamp wins",
			name_local:  "zzz", // pool above "peer-arbitrate" lex
			id:          10,
			localStamp:  now + 100,
			remoteStamp: now,
			expectLocal: false,
			expectRemote: true,
		},
		{
			name:        "local earlier stamp keeps",
			name_local:  "aaa", // pool below "peer-arbitrate" lex but stamp decides
			id:          11,
			localStamp:  now,
			remoteStamp: now + 100,
			expectLocal: true,
			expectRemote: false,
		},
		{
			name:        "equal stamp + lower local hostname keeps",
			name_local:  "a-peer", // below host
			id:          12,
			localStamp:  now,
			remoteStamp: now,
			expectLocal: true,
			expectRemote: false,
		},
		{
			name:        "equal stamp + higher local hostname defers",
			name_local:  "zzz-peer", // above host
			id:          13,
			localStamp:  now,
			remoteStamp: now,
			expectLocal: false,
			expectRemote: true,
		},
		{
			name:        "legacy (no remote stamp) accepts remote",
			name_local:  "zzz", // we'd otherwise keep local
			id:          14,
			localStamp:  now + 100,
			remoteStamp: -1,
			expectLocal: false,
			expectRemote: true,
		},
		{
			name:        "no local grain accepts remote",
			name_local:  "aaa",
			id:          15,
			localStamp:  -1,
			remoteStamp: now,
			expectLocal: false,
			expectRemote: true,
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			// Fresh pool per case so hostname differences don't leak.
			p := newTestPool(t, func(string) (Host[testState], error) { return host, nil })
			p.hostname = tc.name_local

			p.localMu.Lock()
			if tc.localStamp != -1 {
				p.grains[tc.id] = &testGrain{id: tc.id, accessed: time.Unix(0, tc.localStamp)}
			}
			p.localMu.Unlock()

			if err := p.HandleOwnershipChange(host.name, []uint64{tc.id}, []int64{tc.remoteStamp}); err != nil {
				t.Fatal(err)
			}

			p.localMu.RLock()
			_, hasLocal := p.grains[tc.id]
			p.localMu.RUnlock()
			if hasLocal != tc.expectLocal {
				t.Fatalf("local grain presence = %v, want %v", hasLocal, tc.expectLocal)
			}

			owner, ok := p.OwnerHost(tc.id)
			if ok != tc.expectRemote {
				t.Fatalf("OwnerHost ok = %v, want %v", ok, tc.expectRemote)
			}
			if tc.expectRemote && (owner == nil || owner.Name() != host.name) {
				t.Fatalf("Owner = %v, want host %q", owner, host.name)
			}
		})
	}
}

func TestSimpleGrainPoolPurgeEvictsStaleGrains(t *testing.T) {
	pool := newTestPool(t, func(string) (Host[testState], error) {
		return nil, fmt.Errorf("no remotes")
	})
	pool.ttl = time.Minute

	stale := &testGrain{id: 1, accessed: time.Now().Add(-2 * time.Minute)}
	fresh := &testGrain{id: 2, accessed: time.Now()}

	pool.localMu.Lock()
	pool.grains[1] = stale
	pool.grains[2] = fresh
	pool.localMu.Unlock()

	pool.purge()

	pool.localMu.RLock()
	defer pool.localMu.RUnlock()
	if _, ok := pool.grains[1]; ok {
		t.Fatal("stale grain should be purged")
	}
	if _, ok := pool.grains[2]; !ok {
		t.Fatal("fresh grain should remain")
	}
}

func TestSimpleGrainPoolBroadcastOwnership(t *testing.T) {
	host := newMockHost("peer-e")
	pool := newTestPool(t, func(string) (Host[testState], error) { return host, nil })
	if _, err := pool.AddRemote(host.name); err != nil {
		t.Fatal(err)
	}

	pool.broadcastOwnership([]uint64{11, 12})

	host.mu.Lock()
	defer host.mu.Unlock()
	if len(host.ownership) != 2 {
		t.Fatalf("ownership announces = %v, want 2 ids", host.ownership)
	}
}