package dispatcher

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

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

const (
	runtimeDeliveryOwnershipVersion        = 2
	retiredRuntimeDeliveryOwnershipVersion = 3
	maxRuntimeDeliveryOwnershipRecords     = 1024 // 512 live + 512 terminal tombstones
	maxRuntimeDeliveryOwnershipBytes       = 10 << 20
	maxRuntimeDeliveryTerminalRecords      = 512
	maxRuntimeDeliveryTerminalBytes        = 2 << 20
	maxRuntimeDeliveryLiveRecords          = 512
	maxRuntimeDeliveryLiveBytes            = 8 << 20
	maxRuntimeDeliveryRecordBytes          = 512 << 10
)

type sessionDeliveryOwnershipStore struct {
	workDir string
}

func newSessionDeliveryOwnershipStore(workDir string) *sessionDeliveryOwnershipStore {
	if strings.TrimSpace(workDir) == "" {
		return nil
	}
	return &sessionDeliveryOwnershipStore{workDir: workDir}
}

func (s *sessionDeliveryOwnershipStore) LoadDeliveryOwnership() ([]agent.DeliveryOwnershipRecord, error) {
	state, err := loadSessionResumeState(s.workDir)
	if err != nil {
		return nil, err
	}
	switch state.RuntimeDeliveryOwnershipVersion {
	case runtimeDeliveryOwnershipVersion:
		records := append([]agent.DeliveryOwnershipRecord(nil), state.RuntimeDeliveryOwnership...)
		for i := range records {
			records[i] = agent.NormalizeDeliveryOwnershipRecordForRead(records[i])
		}
		for _, record := range records {
			if err := agent.ValidateDeliveryOwnershipRecord(record); err != nil {
				return nil, fmt.Errorf("load runtime delivery ownership: %w", err)
			}
		}
		return records, nil
	case 0:
		// No ownership was ever persisted; legacy records were migrated to the
		// current version during load.
		return nil, nil
	default:
		return nil, unsupportedRuntimeDeliveryOwnershipVersion(state.RuntimeDeliveryOwnershipVersion)
	}
}

func unsupportedRuntimeDeliveryOwnershipVersion(version int) error {
	return fmt.Errorf(
		"unsupported runtime delivery ownership version %d (supported read versions: 0, %d, %d)",
		version,
		runtimeDeliveryOwnershipVersion,
		retiredRuntimeDeliveryOwnershipVersion,
	)
}

func (s *sessionDeliveryOwnershipStore) PutDeliveryOwnership(record agent.DeliveryOwnershipRecord) ([]agent.DeliveryOwnershipRecord, error) {
	record = agent.NormalizeDeliveryOwnershipRecordForPersistence(record)
	if err := agent.ValidateDeliveryOwnershipRecord(record); err != nil {
		return nil, err
	}
	messageID := record.MessageID()
	if messageID == "" {
		return nil, fmt.Errorf("runtime delivery ownership requires message id")
	}
	var transactionReclaimed []agent.DeliveryOwnershipRecord
	err := updateSessionResumeState(s.workDir, func(state *sessionResumeState) error {
		switch state.RuntimeDeliveryOwnershipVersion {
		case 0:
			state.RuntimeDeliveryOwnership = nil
			state.RuntimeDeliveryOwnershipVersion = runtimeDeliveryOwnershipVersion
		case runtimeDeliveryOwnershipVersion:
			// current schema
		default:
			return unsupportedRuntimeDeliveryOwnershipVersion(state.RuntimeDeliveryOwnershipVersion)
		}
		records := append([]agent.DeliveryOwnershipRecord(nil), state.RuntimeDeliveryOwnership...)
		now := time.Now().UTC()
		if record.AcceptedAt.IsZero() {
			record.AcceptedAt = now
		}
		if record.UpdatedAt.IsZero() {
			record.UpdatedAt = now
		}
		replaced := false
		for i := range records {
			if records[i].MessageID() != messageID {
				continue
			}
			if !runtimeDeliveryOwnershipTransitionAllowed(records[i], record) {
				// A Process generation that started this write before an
				// authoritative acked/failed transition may arrive here later.
				// Terminal ownership is absorbing across Process instances; keep the
				// durable tombstone and treat the stale Put as an idempotent no-op.
				return nil
			}
			if !records[i].AcceptedAt.IsZero() {
				record.AcceptedAt = records[i].AcceptedAt
			}
			// Terminal pruning treats slice order as recency. Move every accepted
			// replacement to the tail instead of updating it in place; otherwise a
			// fresh failed->acked authoritative upgrade can retain its old index
			// and be pruned immediately as the apparent oldest acked tombstone.
			records = append(records[:i], records[i+1:]...)
			records = append(records, record)
			replaced = true
			break
		}
		if !replaced {
			records = append(records, record)
		}
		records = pruneRuntimeDeliveryTerminalRecords(records)
		records, reclaimErr := trimRuntimeDeliveryTerminalCapacity(records)
		if reclaimErr != nil {
			return reclaimErr
		}
		var reclaimed []agent.DeliveryOwnershipRecord
		// Capacity reclaim belongs only to a new message's first live admission,
		// whether it starts queued or inflight before a direct runtime mutation.
		// Existing-owner transitions must never choose a new victim: callers may
		// hold Process state locks and cannot safely project a second ownership
		// change from that transition.
		if !replaced && runtimeDeliveryOwnershipIsLive(record) {
			records, reclaimed, reclaimErr = reclaimRuntimeDeliveryCapacity(records, messageID, now)
			if reclaimErr != nil {
				return reclaimErr
			}
		}
		stats, statsErr := runtimeDeliveryOwnershipStats(records)
		if statsErr != nil {
			return statsErr
		}
		if !stats.withinLimits() {
			return fmt.Errorf("live runtime delivery ownership exceeds %d records or %d bytes", maxRuntimeDeliveryLiveRecords, maxRuntimeDeliveryLiveBytes)
		}
		state.RuntimeDeliveryOwnership = records
		if runtimeDeliveryOwnershipIsLive(record) {
			// The same fsync that accepts a queued/runtime-owned mutation must
			// fail-close restart safety. A separate markResumeSafety(false)
			// callback can fail, but the daemon may never retain safe=true on disk
			// after acknowledging an input that the runtime can mutate.
			// Keep epoch/version unchanged so older equal-version callbacks remain
			// idempotent no-ops rather than regaining authority.
			state.SessionResumeSafe = false
		}
		transactionReclaimed = reclaimed
		return nil
	})
	if err != nil {
		return nil, err
	}
	return transactionReclaimed, nil
}

func trimRuntimeDeliveryTerminalCapacity(records []agent.DeliveryOwnershipRecord, protected ...string) ([]agent.DeliveryOwnershipRecord, error) {
	protectedIDs := make(map[string]struct{}, len(protected))
	for _, messageID := range protected {
		if messageID = strings.TrimSpace(messageID); messageID != "" {
			protectedIDs[messageID] = struct{}{}
		}
	}
	for {
		stats, err := runtimeDeliveryOwnershipStats(records)
		if err != nil {
			return nil, err
		}
		if stats.terminalRecords <= maxRuntimeDeliveryTerminalRecords && stats.terminalBytes <= maxRuntimeDeliveryTerminalBytes &&
			stats.totalRecords <= maxRuntimeDeliveryOwnershipRecords && stats.totalBytes <= maxRuntimeDeliveryOwnershipBytes {
			return records, nil
		}
		oldest := -1
		for i := range records {
			if runtimeDeliveryOwnershipIsLive(records[i]) {
				continue
			}
			if _, keep := protectedIDs[records[i].MessageID()]; keep {
				continue
			}
			if oldest < 0 || deliveryOwnershipOlder(records[i], records[oldest]) {
				oldest = i
			}
		}
		if oldest < 0 {
			return records, nil
		}
		records = append(records[:oldest], records[oldest+1:]...)
	}
}

func runtimeDeliveryOwnershipIsLive(record agent.DeliveryOwnershipRecord) bool {
	return agent.IsLiveDeliveryOwnershipState(record.State)
}

func runtimeDeliveryOwnershipTransitionAllowed(current, next agent.DeliveryOwnershipRecord) bool {
	if runtimeDeliveryOwnershipIsLive(current) {
		return true
	}
	if runtimeDeliveryOwnershipIsLive(next) {
		return false
	}
	// Acked is authoritative positive evidence that the runtime acknowledged
	// responsibility for the exact correlated input at its supported boundary.
	// A late failure budget cannot downgrade it. Failed may still be upgraded to
	// acked if an authoritative correlated echo arrives later.
	if current.Outcome == agent.DeliveryOutcomeAcked && next.Outcome == agent.DeliveryOutcomeFailed {
		return false
	}
	return true
}

func (s *sessionDeliveryOwnershipStore) DeleteDeliveryOwnership(messageID string) error {
	messageID = strings.TrimSpace(messageID)
	if messageID == "" {
		return nil
	}
	return updateSessionResumeState(s.workDir, func(state *sessionResumeState) error {
		next := state.RuntimeDeliveryOwnership[:0]
		for _, record := range state.RuntimeDeliveryOwnership {
			if record.MessageID() == messageID && runtimeDeliveryOwnershipIsLive(record) {
				continue
			}
			// Delete is a rollback for a Process-local live admission. Never let a
			// stale rollback erase a terminal tombstone committed by another
			// Process generation after that admission began.
			next = append(next, record)
		}
		state.RuntimeDeliveryOwnership = next
		return nil
	})
}

func pruneRuntimeDeliveryTerminalRecords(records []agent.DeliveryOwnershipRecord) []agent.DeliveryOwnershipRecord {
	count := 0
	for i := len(records) - 1; i >= 0; i-- {
		if runtimeDeliveryOwnershipIsLive(records[i]) {
			continue
		}
		count++
		if count > maxRuntimeDeliveryTerminalRecords {
			records = append(records[:i], records[i+1:]...)
		}
	}
	return records
}

func reclaimRuntimeDeliveryCapacity(records []agent.DeliveryOwnershipRecord, protectedMessageID string, now time.Time) ([]agent.DeliveryOwnershipRecord, []agent.DeliveryOwnershipRecord, error) {
	for i := range records {
		records[i] = agent.NormalizeDeliveryOwnershipRecordForPersistence(records[i])
		if err := agent.ValidateDeliveryOwnershipRecord(records[i]); err != nil {
			return nil, nil, err
		}
	}
	protectedMessageID = strings.TrimSpace(protectedMessageID)
	var reclaimed []agent.DeliveryOwnershipRecord
	for {
		stats, err := runtimeDeliveryOwnershipStats(records)
		if err != nil {
			return nil, nil, err
		}
		if stats.withinLimits() {
			return records, reclaimed, nil
		}
		oldest := -1
		for i := range records {
			if !runtimeDeliveryOwnershipIsLive(records[i]) || records[i].MessageID() == protectedMessageID {
				continue
			}
			if oldest < 0 || deliveryOwnershipOlder(records[i], records[oldest]) {
				oldest = i
			}
		}
		if oldest >= 0 {
			failed := records[oldest]
			failed.State = agent.DeliveryOwnershipTerminal
			failed.InflightReason = ""
			failed.Outcome = agent.DeliveryOutcomeFailed
			failed.UpdatedAt = now
			failed.FailureCode = "ownership_backlog_reclaimed"
			failed.LastError = "older runtime delivery was terminalized to keep newer delivery admission available"
			reclaimed = append(reclaimed, failed)
			// Move the new tombstone to the recency tail, then trim pre-existing
			// terminal evidence first. The record returned to Process must remain in
			// this same durable transaction for same-ID dedup and late evidence.
			records = append(records[:oldest], records[oldest+1:]...)
			records = append(records, agent.NormalizeDeliveryOwnershipRecordForPersistence(failed))
			protectedReclaims := make([]string, 0, len(reclaimed))
			for _, reclaimedRecord := range reclaimed {
				protectedReclaims = append(protectedReclaims, reclaimedRecord.MessageID())
			}
			records, err = trimRuntimeDeliveryTerminalCapacity(records, protectedReclaims...)
			if err != nil {
				return nil, nil, err
			}
			continue
		}
		// No live owner except the incoming record remains. Drop the oldest
		// terminal evidence before rejecting a valid new delivery.
		terminal := -1
		for i := range records {
			if runtimeDeliveryOwnershipIsLive(records[i]) {
				continue
			}
			if terminal < 0 || deliveryOwnershipOlder(records[i], records[terminal]) {
				terminal = i
			}
		}
		if terminal < 0 {
			return nil, nil, fmt.Errorf("runtime delivery ownership %q cannot fit within bounded journal", protectedMessageID)
		}
		records = append(records[:terminal], records[terminal+1:]...)
	}
}

type runtimeDeliveryOwnershipCapacity struct {
	totalRecords, totalBytes, liveRecords, liveBytes, terminalRecords, terminalBytes int
}

func (s runtimeDeliveryOwnershipCapacity) withinLimits() bool {
	return s.totalRecords <= maxRuntimeDeliveryOwnershipRecords && s.totalBytes <= maxRuntimeDeliveryOwnershipBytes &&
		s.liveRecords <= maxRuntimeDeliveryLiveRecords && s.liveBytes <= maxRuntimeDeliveryLiveBytes &&
		s.terminalRecords <= maxRuntimeDeliveryTerminalRecords && s.terminalBytes <= maxRuntimeDeliveryTerminalBytes
}

func runtimeDeliveryOwnershipStats(records []agent.DeliveryOwnershipRecord) (runtimeDeliveryOwnershipCapacity, error) {
	stats := runtimeDeliveryOwnershipCapacity{totalRecords: len(records), totalBytes: 2, liveBytes: 2, terminalBytes: 2}
	for i, record := range records {
		size, err := runtimeDeliveryOwnershipRecordBytes(record)
		if err != nil {
			return stats, err
		}
		if runtimeDeliveryOwnershipIsLive(record) && size > maxRuntimeDeliveryRecordBytes {
			return stats, fmt.Errorf("runtime delivery ownership %q exceeds %d bytes", record.MessageID(), maxRuntimeDeliveryRecordBytes)
		}
		if i > 0 {
			stats.totalBytes++
		}
		stats.totalBytes += size
		if runtimeDeliveryOwnershipIsLive(record) {
			if stats.liveRecords > 0 {
				stats.liveBytes++
			}
			stats.liveRecords++
			stats.liveBytes += size
		} else {
			if stats.terminalRecords > 0 {
				stats.terminalBytes++
			}
			stats.terminalRecords++
			stats.terminalBytes += size
		}
	}
	return stats, nil
}

func deliveryOwnershipOlder(left, right agent.DeliveryOwnershipRecord) bool {
	if left.AcceptedAt.Equal(right.AcceptedAt) {
		return left.MessageID() < right.MessageID()
	}
	if left.AcceptedAt.IsZero() {
		return false
	}
	if right.AcceptedAt.IsZero() {
		return true
	}
	return left.AcceptedAt.Before(right.AcceptedAt)
}

func runtimeDeliveryOwnershipRecordBytes(record agent.DeliveryOwnershipRecord) (int, error) {
	data, err := json.Marshal(agent.NormalizeDeliveryOwnershipRecordForPersistence(record))
	if err != nil {
		return 0, fmt.Errorf("encode runtime delivery ownership %q: %w", record.MessageID(), err)
	}
	return len(data), nil
}

func oldestLiveRuntimeDeliveryOwnership(workDir string) (agent.DeliveryOwnershipRecord, bool, error) {
	store := newSessionDeliveryOwnershipStore(workDir)
	if store == nil {
		return agent.DeliveryOwnershipRecord{}, false, nil
	}
	records, err := store.LoadDeliveryOwnership()
	if err != nil {
		return agent.DeliveryOwnershipRecord{}, false, err
	}
	var oldest agent.DeliveryOwnershipRecord
	found := false
	for _, record := range records {
		if !runtimeDeliveryOwnershipIsLive(record) {
			continue
		}
		if record.MessageID() == "" {
			return agent.DeliveryOwnershipRecord{}, false, fmt.Errorf("live runtime delivery ownership is missing message id")
		}
		if !found || (!record.AcceptedAt.IsZero() && (oldest.AcceptedAt.IsZero() || record.AcceptedAt.Before(oldest.AcceptedAt))) {
			oldest = record
			found = true
		}
	}
	return oldest, found, nil
}

type runtimeDeliveryLaunchRestore struct {
	Message           types.InboxMessage
	HandoffPersisted  bool
	WaitForReconcile  bool
	NewlyPersistedIDs []string
	Reclaimed         []agent.DeliveryOwnershipRecord
}

// prepareRuntimeDeliveryOwnershipLaunch closes the accepted-then-crashed gap.
// Every owner launches without a prompt. A concrete current wake is the
// foreground handoff for this launch; historical queued/inflight owners are
// restored as background backlog and may not replace it with an empty
// reconciliation launch. When no current wake exists, continuation selects the
// oldest live owner. An inflight owner is never replayed speculatively.
func prepareRuntimeDeliveryOwnershipLaunch(workDir string, current types.InboxMessage, requireDurableHandoff bool) (runtimeDeliveryLaunchRestore, error) {
	store := newSessionDeliveryOwnershipStore(workDir)
	var existing []agent.DeliveryOwnershipRecord
	if store != nil {
		var err error
		existing, err = store.LoadDeliveryOwnership()
		if err != nil {
			return runtimeDeliveryLaunchRestore{}, err
		}
	}
	currentID := strings.TrimSpace(current.MessageID)
	var currentRecord agent.DeliveryOwnershipRecord
	currentAlreadyPersisted := false
	for _, record := range existing {
		if record.MessageID() == currentID && currentID != "" {
			currentAlreadyPersisted = true
			currentRecord = record
			break
		}
	}
	result := runtimeDeliveryLaunchRestore{Message: current}
	if currentID != "" && !currentAlreadyPersisted {
		reclaimed, err := store.PutDeliveryOwnership(agent.DeliveryOwnershipRecord{
			Message:    current,
			Input:      drivercontract.RuntimeInput{ID: currentID},
			State:      agent.DeliveryOwnershipQueued,
			AcceptedAt: time.Now().UTC(),
			UpdatedAt:  time.Now().UTC(),
		})
		if err != nil {
			return runtimeDeliveryLaunchRestore{}, fmt.Errorf("persist newer wake behind accepted runtime delivery: %w", err)
		}
		result.NewlyPersistedIDs = append(result.NewlyPersistedIDs, currentID)
		result.Reclaimed = append(result.Reclaimed, reclaimed...)
		currentRecord = agent.DeliveryOwnershipRecord{Message: current, State: agent.DeliveryOwnershipQueued}
		currentAlreadyPersisted = true
	}
	if currentID != "" && currentAlreadyPersisted {
		if !agent.IsLiveDeliveryOwnershipState(currentRecord.State) {
			// Same-ID upstream ACK retry after restart: launch without replaying
			// terminal content so SetDeliveryOwnershipStore can restore the dedup
			// tombstone before the session becomes visible.
			result.Message = types.InboxMessage{}
			result.WaitForReconcile = true
			return result, nil
		}
		switch currentRecord.State {
		case agent.DeliveryOwnershipQueued:
			result.HandoffPersisted = true
			return result, nil
		case agent.DeliveryOwnershipInflight:
			// A duplicate of an already-dispatched identity cannot be replayed.
			// Restore and reconcile it without manufacturing a second input.
			result.Message = types.InboxMessage{}
			result.WaitForReconcile = true
			return result, nil
		}
	}

	oldest, ok, err := oldestLiveRuntimeDeliveryOwnership(workDir)
	if err != nil {
		return runtimeDeliveryLaunchRestore{}, err
	}
	if !ok {
		return result, nil
	}

	switch oldest.State {
	case agent.DeliveryOwnershipQueued:
		result.Message = oldest.Message
		result.HandoffPersisted = true
		return result, nil
	case agent.DeliveryOwnershipInflight:
		result.Message = types.InboxMessage{}
		result.WaitForReconcile = true
		return result, nil
	default:
		return runtimeDeliveryLaunchRestore{}, fmt.Errorf(
			"runtime delivery ownership %q has unsupported live state %q",
			oldest.MessageID(),
			oldest.State,
		)
	}
}

func rollbackRuntimeDeliveryOwnershipLaunch(workDir string, messageIDs []string) {
	store := newSessionDeliveryOwnershipStore(workDir)
	if store == nil {
		return
	}
	for _, messageID := range messageIDs {
		_ = store.DeleteDeliveryOwnership(messageID)
	}
}