package agent

import (
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/consts"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

type reclaimProjectionRaceStore struct {
	mu           sync.Mutex
	records      map[string]DeliveryOwnershipRecord
	freshEntered chan struct{}
	releaseFresh chan struct{}
}

func (s *reclaimProjectionRaceStore) LoadDeliveryOwnership() ([]DeliveryOwnershipRecord, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	result := make([]DeliveryOwnershipRecord, 0, len(s.records))
	for _, record := range s.records {
		result = append(result, record)
	}
	return result, nil
}

func (s *reclaimProjectionRaceStore) PutDeliveryOwnership(record DeliveryOwnershipRecord) ([]DeliveryOwnershipRecord, error) {
	s.mu.Lock()
	if record.MessageID() != "fresh" {
		s.records[record.MessageID()] = record
		s.mu.Unlock()
		return nil, nil
	}
	failed := s.records["old"]
	failed.State = DeliveryOwnershipTerminal
	failed.InflightReason = ""
	failed.Outcome = DeliveryOutcomeFailed
	failed.FailureCode = "ownership_backlog_reclaimed"
	s.records["old"] = failed
	s.records["fresh"] = record
	s.mu.Unlock()

	close(s.freshEntered)
	<-s.releaseFresh
	return []DeliveryOwnershipRecord{failed}, nil
}

func (s *reclaimProjectionRaceStore) DeleteDeliveryOwnership(messageID string) error {
	s.mu.Lock()
	delete(s.records, messageID)
	s.mu.Unlock()
	return nil
}

func TestReviewReclaimProjectionDoesNotOverwriteConcurrentAckedUpgrade(t *testing.T) {
	createdAt := time.Now().UTC()
	old := DeliveryOwnershipRecord{
		Message:        types.InboxMessage{MessageID: "old", Target: "#ops"},
		Input:          drivercontract.RuntimeInput{ID: "old"},
		State:          DeliveryOwnershipInflight,
		InflightReason: DeliveryInflightOutcomeUnknown,
		AcceptedAt:     createdAt,
	}
	store := &reclaimProjectionRaceStore{
		records:      map[string]DeliveryOwnershipRecord{"old": old},
		freshEntered: make(chan struct{}),
		releaseFresh: make(chan struct{}),
	}
	ap := newInboxRetryTestProcess(&ackedRuntimeInputStubDriver{supports: true}, time.Hour)
	ap.deliveryOwnershipStore = store
	ap.pendingRuntimeInputs = map[string]pendingRuntimeInput{
		"old": {
			Message:        old.Message,
			Input:          old.Input,
			SessionID:      ap.SessionID,
			CreatedAt:      createdAt,
			InflightReason: DeliveryInflightOutcomeUnknown,
		},
	}
	// A stale reclaim projection must not re-terminalize a concurrently
	// acked owner. Give the projection every side-effect target the buggy
	// implementation touches: a live active wake, the workflow callback turn
	// owner, and the wake_failed runtime event channel. The assertions below
	// require all three to survive the stale projection unchanged.
	ap.pendingWake = &pendingDelivery{Wake: WakeInfo{MessageID: "old", Target: "#ops"}, Status: pendingDeliveryAwaitingCheck}
	ap.pendingWakeToken = "tok-old"
	var staleWakeFailures []WakeStatusEvent
	var staleRuntimeFailures []RuntimeEvent
	ap.onWakeStatus = func(ev WakeStatusEvent) {
		if ev.MessageID == "old" && ev.Status == consts.WakeStatusFailed {
			staleWakeFailures = append(staleWakeFailures, ev)
		}
	}
	ap.onRuntimeEvent = func(ev RuntimeEvent) {
		if ev.MessageID == "old" && ev.EventClass == consts.RuntimeEventClassWakeFailed {
			staleRuntimeFailures = append(staleRuntimeFailures, ev)
		}
	}
	t.Cleanup(func() { stopInboxRetryTestProcess(ap) })

	admissionDone := make(chan error, 1)
	go func() {
		admissionDone <- ap.persistQueuedDelivery(types.InboxMessage{MessageID: "fresh"}, time.Time{}, 0, nil)
	}()
	<-store.freshEntered

	ap.mu.Lock()
	pending, token, ok := ap.claimPendingRuntimeInputTransitionLocked("old", createdAt, nil)
	ap.mu.Unlock()
	require.True(t, ok)
	acked := ap.ownershipRecord(pending, deliveryOwnershipAcked, 0, nil)
	ackedDone := make(chan error, 1)
	go func() {
		_, err := ap.commitPendingRuntimeInputTransition(token, acked, true, func(current pendingRuntimeInput) {
			delete(ap.pendingRuntimeInputs, "old")
			ap.rememberRuntimeAckedMessageLocked("old")
		})
		ackedDone <- err
	}()

	require.Eventually(t, func() bool {
		if ap.mu.TryLock() {
			ap.mu.Unlock()
			return false
		}
		return true
	}, time.Second, time.Millisecond, "acked transition never blocked behind admission mutation")
	close(store.releaseFresh)
	require.NoError(t, <-ackedDone)
	require.NoError(t, <-admissionDone)

	store.mu.Lock()
	durable := store.records["old"]
	store.mu.Unlock()
	require.True(t, durable.IsAcked())
	ap.mu.Lock()
	_, ackedLocally := ap.ackedSet["old"]
	_, failedLocally := ap.failedSet["old"]
	ap.mu.Unlock()
	require.True(t, ackedLocally)
	require.False(t, failedLocally, "stale reclaim projection must not re-mark a concurrently acked owner failed")
	require.Empty(t, staleWakeFailures, "stale reclaim projection must not re-terminate the concurrently acked owner's live wake")
	require.Empty(t, staleRuntimeFailures, "stale reclaim projection must not emit wake_failed for a concurrently acked owner")
	ap.mu.Lock()
	require.NotNil(t, ap.pendingWake, "stale reclaim projection must not clear the concurrently acked owner's live wake")
	require.Equal(t, "old", ap.pendingWake.Wake.MessageID)
	ap.mu.Unlock()
}

// TestApplyReclaimedDeliveryOwnershipProjectsFullFailureForCurrentOwner is the
// positive counterpart to the stale-projection race: when the reclaimed record
// still matches the current live owner (no concurrent acked upgrade won),
// the projection must apply the full failure side effects -- terminate the
// wake, release the workflow callback turn owner, mark the owner failed, and
// emit wake_failed. This guards against the over-correction of making every
// reclaim projection a no-op.
func TestApplyReclaimedDeliveryOwnershipProjectsFullFailureForCurrentOwner(t *testing.T) {
	createdAt := time.Now().UTC()
	old := DeliveryOwnershipRecord{
		Message:        types.InboxMessage{MessageID: "old", Target: "#ops"},
		Input:          drivercontract.RuntimeInput{ID: "old"},
		State:          DeliveryOwnershipInflight,
		InflightReason: DeliveryInflightOutcomeUnknown,
		AcceptedAt:     createdAt,
	}
	ap := newInboxRetryTestProcess(&ackedRuntimeInputStubDriver{supports: true}, time.Hour)
	ap.pendingRuntimeInputs = map[string]pendingRuntimeInput{
		"old": {
			Message:        old.Message,
			Input:          old.Input,
			SessionID:      ap.SessionID,
			CreatedAt:      createdAt,
			InflightReason: DeliveryInflightOutcomeUnknown,
		},
	}
	ap.pendingWake = &pendingDelivery{Wake: WakeInfo{MessageID: "old", Target: "#ops"}, Status: pendingDeliveryAwaitingCheck}
	ap.pendingWakeToken = "tok-old"
	var wakeFailures []WakeStatusEvent
	var runtimeFailures []RuntimeEvent
	ap.onWakeStatus = func(ev WakeStatusEvent) {
		if ev.MessageID == "old" && ev.Status == consts.WakeStatusFailed {
			wakeFailures = append(wakeFailures, ev)
		}
	}
	ap.onRuntimeEvent = func(ev RuntimeEvent) {
		if ev.MessageID == "old" && ev.EventClass == consts.RuntimeEventClassWakeFailed {
			runtimeFailures = append(runtimeFailures, ev)
		}
	}
	t.Cleanup(func() { stopInboxRetryTestProcess(ap) })

	ap.applyReclaimedDeliveryOwnership([]DeliveryOwnershipRecord{{
		Message:     old.Message,
		Input:       old.Input,
		State:       DeliveryOwnershipTerminal,
		Outcome:     DeliveryOutcomeFailed,
		FailureCode: "ownership_backlog_reclaimed",
		AcceptedAt:  createdAt,
	}})

	require.Len(t, wakeFailures, 1, "current-owner reclaim must terminate the live wake")
	require.Len(t, runtimeFailures, 1, "current-owner reclaim must emit wake_failed")
	ap.mu.Lock()
	require.Nil(t, ap.pendingWake, "current-owner reclaim must clear the live wake")
	_, failedLocally := ap.failedSet["old"]
	ap.mu.Unlock()
	require.True(t, failedLocally, "current-owner reclaim must mark the owner failed")
}