package dispatcher

import (
	"bytes"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/tidbcloud/AgentTeam/desktop/internal/vault"
)

const (
	legacyUserPreferencesMigrationMarkerPrefix = "legacy-user-preferences-"
	legacyUserPreferencesMigrationStateDir     = ".memory-migrations"
	legacyUserPreferencesMigrationReadySuffix  = ".ready"
	legacyUserPreferencesMigrationDoneSuffix   = ".done"
	legacyUserPreferencesMigrationReadySchema  = 2

	legacyUserPreferencesMigrationPhaseGeneralPending = "general_import_pending"
	legacyUserPreferencesMigrationPhaseCutoverReady   = "legacy_cutover_ready"
)

type legacyUserPreferencesMigrationReadyState struct {
	Schema               int    `json:"schema"`
	Phase                string `json:"phase"`
	GeneralImportedBytes int64  `json:"generalImportedBytes"`
}

func legacyUserPreferencesMigrationComplete(workDir, agentID string, clone *vault.Clone) (bool, error) {
	if clone == nil {
		return false, nil
	}
	_, err := os.Stat(legacyUserPreferencesMigrationStatePath(workDir, agentID, clone.MemoryID(), clone.Branch(), legacyUserPreferencesMigrationDoneSuffix))
	if os.IsNotExist(err) {
		return false, nil
	}
	return err == nil, err
}

func legacyUserPreferencesMigrationReady(workDir, agentID string, clone *vault.Clone) (*legacyUserPreferencesMigrationReadyState, error) {
	if clone == nil {
		return nil, nil
	}
	content, err := os.ReadFile(legacyUserPreferencesMigrationStatePath(workDir, agentID, clone.MemoryID(), clone.Branch(), legacyUserPreferencesMigrationReadySuffix))
	if os.IsNotExist(err) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	// Pre-release empty markers did not record the general-stage budget. Treat
	// them as fully consumed so an upgrade cannot exceed the cumulative bound.
	if len(bytes.TrimSpace(content)) == 0 {
		return &legacyUserPreferencesMigrationReadyState{
			Schema:               legacyUserPreferencesMigrationReadySchema,
			Phase:                legacyUserPreferencesMigrationPhaseCutoverReady,
			GeneralImportedBytes: existingLocalMemoryImportMaxBytes,
		}, nil
	}
	var stored struct {
		Schema               int    `json:"schema"`
		Phase                string `json:"phase"`
		GeneralImportedBytes int64  `json:"generalImportedBytes"`
	}
	if err := json.Unmarshal(content, &stored); err != nil {
		return nil, err
	}
	// Schema 1 was briefly written before the general-import phase became
	// explicit. It was only persisted after the general push, so it is equivalent
	// to the cutover-ready phase.
	if stored.Schema == 1 {
		stored.Schema = legacyUserPreferencesMigrationReadySchema
		stored.Phase = legacyUserPreferencesMigrationPhaseCutoverReady
	}
	if stored.Schema != legacyUserPreferencesMigrationReadySchema ||
		(stored.Phase != legacyUserPreferencesMigrationPhaseGeneralPending && stored.Phase != legacyUserPreferencesMigrationPhaseCutoverReady) ||
		stored.GeneralImportedBytes < 0 || stored.GeneralImportedBytes > existingLocalMemoryImportMaxBytes {
		return nil, errors.New("invalid legacy user preference migration ready state")
	}
	return &legacyUserPreferencesMigrationReadyState{
		Schema:               stored.Schema,
		Phase:                stored.Phase,
		GeneralImportedBytes: stored.GeneralImportedBytes,
	}, nil
}

func markLegacyUserPreferencesGeneralImportPending(workDir, agentID, memoryID, branch string, generalImportedBytes int64) error {
	return writeLegacyUserPreferencesMigrationState(
		workDir, agentID, memoryID, branch,
		legacyUserPreferencesMigrationPhaseGeneralPending,
		generalImportedBytes,
	)
}

func markLegacyUserPreferencesMigrationReady(workDir, agentID, memoryID, branch string, generalImportedBytes int64) error {
	return writeLegacyUserPreferencesMigrationState(
		workDir, agentID, memoryID, branch,
		legacyUserPreferencesMigrationPhaseCutoverReady,
		generalImportedBytes,
	)
}

func writeLegacyUserPreferencesMigrationState(workDir, agentID, memoryID, branch, phase string, generalImportedBytes int64) error {
	content, err := json.Marshal(legacyUserPreferencesMigrationReadyState{
		Schema:               legacyUserPreferencesMigrationReadySchema,
		Phase:                phase,
		GeneralImportedBytes: generalImportedBytes,
	})
	if err != nil {
		return err
	}
	return writeLegacyUserPreferencesMigrationMarker(
		legacyUserPreferencesMigrationStatePath(workDir, agentID, memoryID, branch, legacyUserPreferencesMigrationReadySuffix),
		content,
	)
}

func markLegacyUserPreferencesMigrationComplete(workDir, agentID, memoryID, branch string) error {
	markerPath := legacyUserPreferencesMigrationStatePath(workDir, agentID, memoryID, branch, legacyUserPreferencesMigrationDoneSuffix)
	readyPath := legacyUserPreferencesMigrationStatePath(workDir, agentID, memoryID, branch, legacyUserPreferencesMigrationReadySuffix)
	if _, err := os.Stat(markerPath); err == nil {
		return nil
	} else if !os.IsNotExist(err) {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(markerPath), 0o700); err != nil {
		return err
	}
	if err := os.Rename(readyPath, markerPath); err == nil {
		return syncSessionResumeStateDir(filepath.Dir(markerPath))
	} else if !os.IsNotExist(err) {
		return err
	}
	return writeLegacyUserPreferencesMigrationMarker(markerPath, nil)
}

func writeLegacyUserPreferencesMigrationMarker(markerPath string, content []byte) error {
	if err := os.MkdirAll(filepath.Dir(markerPath), 0o700); err != nil {
		return err
	}
	tmp, err := os.CreateTemp(filepath.Dir(markerPath), ".legacy-user-preferences-*")
	if err != nil {
		return err
	}
	tmpPath := tmp.Name()
	defer os.Remove(tmpPath)
	if err := tmp.Chmod(0o600); err != nil {
		_ = tmp.Close()
		return err
	}
	if _, err := tmp.Write(content); 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, markerPath); err != nil {
		return err
	}
	return syncSessionResumeStateDir(filepath.Dir(markerPath))
}

func legacyUserPreferencesMigrationStatePath(workDir, agentID, memoryID, branch, suffix string) string {
	binding := fmt.Sprintf("legacy-user-preferences-v1\x00%s\x00%s\x00%s", strings.TrimSpace(agentID), strings.TrimSpace(memoryID), strings.TrimSpace(branch))
	digest := sha256.Sum256([]byte(binding))
	return filepath.Join(workDir, legacyUserPreferencesMigrationStateDir, fmt.Sprintf("%s%x%s", legacyUserPreferencesMigrationMarkerPrefix, digest, suffix))
}