package weed_server

import (
	"context"
	"errors"
	"fmt"
	"path/filepath"
	"reflect"
	"testing"
	"time"

	engine "github.com/seaweedfs/seaweedfs/sw-block/engine/replication"
	rt "github.com/seaweedfs/seaweedfs/sw-block/engine/replication/runtime"
	"github.com/seaweedfs/seaweedfs/weed/storage"
	"github.com/seaweedfs/seaweedfs/weed/storage/blockvol"
	"github.com/seaweedfs/seaweedfs/weed/storage/blockvol/v2bridge"
)

// ============================================================
// Phase 09 P4: Stronger live runtime ownership
//
// Proofs:
//   1. Live-path with real vol: ProcessAssignments → plan → executor
//   2. Serialized replacement: old drained before new starts
//   3. Shutdown drains all tasks
//   4. Rebuild address scoped
//   5. No split ownership under replacement
// ============================================================

func createTestBlockServiceWithVol(t *testing.T) (*BlockService, string) {
	t.Helper()
	dir := t.TempDir()

	// Create a real blockvol.
	volPath := filepath.Join(dir, "vol1.blk")
	vol, err := blockvol.CreateBlockVol(volPath, blockvol.CreateOptions{
		VolumeSize: 1 * 1024 * 1024,
		BlockSize:  4096,
		WALSize:    256 * 1024,
	})
	if err != nil {
		t.Fatalf("CreateBlockVol: %v", err)
	}
	vol.Close()

	// Build BlockService with a real BlockVolumeStore.
	store := storage.NewBlockVolumeStore()
	if _, err := store.AddBlockVolume(volPath, ""); err != nil {
		t.Fatalf("AddBlockVolume: %v", err)
	}

	bs := &BlockService{
		blockStore:     store,
		blockDir:       dir,
		listenAddr:     "127.0.0.1:3260",
		localServerID:  "test-server-1",
		v2Bridge:       v2bridge.NewControlBridge(),
		v2Orchestrator: engine.NewRecoveryOrchestrator(),
	}
	bs.v2Recovery = NewRecoveryManager(bs)

	t.Cleanup(func() {
		bs.v2Recovery.Shutdown()
		store.Close()
	})

	return bs, volPath
}

func createTestBlockServiceWithVolCoreNoRecovery(t *testing.T) (*BlockService, string) {
	t.Helper()
	dir := t.TempDir()

	volPath := filepath.Join(dir, "vol1.blk")
	vol, err := blockvol.CreateBlockVol(volPath, blockvol.CreateOptions{
		VolumeSize: 1 * 1024 * 1024,
		BlockSize:  4096,
		WALSize:    256 * 1024,
	})
	if err != nil {
		t.Fatalf("CreateBlockVol: %v", err)
	}
	vol.Close()

	store := storage.NewBlockVolumeStore()
	if _, err := store.AddBlockVolume(volPath, ""); err != nil {
		t.Fatalf("AddBlockVolume: %v", err)
	}

	bs := &BlockService{
		blockStore:     store,
		blockDir:       dir,
		listenAddr:     "127.0.0.1:3260",
		localServerID:  "test-server-1",
		v2Bridge:       v2bridge.NewControlBridge(),
		v2Orchestrator: engine.NewRecoveryOrchestrator(),
		v2Core:         engine.NewCoreEngine(),
		coreProj:       make(map[string]engine.PublicationProjection),
		replStates:     make(map[string]*volReplState),
	}

	t.Cleanup(func() {
		store.Close()
	})

	return bs, volPath
}

func installTestSession(t *testing.T, bs *BlockService, replicaID string, kind engine.SessionKind) uint64 {
	t.Helper()
	rm := bs.v2Recovery
	if rm == nil {
		rm = NewRecoveryManager(bs)
	}
	if err := rm.installSession(replicaID, kind); err != nil {
		t.Fatalf("install %s session for %s: %v", kind, replicaID, err)
	}
	sender := bs.v2Orchestrator.Registry.Sender(replicaID)
	if sender == nil {
		t.Fatalf("sender %s not found after install", replicaID)
	}
	snap := sender.SessionSnapshot()
	if snap == nil || !snap.Active || snap.Kind != kind {
		t.Fatalf("session=%+v, want active %s", snap, kind)
	}
	return snap.ID
}

type fakeRebuildIO struct {
	achievedLSN uint64
}

func (f fakeRebuildIO) TransferFullBase(committedLSN uint64) (uint64, error) {
	if f.achievedLSN > 0 {
		return f.achievedLSN, nil
	}
	return committedLSN, nil
}

func (f fakeRebuildIO) TransferSnapshot(snapshotLSN uint64) error {
	return nil
}

func (f fakeRebuildIO) StreamWALEntries(startExclusive, endInclusive uint64) (uint64, error) {
	return endInclusive, nil
}

type errorRebuildIO struct {
	err error
}

func (e errorRebuildIO) TransferFullBase(committedLSN uint64) (uint64, error) {
	if e.err == nil {
		return 0, errors.New("rebuild execution failed")
	}
	return 0, e.err
}

func (e errorRebuildIO) TransferSnapshot(snapshotLSN uint64) error {
	if e.err == nil {
		return errors.New("rebuild execution failed")
	}
	return e.err
}

func (e errorRebuildIO) StreamWALEntries(startExclusive, endInclusive uint64) (uint64, error) {
	if e.err == nil {
		return 0, errors.New("rebuild execution failed")
	}
	return 0, e.err
}

type completedAckThenErrorRebuildIO struct {
	rm          *RecoveryManager
	replicaID   string
	achievedLSN uint64
	err         error
}

func (c completedAckThenErrorRebuildIO) TransferFullBase(committedLSN uint64) (uint64, error) {
	if c.rm != nil {
		achieved := c.achievedLSN
		if achieved == 0 {
			achieved = committedLSN
		}
		c.rm.mu.Lock()
		if c.rm.remoteRebuildAchieved == nil {
			c.rm.remoteRebuildAchieved = make(map[string]uint64)
		}
		c.rm.remoteRebuildAchieved[c.replicaID] = achieved
		c.rm.mu.Unlock()
	}
	if c.err == nil {
		return 0, errors.New("post_ack_error")
	}
	return 0, c.err
}

func (c completedAckThenErrorRebuildIO) TransferSnapshot(snapshotLSN uint64) error {
	if c.err == nil {
		return errors.New("post_ack_error")
	}
	return c.err
}

func (c completedAckThenErrorRebuildIO) StreamWALEntries(startExclusive, endInclusive uint64) (uint64, error) {
	if c.err == nil {
		return 0, errors.New("post_ack_error")
	}
	return 0, c.err
}

// --- Live-path with real vol: reaches planning ---

func TestP4_LivePath_RealVol_ReachesPlan(t *testing.T) {
	bs, volPath := createTestBlockServiceWithVol(t)

	// Write data to the vol so recovery has something to plan against.
	if err := bs.blockStore.WithVolume(volPath, func(vol *blockvol.BlockVol) error {
		for i := 0; i < 5; i++ {
			vol.WriteLBA(uint64(i), make([]byte, 4096))
		}
		return nil
	}); err != nil {
		t.Fatalf("write: %v", err)
	}

	// Process assignment through real path.
	bs.ProcessAssignments([]blockvol.BlockVolumeAssignment{
		{
			Path:            volPath,
			Epoch:           1,
			Role:            uint32(blockvol.RolePrimary),
			ReplicaServerID: "vs2",
			ReplicaDataAddr: "10.0.0.2:9333",
			ReplicaCtrlAddr: "10.0.0.2:9334",
		},
	})

	// Give recovery goroutine time to reach planning.
	time.Sleep(200 * time.Millisecond)

	// Verify: sender exists and engine processed it.
	replicaID := volPath + "/vs2"
	s := bs.v2Orchestrator.Registry.Sender(replicaID)
	if s == nil {
		t.Fatal("sender not created")
	}

	// Assert the full chain completed: plan → executor → in_sync.
	events := bs.v2Orchestrator.Log.EventsFor(replicaID)
	required := map[string]bool{
		"plan_catchup":         false,
		"exec_catchup_started": false,
		"exec_completed":       false,
	}
	for _, ev := range events {
		if _, ok := required[ev.Event]; ok {
			required[ev.Event] = true
		}
	}
	for event, found := range required {
		if !found {
			t.Fatalf("missing required event: %s (events=%d)", event, len(events))
		}
	}

	// Assert final sender state.
	if s.State() != engine.StateInSync {
		t.Fatalf("sender state=%s, want in_sync", s.State())
	}

	t.Log("P4 live-path: ProcessAssignments → plan_catchup → exec_catchup_started → exec_completed → in_sync")
}

func TestP16B_RunCatchUp_UpdatesCoreProjectionFromLiveRecovery(t *testing.T) {
	bs, volPath := createTestBlockServiceWithVolCoreNoRecovery(t)

	if err := bs.blockStore.WithVolume(volPath, func(vol *blockvol.BlockVol) error {
		for i := 0; i < 5; i++ {
			if err := vol.WriteLBA(uint64(i), make([]byte, 4096)); err != nil {
				return err
			}
		}
		return nil
	}); err != nil {
		t.Fatalf("write: %v", err)
	}

	bs.ProcessAssignments([]blockvol.BlockVolumeAssignment{
		{
			Path:            volPath,
			Epoch:           1,
			Role:            uint32(blockvol.RolePrimary),
			ReplicaServerID: "vs2",
			ReplicaDataAddr: "10.0.0.2:9333",
			ReplicaCtrlAddr: "10.0.0.2:9334",
		},
	})

	replicaID := volPath + "/vs2"
	sender := bs.v2Orchestrator.Registry.Sender(replicaID)
	if sender == nil || !sender.HasActiveSession() {
		t.Fatal("expected active sender session before catch-up")
	}

	rm := NewRecoveryManager(bs)
	bs.v2Recovery = rm
	rm.OnPendingExecution = func(volumeID string, pending *rt.PendingExecution) {
		if volumeID == volPath && pending != nil && pending.Plan != nil {
			pending.CatchUpIO = fakeCatchUpIO{transferredTo: pending.Plan.CatchUpTarget}
		}
	}
	rm.runCatchUp(context.Background(), replicaID, nil)

	proj, ok := bs.CoreProjection(volPath)
	if !ok {
		t.Fatal("expected cached core projection after live catch-up")
	}
	if proj.Boundary.TargetLSN == 0 {
		t.Fatalf("target_lsn=%d", proj.Boundary.TargetLSN)
	}
	if proj.Boundary.AchievedLSN == 0 {
		t.Fatalf("achieved_lsn=%d", proj.Boundary.AchievedLSN)
	}
	if proj.Boundary.DurableLSN == 0 {
		t.Fatalf("durable_lsn=%d", proj.Boundary.DurableLSN)
	}
	if proj.Recovery.Phase != engine.RecoveryIdle {
		t.Fatalf("recovery_phase=%s", proj.Recovery.Phase)
	}
	replicaSync, ok := proj.ReplicaSync[replicaID]
	if !ok {
		t.Fatalf("missing replica sync for %s", replicaID)
	}
	if replicaSync.AckKind != engine.SyncAckTimedOut {
		t.Fatalf("sync_ack_kind=%s", replicaSync.AckKind)
	}
	if replicaSync.Action != engine.SyncActionCatchUp {
		t.Fatalf("sync_action=%s", replicaSync.Action)
	}
	if got := bs.ExecutedCoreCommands(volPath); len(got) == 0 || got[len(got)-1] != "start_catchup" {
		t.Fatalf("expected start_catchup execution, got %v", got)
	}
	if got := countCommandName(bs.ExecutedCoreCommands(volPath), "start_rebuild"); got != 0 {
		t.Fatalf("start_rebuild count=%d, want 0 on catch-up path", got)
	}
}

func TestP16B_RunCatchUp_UsesUnifiedFactEntryToStartRebuild(t *testing.T) {
	bs, volPath := createTestBlockServiceWithVolCoreNoRecovery(t)

	if err := bs.blockStore.WithVolume(volPath, func(vol *blockvol.BlockVol) error {
		for i := 0; i < 5; i++ {
			if err := vol.WriteLBA(uint64(i), make([]byte, 4096)); err != nil {
				return err
			}
		}
		return vol.ForceFlush()
	}); err != nil {
		t.Fatalf("write+flush: %v", err)
	}

	bs.ProcessAssignments([]blockvol.BlockVolumeAssignment{
		{
			Path:            volPath,
			Epoch:           1,
			Role:            uint32(blockvol.RolePrimary),
			ReplicaServerID: "vs2",
			ReplicaDataAddr: "10.0.0.2:9333",
			ReplicaCtrlAddr: "10.0.0.2:9334",
		},
	})

	replicaID := volPath + "/vs2"
	sender := bs.v2Orchestrator.Registry.Sender(replicaID)
	if sender == nil || !sender.HasActiveSession() {
		t.Fatal("expected active sender session before needs_rebuild planning")
	}

	rm := NewRecoveryManager(bs)
	bs.v2Recovery = rm
	rm.OnPendingExecution = func(volumeID string, pending *rt.PendingExecution) {
		if volumeID != volPath || pending == nil || pending.Plan == nil {
			return
		}
		pending.RebuildIO = fakeRebuildIO{achievedLSN: pending.Plan.RebuildTargetLSN}
	}
	_, _, rebuildPort := bs.ReplicationPorts(volPath)
	rebuildAddr := fmt.Sprintf("127.0.0.1:%d", rebuildPort)
	rm.runCatchUp(context.Background(), replicaID, []blockvol.BlockVolumeAssignment{{Path: volPath, RebuildAddr: rebuildAddr}})

	proj, ok := bs.CoreProjection(volPath)
	if !ok {
		t.Fatal("expected cached core projection after needs_rebuild escalation")
	}
	replicaSync, ok := proj.ReplicaSync[replicaID]
	if !ok {
		t.Fatalf("missing replica sync for %s", replicaID)
	}
	if replicaSync.AckKind != engine.SyncAckTimedOut {
		t.Fatalf("sync_ack_kind=%s", replicaSync.AckKind)
	}
	if replicaSync.Action != engine.SyncActionRebuild {
		t.Fatalf("sync_action=%s", replicaSync.Action)
	}
	if proj.Recovery.Phase != engine.RecoveryIdle {
		t.Fatalf("recovery_phase=%s, want %s", proj.Recovery.Phase, engine.RecoveryIdle)
	}
	if got := countCommandName(bs.ExecutedCoreCommands(volPath), "start_catchup"); got != 0 {
		t.Fatalf("start_catchup count=%d, want 0 on rebuild path", got)
	}
	if got := countCommandName(bs.ExecutedCoreCommands(volPath), "start_rebuild"); got != 1 {
		t.Fatalf("start_rebuild count=%d, want 1", got)
	}
	sender = bs.v2Orchestrator.Registry.Sender(replicaID)
	if sender == nil {
		t.Fatal("sender missing after rebuild start")
	}
	if sender.State() != engine.StateInSync {
		t.Fatalf("sender state=%s, want %s", sender.State(), engine.StateInSync)
	}
}

func TestP16B_OnCatchUpFailed_ReentersFactDecisionForRebuild(t *testing.T) {
	bs, volPath := createTestBlockServiceWithVolCoreNoRecovery(t)

	if err := bs.blockStore.WithVolume(volPath, func(vol *blockvol.BlockVol) error {
		for i := 0; i < 5; i++ {
			if err := vol.WriteLBA(uint64(i), make([]byte, 4096)); err != nil {
				return err
			}
		}
		return vol.ForceFlush()
	}); err != nil {
		t.Fatalf("write+flush: %v", err)
	}

	bs.ProcessAssignments([]blockvol.BlockVolumeAssignment{
		{
			Path:            volPath,
			Epoch:           1,
			Role:            uint32(blockvol.RolePrimary),
			ReplicaServerID: "vs2",
			ReplicaDataAddr: "10.0.0.2:9333",
			ReplicaCtrlAddr: "10.0.0.2:9334",
		},
	})

	replicaID := volPath + "/vs2"
	sender := bs.v2Orchestrator.Registry.Sender(replicaID)
	if sender == nil || !sender.HasActiveSession() {
		t.Fatal("expected active sender session before catch-up failure")
	}

	rm := NewRecoveryManager(bs)
	bs.v2Recovery = rm
	rm.OnPendingExecution = func(volumeID string, pending *rt.PendingExecution) {
		if volumeID != volPath || pending == nil || pending.Plan == nil {
			return
		}
		pending.RebuildIO = fakeRebuildIO{achievedLSN: pending.Plan.RebuildTargetLSN}
	}
	rm.OnCatchUpFailed(volPath, replicaID, "recoverability_lost")

	proj, ok := bs.CoreProjection(volPath)
	if !ok {
		t.Fatal("expected cached core projection after catch-up failure")
	}
	replicaSync, ok := proj.ReplicaSync[replicaID]
	if !ok {
		t.Fatalf("missing replica sync for %s", replicaID)
	}
	if replicaSync.AckKind != engine.SyncAckTimedOut {
		t.Fatalf("sync_ack_kind=%s", replicaSync.AckKind)
	}
	if replicaSync.Action != engine.SyncActionRebuild {
		t.Fatalf("sync_action=%s", replicaSync.Action)
	}
	if replicaSync.Reason == "" {
		t.Fatal("expected final sync reason after fresh re-decision")
	}
	if proj.Recovery.Phase != engine.RecoveryIdle {
		t.Fatalf("recovery_phase=%s, want %s", proj.Recovery.Phase, engine.RecoveryIdle)
	}
	sender = bs.v2Orchestrator.Registry.Sender(replicaID)
	if sender == nil {
		t.Fatal("sender missing after failure-driven rebuild")
	}
	if sender.State() != engine.StateInSync {
		t.Fatalf("sender state=%s, want %s", sender.State(), engine.StateInSync)
	}
	if got := countCommandName(bs.ExecutedCoreCommands(volPath), "start_rebuild"); got != 1 {
		t.Fatalf("start_rebuild count=%d, want 1", got)
	}
}

func TestP16B_RunRebuild_UsesCoreStartRebuildCommandOnLivePath(t *testing.T) {
	bs, volPath := createTestBlockServiceWithVolCoreNoRecovery(t)

	if err := bs.blockStore.WithVolume(volPath, func(vol *blockvol.BlockVol) error {
		for i := 0; i < 5; i++ {
			if err := vol.WriteLBA(uint64(i), make([]byte, 4096)); err != nil {
				return err
			}
		}
		return vol.ForceFlush()
	}); err != nil {
		t.Fatalf("write+flush: %v", err)
	}

	bs.ProcessAssignments([]blockvol.BlockVolumeAssignment{{
		Path:            volPath,
		Epoch:           1,
		Role:            uint32(blockvol.RolePrimary),
		ReplicaServerID: "vs2",
		ReplicaDataAddr: "10.0.0.2:9333",
		ReplicaCtrlAddr: "10.0.0.2:9334",
	}})

	replicaID := volPath + "/vs2"
	installTestSession(t, bs, replicaID, engine.SessionRebuild)
	sender := bs.v2Orchestrator.Registry.Sender(replicaID)
	if sender == nil {
		t.Fatal("sender not found")
	}
	snap := sender.SessionSnapshot()
	if snap == nil || snap.Kind != engine.SessionRebuild {
		t.Fatalf("session=%+v", snap)
	}

	rm := NewRecoveryManager(bs)
	bs.v2Recovery = rm
	rm.OnPendingExecution = func(volumeID string, pending *rt.PendingExecution) {
		if volumeID != volPath || pending == nil || pending.Plan == nil {
			return
		}
		pending.RebuildIO = fakeRebuildIO{achievedLSN: pending.Plan.RebuildTargetLSN}
	}
	_, _, rebuildPort := bs.ReplicationPorts(volPath)
	rebuildAddr := fmt.Sprintf("127.0.0.1:%d", rebuildPort)
	rm.runRebuild(context.Background(), replicaID, []blockvol.BlockVolumeAssignment{{Path: volPath, RebuildAddr: rebuildAddr}})

	proj, ok := bs.CoreProjection(volPath)
	if !ok {
		t.Fatal("expected cached core projection after live command-driven rebuild")
	}
	if proj.Recovery.Phase != engine.RecoveryIdle {
		t.Fatalf("recovery_phase=%s", proj.Recovery.Phase)
	}
	if sender.State() != engine.StateInSync {
		t.Fatalf("sender state=%s, want in_sync", sender.State())
	}
	if got := bs.ExecutedCoreCommands(volPath); len(got) == 0 || got[len(got)-1] != "start_rebuild" {
		t.Fatalf("expected start_rebuild execution, got %v", got)
	}
}

func TestP16B_RunRebuild_FailClosedWithoutFreshStartRebuildCommand(t *testing.T) {
	bs, volPath := createTestBlockServiceWithVolCoreNoRecovery(t)

	var targetLSN uint64
	if err := bs.blockStore.WithVolume(volPath, func(vol *blockvol.BlockVol) error {
		for i := 0; i < 5; i++ {
			if err := vol.WriteLBA(uint64(i), make([]byte, 4096)); err != nil {
				return err
			}
		}
		if err := vol.ForceFlush(); err != nil {
			return err
		}
		targetLSN = vol.StatusSnapshot().CommittedLSN
		return nil
	}); err != nil {
		t.Fatalf("write+flush: %v", err)
	}

	bs.ProcessAssignments([]blockvol.BlockVolumeAssignment{{
		Path:            volPath,
		Epoch:           1,
		Role:            uint32(blockvol.RolePrimary),
		ReplicaServerID: "vs2",
		ReplicaDataAddr: "10.0.0.2:9333",
		ReplicaCtrlAddr: "10.0.0.2:9334",
	}})

	replicaID := volPath + "/vs2"
	installTestSession(t, bs, replicaID, engine.SessionRebuild)

	// Prime the core with the same rebuild target before wiring recovery,
	// so the subsequent live run does not emit a fresh start_rebuild command.
	bs.applyCoreEvent(engine.RebuildStarted{ID: volPath, TargetLSN: targetLSN})
	before := bs.ExecutedCoreCommands(volPath)

	rm := NewRecoveryManager(bs)
	bs.v2Recovery = rm
	_, _, rebuildPort := bs.ReplicationPorts(volPath)
	rebuildAddr := fmt.Sprintf("127.0.0.1:%d", rebuildPort)
	rm.runRebuild(context.Background(), replicaID, []blockvol.BlockVolumeAssignment{{Path: volPath, RebuildAddr: rebuildAddr}})

	after := bs.ExecutedCoreCommands(volPath)
	if !reflect.DeepEqual(after, before) {
		t.Fatalf("rebuild should fail closed without fresh start_rebuild command: before=%v after=%v", before, after)
	}
	sender := bs.v2Orchestrator.Registry.Sender(replicaID)
	if sender == nil {
		t.Fatal("sender not found")
	}
	if sender.State() == engine.StateInSync {
		t.Fatalf("sender should not become in_sync without executing start_rebuild, state=%s", sender.State())
	}
}

func TestP16B_RunRebuild_ExecutionFailureAllowsFreshRetry(t *testing.T) {
	bs, volPath := createTestBlockServiceWithVolCoreNoRecovery(t)

	if err := bs.blockStore.WithVolume(volPath, func(vol *blockvol.BlockVol) error {
		for i := 0; i < 5; i++ {
			if err := vol.WriteLBA(uint64(i), make([]byte, 4096)); err != nil {
				return err
			}
		}
		return vol.ForceFlush()
	}); err != nil {
		t.Fatalf("write+flush: %v", err)
	}

	bs.ProcessAssignments([]blockvol.BlockVolumeAssignment{{
		Path:            volPath,
		Epoch:           1,
		Role:            uint32(blockvol.RolePrimary),
		ReplicaServerID: "vs2",
		ReplicaDataAddr: "10.0.0.2:9333",
		ReplicaCtrlAddr: "10.0.0.2:9334",
	}})

	replicaID := volPath + "/vs2"
	installTestSession(t, bs, replicaID, engine.SessionRebuild)

	rm := NewRecoveryManager(bs)
	bs.v2Recovery = rm

	attempts := 0
	rm.OnPendingExecution = func(volumeID string, pending *rt.PendingExecution) {
		if volumeID != volPath || pending == nil || pending.Plan == nil {
			return
		}
		attempts++
		if attempts == 1 {
			pending.RebuildIO = errorRebuildIO{err: errors.New("mid_rebuild_failure")}
			return
		}
		pending.RebuildIO = fakeRebuildIO{achievedLSN: pending.Plan.RebuildTargetLSN}
	}
	_, _, rebuildPort := bs.ReplicationPorts(volPath)
	rebuildAddr := fmt.Sprintf("127.0.0.1:%d", rebuildPort)

	rm.runRebuild(context.Background(), replicaID, []blockvol.BlockVolumeAssignment{{Path: volPath, RebuildAddr: rebuildAddr}})

	sender := bs.v2Orchestrator.Registry.Sender(replicaID)
	if sender == nil {
		t.Fatal("sender missing after failed rebuild")
	}
	if sender.State() == engine.StateInSync {
		t.Fatalf("sender state=%s, want non-in_sync after failed rebuild", sender.State())
	}
	firstProj, ok := bs.CoreProjection(volPath)
	if !ok {
		t.Fatal("expected core projection after failed rebuild")
	}
	if firstProj.Mode.Name != engine.ModeDegraded {
		t.Fatalf("mode=%s, want degraded after failed rebuild", firstProj.Mode.Name)
	}

	installTestSession(t, bs, replicaID, engine.SessionRebuild)
	rm.runRebuild(context.Background(), replicaID, []blockvol.BlockVolumeAssignment{{Path: volPath, RebuildAddr: rebuildAddr}})

	sender = bs.v2Orchestrator.Registry.Sender(replicaID)
	if sender == nil {
		t.Fatal("sender missing after retry rebuild")
	}
	if sender.State() != engine.StateInSync {
		t.Fatalf("sender state=%s, want %s after retry", sender.State(), engine.StateInSync)
	}
	finalProj, ok := bs.CoreProjection(volPath)
	if !ok {
		t.Fatal("expected final core projection after retry rebuild")
	}
	if finalProj.Recovery.Phase != engine.RecoveryIdle {
		t.Fatalf("recovery_phase=%s, want %s", finalProj.Recovery.Phase, engine.RecoveryIdle)
	}
	if attempts != 2 {
		t.Fatalf("attempts=%d, want 2", attempts)
	}
}

func TestP16B_RunRebuild_PostAckErrorStillFailsAndClearsRemoteMarker(t *testing.T) {
	bs, volPath := createTestBlockServiceWithVolCoreNoRecovery(t)

	if err := bs.blockStore.WithVolume(volPath, func(vol *blockvol.BlockVol) error {
		for i := 0; i < 5; i++ {
			if err := vol.WriteLBA(uint64(i), make([]byte, 4096)); err != nil {
				return err
			}
		}
		return vol.ForceFlush()
	}); err != nil {
		t.Fatalf("write+flush: %v", err)
	}

	bs.ProcessAssignments([]blockvol.BlockVolumeAssignment{{
		Path:            volPath,
		Epoch:           1,
		Role:            uint32(blockvol.RolePrimary),
		ReplicaServerID: "vs2",
		ReplicaDataAddr: "10.0.0.2:9333",
		ReplicaCtrlAddr: "10.0.0.2:9334",
	}})

	replicaID := volPath + "/vs2"
	installTestSession(t, bs, replicaID, engine.SessionRebuild)

	rm := NewRecoveryManager(bs)
	bs.v2Recovery = rm
	rm.OnPendingExecution = func(volumeID string, pending *rt.PendingExecution) {
		if volumeID != volPath || pending == nil || pending.Plan == nil {
			return
		}
		pending.RebuildIO = completedAckThenErrorRebuildIO{
			rm:          rm,
			replicaID:   replicaID,
			achievedLSN: pending.Plan.RebuildTargetLSN,
			err:         errors.New("post_ack_error"),
		}
	}
	_, _, rebuildPort := bs.ReplicationPorts(volPath)
	rebuildAddr := fmt.Sprintf("127.0.0.1:%d", rebuildPort)

	rm.runRebuild(context.Background(), replicaID, []blockvol.BlockVolumeAssignment{{Path: volPath, RebuildAddr: rebuildAddr}})

	proj, ok := bs.CoreProjection(volPath)
	if !ok {
		t.Fatal("expected core projection after post-ack error")
	}
	if proj.Mode.Name != engine.ModeDegraded {
		t.Fatalf("mode=%s, want degraded after post-ack error", proj.Mode.Name)
	}
	if proj.Recovery.Reason != "post_ack_error" {
		t.Fatalf("recovery_reason=%q, want post_ack_error", proj.Recovery.Reason)
	}
	rm.mu.Lock()
	_, stillMarked := rm.remoteRebuildAchieved[replicaID]
	rm.mu.Unlock()
	if stillMarked {
		t.Fatal("remote rebuild achieved marker should be cleared on failure")
	}
}

func TestP16B_OnRebuildCompleted_RemotePathEmitsSessionCompleted(t *testing.T) {
	bs, volPath := createTestBlockServiceWithVolCoreNoRecovery(t)

	bs.ProcessAssignments([]blockvol.BlockVolumeAssignment{{
		Path:            volPath,
		Epoch:           1,
		Role:            uint32(blockvol.RolePrimary),
		ReplicaServerID: "vs2",
		ReplicaDataAddr: "10.0.0.2:9333",
		ReplicaCtrlAddr: "10.0.0.2:9334",
	}})

	replicaID := volPath + "/vs2"
	sessionID := installTestSession(t, bs, replicaID, engine.SessionRebuild)
	if sessionID == 0 {
		t.Fatal("expected rebuild session")
	}
	targetLSN := uint64(123)
	bs.applyCoreEvent(engine.RebuildStarted{ID: volPath, ReplicaID: replicaID, TargetLSN: targetLSN})

	rm := NewRecoveryManager(bs)
	bs.v2Recovery = rm
	rm.mu.Lock()
	rm.remoteRebuildAchieved[replicaID] = targetLSN + 7
	rm.mu.Unlock()

	rm.OnRebuildCompleted(volPath, replicaID, &engine.RecoveryPlan{RebuildTargetLSN: targetLSN})

	proj, ok := bs.CoreProjection(volPath)
	if !ok {
		t.Fatal("expected core projection after remote rebuild completion")
	}
	if proj.Recovery.Phase != engine.RecoveryIdle {
		t.Fatalf("recovery_phase=%s, want %s", proj.Recovery.Phase, engine.RecoveryIdle)
	}
	if proj.Boundary.AchievedLSN != targetLSN+7 {
		t.Fatalf("achieved_lsn=%d, want %d", proj.Boundary.AchievedLSN, targetLSN+7)
	}
	rm.mu.Lock()
	_, stillMarked := rm.remoteRebuildAchieved[replicaID]
	rm.mu.Unlock()
	if stillMarked {
		t.Fatal("remote rebuild achieved marker should be consumed on completion")
	}
}

func TestP16B_FactTriggeredRebuildCycle_AutoInstallsRebuildAndReachesInSync(t *testing.T) {
	bs, volPath := createTestBlockServiceWithVolCoreNoRecovery(t)

	if err := bs.blockStore.WithVolume(volPath, func(vol *blockvol.BlockVol) error {
		for i := 0; i < 5; i++ {
			if err := vol.WriteLBA(uint64(i), make([]byte, 4096)); err != nil {
				return err
			}
		}
		return vol.ForceFlush()
	}); err != nil {
		t.Fatalf("write+flush: %v", err)
	}

	bs.ProcessAssignments([]blockvol.BlockVolumeAssignment{{
		Path:            volPath,
		Epoch:           1,
		Role:            uint32(blockvol.RolePrimary),
		ReplicaServerID: "vs2",
		ReplicaDataAddr: "10.0.0.2:9333",
		ReplicaCtrlAddr: "10.0.0.2:9334",
	}})

	replicaID := volPath + "/vs2"
	sender := bs.v2Orchestrator.Registry.Sender(replicaID)
	if sender == nil || !sender.HasActiveSession() {
		t.Fatal("expected active sender session before timeout-driven rebuild")
	}

	rm := NewRecoveryManager(bs)
	bs.v2Recovery = rm
	rm.OnPendingExecution = func(volumeID string, pending *rt.PendingExecution) {
		if volumeID != volPath || pending == nil || pending.Plan == nil {
			return
		}
		pending.RebuildIO = fakeRebuildIO{achievedLSN: pending.Plan.RebuildTargetLSN}
	}
	_, _, rebuildPort := bs.ReplicationPorts(volPath)
	rebuildAddr := fmt.Sprintf("127.0.0.1:%d", rebuildPort)
	rm.runCatchUp(context.Background(), replicaID, []blockvol.BlockVolumeAssignment{{Path: volPath, RebuildAddr: rebuildAddr}})

	proj, ok := bs.CoreProjection(volPath)
	if !ok {
		t.Fatal("expected core projection after rebuild completion")
	}
	replicaSync, ok := proj.ReplicaSync[replicaID]
	if !ok {
		t.Fatalf("missing replica sync for %s", replicaID)
	}
	if replicaSync.AckKind != engine.SyncAckTimedOut {
		t.Fatalf("sync_ack_kind=%s, want %s", replicaSync.AckKind, engine.SyncAckTimedOut)
	}
	if replicaSync.Action != engine.SyncActionRebuild {
		t.Fatalf("sync_action=%s, want %s", replicaSync.Action, engine.SyncActionRebuild)
	}
	if proj.Recovery.Phase != engine.RecoveryIdle {
		t.Fatalf("recovery_phase=%s, want %s", proj.Recovery.Phase, engine.RecoveryIdle)
	}
	sender = bs.v2Orchestrator.Registry.Sender(replicaID)
	if sender == nil {
		t.Fatal("sender missing after auto rebuild cycle")
	}
	if sender.State() != engine.StateInSync {
		t.Fatalf("sender state=%s, want %s", sender.State(), engine.StateInSync)
	}
	state, ok := bs.ProtocolExecutionState(volPath)
	if !ok {
		t.Fatal("expected protocol execution state after rebuild completion")
	}
	replicaExec, ok := state.Replicas[replicaID]
	if !ok {
		t.Fatalf("missing protocol execution state for %s", replicaID)
	}
	if replicaExec.SessionActive {
		t.Fatal("SessionActive=true, want false after rebuild completion")
	}
	if !replicaExec.LiveEligible {
		t.Fatal("LiveEligible=false, want true after rebuild completion")
	}
	if got := countCommandName(bs.ExecutedCoreCommands(volPath), "start_rebuild"); got != 1 {
		t.Fatalf("start_rebuild count=%d, want 1", got)
	}
}

// --- Serialized replacement: old drained before new starts ---

func TestP4_SerializedReplacement_DrainsBeforeStart(t *testing.T) {
	bs, volPath := createTestBlockServiceWithVol(t)
	rm := bs.v2Recovery

	// Write data so recovery has work to do.
	if err := bs.blockStore.WithVolume(volPath, func(vol *blockvol.BlockVol) error {
		for i := 0; i < 10; i++ {
			vol.WriteLBA(uint64(i), make([]byte, 4096))
		}
		return nil
	}); err != nil {
		t.Fatalf("write: %v", err)
	}

	// Hook: block the first recovery goroutine so it is DEFINITELY still
	// alive when the supersede arrives. Release it via channel.
	holdFirst := make(chan struct{})
	firstReached := make(chan struct{}, 1)
	callCount := 0
	rm.OnBeforeExecute = func(replicaID string) {
		callCount++
		if callCount == 1 {
			// First call: signal that we're alive, then block.
			firstReached <- struct{}{}
			<-holdFirst // block until test releases
		}
		// Second call (replacement): proceed immediately.
	}

	replicaID := volPath + "/vs2"

	// Epoch 1: start recovery goroutine (will block at OnBeforeExecute).
	bs.ProcessAssignments([]blockvol.BlockVolumeAssignment{
		{Path: volPath, Epoch: 1, Role: uint32(blockvol.RolePrimary),
			ReplicaServerID: "vs2", ReplicaDataAddr: "10.0.0.2:9333", ReplicaCtrlAddr: "10.0.0.2:9334"},
	})

	// Wait for the first goroutine to reach the hook (still alive).
	select {
	case <-firstReached:
		t.Log("first recovery goroutine is alive and blocked at OnBeforeExecute")
	case <-time.After(5 * time.Second):
		t.Fatal("first goroutine did not reach OnBeforeExecute")
	}

	// Capture the old task's done channel WHILE it is still running.
	rm.mu.Lock()
	oldTask := rm.tasks[replicaID]
	rm.mu.Unlock()
	if oldTask == nil {
		t.Fatal("old task must exist while goroutine is blocked")
	}
	oldDone := oldTask.done

	// Verify: old done channel is NOT closed yet.
	select {
	case <-oldDone:
		t.Fatal("old task done should NOT be closed yet")
	default:
		t.Log("confirmed: old task still running (done channel open)")
	}

	// Epoch 2: supersede. cancelAndDrain will cancel the old context and
	// block on oldDone. We release the hold from another goroutine so
	// cancelAndDrain can complete.
	go func() {
		time.Sleep(50 * time.Millisecond)
		close(holdFirst) // release the blocked first goroutine
	}()

	// This call blocks inside cancelAndDrain until the old goroutine exits.
	bs.ProcessAssignments([]blockvol.BlockVolumeAssignment{
		{Path: volPath, Epoch: 2, Role: uint32(blockvol.RolePrimary),
			ReplicaServerID: "vs2", ReplicaDataAddr: "10.0.0.2:9333", ReplicaCtrlAddr: "10.0.0.2:9334"},
	})

	// After ProcessAssignments returns, the old goroutine MUST be drained.
	select {
	case <-oldDone:
		t.Log("confirmed: old task drained (done channel closed) before replacement started")
	default:
		t.Fatal("old task done channel still open after ProcessAssignments returned")
	}

	time.Sleep(200 * time.Millisecond)

	// At most 1 active task.
	count := rm.ActiveTaskCount()
	if count > 1 {
		t.Fatalf("overlap: %d tasks after replacement", count)
	}

	t.Log("P4 serialized: old goroutine blocked → supersede → drain waited → old exited → replacement started")
}

// --- Shutdown drains all ---

func TestP4_ShutdownDrain(t *testing.T) {
	bs, volPath := createTestBlockServiceWithVol(t)
	rm := bs.v2Recovery

	// Hook: block the goroutine so it's alive when shutdown arrives.
	holdTask := make(chan struct{})
	taskReached := make(chan struct{}, 1)
	rm.OnBeforeExecute = func(replicaID string) {
		taskReached <- struct{}{}
		<-holdTask
	}

	// Write data so recovery has work.
	if err := bs.blockStore.WithVolume(volPath, func(vol *blockvol.BlockVol) error {
		for i := 0; i < 5; i++ {
			vol.WriteLBA(uint64(i), make([]byte, 4096))
		}
		return nil
	}); err != nil {
		t.Fatalf("write: %v", err)
	}

	bs.ProcessAssignments([]blockvol.BlockVolumeAssignment{
		{Path: volPath, Epoch: 1, Role: uint32(blockvol.RolePrimary),
			ReplicaServerID: "vs2", ReplicaDataAddr: "10.0.0.2:9333", ReplicaCtrlAddr: "10.0.0.2:9334"},
	})

	// Wait for task to be alive.
	select {
	case <-taskReached:
		t.Log("task is alive and blocked before shutdown")
	case <-time.After(5 * time.Second):
		t.Fatal("task did not reach hook")
	}

	// Assert: a live task exists.
	if rm.ActiveTaskCount() == 0 {
		t.Fatal("expected live task before shutdown")
	}

	// Release the blocked goroutine from another goroutine so Shutdown can drain.
	go func() {
		time.Sleep(50 * time.Millisecond)
		close(holdTask)
	}()

	done := make(chan bool, 1)
	go func() {
		rm.Shutdown()
		done <- true
	}()

	select {
	case <-done:
	case <-time.After(5 * time.Second):
		t.Fatal("Shutdown did not complete within 5 seconds")
	}

	if rm.ActiveTaskCount() != 0 {
		t.Fatalf("expected 0 active tasks, got %d", rm.ActiveTaskCount())
	}

	t.Log("P4 shutdown: live task existed → shutdown drained it → 0 active")
}

// --- Rebuild address scoped ---

// ============================================================
// Probe → rebuild path: component tests for StartRebuildFromProbe
// and the full handleReplicaProbeResult → recovery chain.
//
// These tests close the gap where bugs were only caught on hardware.
// ============================================================

// TestStartRebuildFromProbe_InstallsSessionAndReachesPlan verifies the new
// primary-direct rebuild entry point: installSession(SessionRebuild) + startTask.
// The goroutine must find the session, derive the rebuild address from
// ReplicationPorts (no assignments), and reach rebuild planning.
func TestStartRebuildFromProbe_InstallsSessionAndReachesPlan(t *testing.T) {
	bs, volPath := createTestBlockServiceWithVolCoreNoRecovery(t)

	// Write data so the volume has content for rebuild planning.
	if err := bs.blockStore.WithVolume(volPath, func(vol *blockvol.BlockVol) error {
		for i := 0; i < 5; i++ {
			if err := vol.WriteLBA(uint64(i), make([]byte, 4096)); err != nil {
				return err
			}
		}
		return vol.ForceFlush()
	}); err != nil {
		t.Fatalf("write+flush: %v", err)
	}

	// Set up primary assignment with a replica.
	bs.ProcessAssignments([]blockvol.BlockVolumeAssignment{{
		Path:            volPath,
		Epoch:           1,
		Role:            uint32(blockvol.RolePrimary),
		ReplicaServerID: "vs2",
		ReplicaDataAddr: "10.0.0.2:9333",
		ReplicaCtrlAddr: "10.0.0.2:9334",
	}})

	replicaID := volPath + "/vs2"
	sender := bs.v2Orchestrator.Registry.Sender(replicaID)
	if sender == nil {
		t.Fatal("expected sender for replica after assignment")
	}

	// Fire NeedsRebuildObserved (what handleReplicaProbeResult does before calling us).
	bs.applyCoreEvent(engine.NeedsRebuildObserved{
		ID:        volPath,
		ReplicaID: replicaID,
		Reason:    "gap_exceeds_retained_wal",
	})

	// Create recovery manager with fake rebuild IO so planning completes.
	rm := NewRecoveryManager(bs)
	bs.v2Recovery = rm
	t.Cleanup(func() { rm.Shutdown() })

	rebuildReached := make(chan string, 1)
	rm.OnPendingExecution = func(volumeID string, pending *rt.PendingExecution) {
		if pending != nil && pending.Plan != nil {
			pending.RebuildIO = fakeRebuildIO{achievedLSN: pending.Plan.RebuildTargetLSN}
			select {
			case rebuildReached <- volumeID:
			default:
			}
		}
	}

	// Call StartRebuildFromProbe — the method under test.
	rm.StartRebuildFromProbe(replicaID)

	// 1. Session must be installed.
	sender = bs.v2Orchestrator.Registry.Sender(replicaID)
	if sender == nil {
		t.Fatal("sender missing after StartRebuildFromProbe")
	}
	snap := sender.SessionSnapshot()
	if snap == nil || !snap.Active || snap.Kind != engine.SessionRebuild {
		t.Fatalf("session=%+v, want active rebuild session", snap)
	}

	// 2. Recovery goroutine must reach rebuild planning.
	select {
	case vol := <-rebuildReached:
		if vol != volPath {
			t.Fatalf("rebuild reached for %s, want %s", vol, volPath)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("rebuild planning not reached within 5s")
	}
}

// TestStartRebuildFromProbe_DeriveRebuildAddr_NilAssignments verifies that
// deriveRebuildAddr produces a valid address when assignments are nil,
// using the ReplicationPorts fallback.
func TestStartRebuildFromProbe_DeriveRebuildAddr_NilAssignments(t *testing.T) {
	bs, volPath := createTestBlockServiceWithVol(t)
	rm := bs.v2Recovery

	replicaID := volPath + "/vs2"
	addr := rm.deriveRebuildAddr(replicaID, nil)
	if addr == "" {
		t.Fatal("deriveRebuildAddr with nil assignments returned empty — fallback broken")
	}

	// The fallback should use the same port as ReplicationPorts.
	_, _, rebuildPort := bs.ReplicationPorts(volPath)
	expected := fmt.Sprintf("127.0.0.1:%d", rebuildPort)
	if addr != expected {
		t.Fatalf("deriveRebuildAddr=%s, want %s", addr, expected)
	}
}

// TestHandleReplicaProbeResult_RebuildRequired_FullPath exercises the full
// probe → NeedsRebuildObserved → StartRebuildFromProbe → recovery chain.
// This is the integration seam that was previously untested and required
// hardware runs to find bugs.
func TestHandleReplicaProbeResult_RebuildRequired_FullPath(t *testing.T) {
	bs, volPath := createTestBlockServiceWithVolCoreNoRecovery(t)

	if err := bs.blockStore.WithVolume(volPath, func(vol *blockvol.BlockVol) error {
		for i := 0; i < 5; i++ {
			if err := vol.WriteLBA(uint64(i), make([]byte, 4096)); err != nil {
				return err
			}
		}
		return vol.ForceFlush()
	}); err != nil {
		t.Fatalf("write+flush: %v", err)
	}

	bs.ProcessAssignments([]blockvol.BlockVolumeAssignment{{
		Path:            volPath,
		Epoch:           1,
		Role:            uint32(blockvol.RolePrimary),
		ReplicaServerID: "vs2",
		ReplicaDataAddr: "10.0.0.2:9333",
		ReplicaCtrlAddr: "10.0.0.2:9334",
	}})

	replicaID := volPath + "/vs2"

	rm := NewRecoveryManager(bs)
	bs.v2Recovery = rm
	t.Cleanup(func() { rm.Shutdown() })

	rebuildReached := make(chan string, 1)
	rm.OnPendingExecution = func(volumeID string, pending *rt.PendingExecution) {
		if pending != nil && pending.Plan != nil {
			pending.RebuildIO = fakeRebuildIO{achievedLSN: pending.Plan.RebuildTargetLSN}
			select {
			case rebuildReached <- volumeID:
			default:
			}
		}
	}

	// Simulate what the shipper's ProbeReconnect returns.
	bs.handleReplicaProbeResult(volPath, blockvol.ReplicaProbeResult{
		ReplicaID:         "vs2", // shipper format, not engine format
		Outcome:           blockvol.ProbeRebuildRequired,
		ReplicaFlushedLSN: 0,
	})

	// 1. Engine should record NeedsRebuild.
	proj, ok := bs.CoreProjection(volPath)
	if !ok {
		t.Fatal("expected core projection")
	}
	if proj.Recovery.Phase != engine.RecoveryNeedsRebuild {
		t.Fatalf("recovery.phase=%s, want %s", proj.Recovery.Phase, engine.RecoveryNeedsRebuild)
	}

	// 2. Rebuild session should be installed.
	sender := bs.v2Orchestrator.Registry.Sender(replicaID)
	if sender == nil {
		t.Fatal("sender missing after probe result")
	}
	snap := sender.SessionSnapshot()
	if snap == nil || !snap.Active || snap.Kind != engine.SessionRebuild {
		t.Fatalf("session=%+v, want active rebuild", snap)
	}

	// 3. Recovery goroutine should reach planning.
	select {
	case vol := <-rebuildReached:
		if vol != volPath {
			t.Fatalf("rebuild for %s, want %s", vol, volPath)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("rebuild planning not reached within 5s")
	}
}

func TestP4_RebuildAddrScoped(t *testing.T) {
	bs, _ := createTestBlockServiceWithVol(t)
	rm := bs.v2Recovery

	assignments := []blockvol.BlockVolumeAssignment{
		{Path: "/data/vol1.blk", RebuildAddr: "10.0.0.1:5000"},
		{Path: "/data/vol2.blk", RebuildAddr: "10.0.0.2:5000"},
	}

	if addr := rm.deriveRebuildAddr("/data/vol1.blk/vs2", assignments); addr != "10.0.0.1:5000" {
		t.Fatalf("vol1 addr=%s", addr)
	}
	if addr := rm.deriveRebuildAddr("/data/vol2.blk/vs3", assignments); addr != "10.0.0.2:5000" {
		t.Fatalf("vol2 addr=%s", addr)
	}
	if addr := rm.deriveRebuildAddr("/data/vol3.blk/vs4", assignments); addr != "" {
		t.Fatalf("vol3 addr=%s", addr)
	}
}