package dispatcher

import (
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

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

const sessionResumeStateFilename = ".daemon-state.json"

type sessionResumeState struct {
	LastSessionID                   string                          `json:"lastSessionId,omitempty"`
	SessionResumeSafe               bool                            `json:"sessionResumeSafe,omitempty"`
	SessionResumeSafetyEpoch        uint64                          `json:"sessionResumeSafetyEpoch,omitempty"`
	SessionResumeSafetyVersion      uint64                          `json:"sessionResumeSafetyVersion,omitempty"`
	LastStartConfig                 *types.AgentStartConfig         `json:"lastStartConfig,omitempty"`
	LastRuntimePolicy               *types.RuntimePolicySnapshot    `json:"lastRuntimePolicy,omitempty"`
	LastRuntimePolicyFingerprint    string                          `json:"lastRuntimePolicyFingerprint,omitempty"`
	LastRuntimePolicyGeneration     int64                           `json:"lastRuntimePolicyGeneration,omitempty"`
	ActiveTaskTarget                types.ActiveTaskTarget          `json:"activeTaskTarget,omitempty,omitzero"`
	DeferredDeliveries              []deferredDeliveryEntry         `json:"deferredDeliveries,omitempty"`
	RuntimeDeliveryOwnershipVersion int                             `json:"runtimeDeliveryOwnershipVersion,omitempty"`
	RuntimeDeliveryOwnership        []agent.DeliveryOwnershipRecord `json:"runtimeDeliveryOwnership,omitempty"`
}

type sessionResumeStateStore struct {
	locks   sync.Map // map[cleanWorkDir]*sync.Mutex
	syncDir func(string) error
}

var defaultSessionResumeStateStore = &sessionResumeStateStore{}

func cloneAgentStartConfig(cfg *types.AgentStartConfig) *types.AgentStartConfig {
	if cfg == nil {
		return nil
	}
	cloned := *cfg
	return &cloned
}

func loadSessionResumeState(workDir string) (sessionResumeState, error) {
	return defaultSessionResumeStateStore.Load(workDir)
}

func saveSessionResumeState(workDir string, state sessionResumeState) error {
	return defaultSessionResumeStateStore.Save(workDir, state)
}

func updateSessionResumeState(workDir string, mutate func(*sessionResumeState) error) error {
	return defaultSessionResumeStateStore.Update(workDir, mutate)
}

func (s *sessionResumeStateStore) Load(workDir string) (sessionResumeState, error) {
	lock := s.lockFor(workDir)
	lock.Lock()
	defer lock.Unlock()
	return s.loadUnlocked(workDir)
}

func (s *sessionResumeStateStore) Save(workDir string, state sessionResumeState) error {
	lock := s.lockFor(workDir)
	lock.Lock()
	defer lock.Unlock()
	return s.saveUnlocked(workDir, state)
}

func (s *sessionResumeStateStore) Update(workDir string, mutate func(*sessionResumeState) error) error {
	if mutate == nil {
		return nil
	}

	lock := s.lockFor(workDir)
	lock.Lock()
	defer lock.Unlock()

	state, err := s.loadUnlocked(workDir)
	if err != nil {
		return err
	}
	if err := mutate(&state); err != nil {
		return err
	}
	return s.saveUnlocked(workDir, state)
}

func (s *sessionResumeStateStore) Delete(workDir string) error {
	lock := s.lockFor(workDir)
	lock.Lock()
	defer lock.Unlock()

	path := sessionResumeStatePath(workDir)
	if err := os.Remove(path); err != nil {
		if !errors.Is(err, os.ErrNotExist) {
			return err
		}
	}
	// Even NotExist must re-sync the parent. A prior attempt may have removed
	// the entry successfully and then failed directory fsync; treating the
	// retry's NotExist as durable would silently skip the missing boundary.
	if err := s.syncDirectory(filepath.Dir(path)); err != nil {
		return fmt.Errorf("sync session resume state directory after delete: %w", err)
	}
	return nil
}

func (s *sessionResumeStateStore) lockFor(workDir string) *sync.Mutex {
	key := strings.TrimSpace(workDir)
	if key == "" {
		key = "."
	}
	key = filepath.Clean(key)
	lock, _ := s.locks.LoadOrStore(key, &sync.Mutex{})
	return lock.(*sync.Mutex)
}

func (s *sessionResumeStateStore) loadUnlocked(workDir string) (sessionResumeState, error) {
	var state sessionResumeState
	data, err := os.ReadFile(sessionResumeStatePath(workDir))
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return state, nil
		}
		return state, err
	}
	if len(data) == 0 {
		return state, nil
	}
	if err := json.Unmarshal(data, &state); err != nil {
		return state, fmt.Errorf("decode %s: %w", sessionResumeStateFilename, err)
	}
	switch state.RuntimeDeliveryOwnershipVersion {
	case 0:
		if len(state.RuntimeDeliveryOwnership) > 0 {
			migrated, migrateErr := migrateLegacyRuntimeDeliveryOwnership(data)
			if migrateErr != nil {
				return state, migrateErr
			}
			state.RuntimeDeliveryOwnership = migrated
			state.RuntimeDeliveryOwnershipVersion = runtimeDeliveryOwnershipVersion
		}
	case runtimeDeliveryOwnershipVersion:
		// Current format.
	case retiredRuntimeDeliveryOwnershipVersion:
		if err := migrateRuntimeDeliveryOwnershipV3(data, &state); err != nil {
			return state, err
		}
	}
	if state.RuntimeDeliveryOwnershipVersion == runtimeDeliveryOwnershipVersion {
		for i := range state.RuntimeDeliveryOwnership {
			state.RuntimeDeliveryOwnership[i] = agent.NormalizeDeliveryOwnershipRecordForRead(
				state.RuntimeDeliveryOwnership[i],
			)
		}
	}
	// Unknown future ownership versions are preserved verbatim: a newer daemon
	// may have acknowledged queued/inflight owners in a schema this binary cannot
	// interpret. Clearing them here would silently drop acknowledged ownership
	// and let the next save make that loss durable. LoadDeliveryOwnership and the
	// save path fail closed on such versions instead.
	return state, nil
}

func migrateRuntimeDeliveryOwnershipV3(_ []byte, state *sessionResumeState) error {
	// providerSessionSeal belonged to a retired provider-session rotation
	// feature. It is intentionally ignored: provider-session compatibility is
	// guarded only by current runtime policy and semantic instructions, while
	// delivery ownership keeps its own reconciliation journal. The outer state
	// decode already validated the JSON envelope; retired metadata is not a
	// second schema that this daemon needs to parse.
	state.RuntimeDeliveryOwnershipVersion = runtimeDeliveryOwnershipVersion
	return nil
}

func migrateLegacyRuntimeDeliveryOwnership(data []byte) ([]agent.DeliveryOwnershipRecord, error) {
	var envelope struct {
		Records []struct {
			Message     types.InboxMessage          `json:"message"`
			Input       drivercontract.RuntimeInput `json:"input"`
			SessionID   string                      `json:"sessionId"`
			WakeToken   string                      `json:"wakeToken"`
			Phase       string                      `json:"phase"`
			AcceptedAt  time.Time                   `json:"acceptedAt"`
			UpdatedAt   time.Time                   `json:"updatedAt"`
			Attempts    int                         `json:"attempts"`
			LastError   string                      `json:"lastError"`
			FailureCode string                      `json:"failureCode"`
		} `json:"runtimeDeliveryOwnership"`
	}
	if err := json.Unmarshal(data, &envelope); err != nil {
		return nil, fmt.Errorf("decode legacy runtime delivery ownership: %w", err)
	}
	records := make([]agent.DeliveryOwnershipRecord, 0, len(envelope.Records))
	for _, legacy := range envelope.Records {
		record := agent.DeliveryOwnershipRecord{
			Message: legacy.Message, Input: legacy.Input, SessionID: legacy.SessionID,
			WakeToken: legacy.WakeToken, AcceptedAt: legacy.AcceptedAt, UpdatedAt: legacy.UpdatedAt,
			Attempts: legacy.Attempts, LastError: legacy.LastError, FailureCode: legacy.FailureCode,
		}
		switch strings.TrimSpace(legacy.Phase) {
		case "queued":
			record.State = agent.DeliveryOwnershipQueued
		case "pending":
			record.State = agent.DeliveryOwnershipInflight
			record.InflightReason = agent.DeliveryInflightDispatching
		case "uncertain":
			record.State = agent.DeliveryOwnershipInflight
			record.InflightReason = agent.DeliveryInflightOutcomeUnknown
		case "surfaced", "acked":
			record.State = agent.DeliveryOwnershipTerminal
			record.Outcome = agent.DeliveryOutcomeAcked
		case "failed":
			record.State = agent.DeliveryOwnershipTerminal
			record.Outcome = agent.DeliveryOutcomeFailed
		default:
			// The legacy v0 phase vocabulary is closed: only daemons that predate
			// the versioned schema ever wrote it, and early builds persisted
			// accepted deliveries with an empty phase. Failing closed here would
			// block every future agent start on stale data that can never heal
			// itself. Preserve the record as a terminal failed tombstone instead:
			// launch ignores terminal owners while deduplication identity and
			// diagnostics stay intact.
			if record.MessageID() == "" {
				slog.Warn("dropping legacy runtime delivery ownership without message id",
					"phase", legacy.Phase)
				continue
			}
			record.State = agent.DeliveryOwnershipTerminal
			record.Outcome = agent.DeliveryOutcomeFailed
			if record.FailureCode == "" {
				record.FailureCode = "legacy_phase_unsupported"
			}
			if record.LastError == "" {
				record.LastError = fmt.Sprintf("legacy runtime delivery ownership phase %q could not be migrated", legacy.Phase)
			}
			slog.Warn("legacy runtime delivery ownership has unsupported phase; migrated to terminal failed",
				"message_id", record.MessageID(), "phase", legacy.Phase)
		}
		record = agent.NormalizeDeliveryOwnershipRecordForPersistence(record)
		if err := agent.ValidateDeliveryOwnershipRecord(record); err != nil {
			return nil, fmt.Errorf("migrate legacy runtime delivery ownership: %w", err)
		}
		records = append(records, record)
	}
	return records, nil
}

func (s *sessionResumeStateStore) saveUnlocked(workDir string, state sessionResumeState) error {
	if strings.TrimSpace(workDir) == "" {
		return fmt.Errorf("session resume state workDir is required")
	}
	// Fail closed on any unknown non-zero ownership version unconditionally, even
	// when the record slice is empty or holds a payload this binary could not
	// interpret. A newer daemon may have persisted ownership in a schema we cannot
	// decode; an unrelated resume-state Update that skipped this check would
	// rewrite the file and make the silent loss durable. The version stamp below
	// only applies when we are persisting records we actually understand.
	if state.RuntimeDeliveryOwnershipVersion != 0 && state.RuntimeDeliveryOwnershipVersion != runtimeDeliveryOwnershipVersion {
		return unsupportedRuntimeDeliveryOwnershipVersion(state.RuntimeDeliveryOwnershipVersion)
	}
	if len(state.RuntimeDeliveryOwnership) > 0 {
		for _, record := range state.RuntimeDeliveryOwnership {
			if err := agent.ValidateDeliveryOwnershipRecord(record); err != nil {
				return fmt.Errorf("save runtime delivery ownership: %w", err)
			}
		}
		state.RuntimeDeliveryOwnershipVersion = runtimeDeliveryOwnershipVersion
	}
	data, err := json.Marshal(state)
	if err != nil {
		return err
	}

	path := sessionResumeStatePath(workDir)
	dir := filepath.Dir(path)
	tmp, err := os.CreateTemp(dir, sessionResumeStateFilename+".tmp-*")
	if err != nil {
		return err
	}
	tmpPath := tmp.Name()
	cleanup := true
	defer func() {
		if cleanup {
			_ = os.Remove(tmpPath)
		}
	}()

	if _, err := tmp.Write(data); err != nil {
		_ = tmp.Close()
		return err
	}
	if err := tmp.Sync(); err != nil {
		_ = tmp.Close()
		return err
	}
	if err := tmp.Close(); err != nil {
		return err
	}
	if err := os.Rename(tmpPath, path); err != nil {
		return err
	}
	cleanup = false
	if err := s.syncDirectory(dir); err != nil {
		return fmt.Errorf("sync session resume state directory: %w", err)
	}
	return nil
}

func sessionResumeStatePath(workDir string) string {
	return filepath.Join(workDir, sessionResumeStateFilename)
}

func (s *sessionResumeStateStore) syncDirectory(dir string) error {
	if s != nil && s.syncDir != nil {
		return s.syncDir(dir)
	}
	return syncSessionResumeStateDir(dir)
}