CLIProxyAPI/internal/util/cookie_snapshot.go

package util

import (
	"encoding/json"
	"errors"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/luispater/CLIProxyAPI/v5/internal/misc"
)

const cookieSnapshotExt = ".cookie"

// CookieSnapshotPath derives the cookie snapshot file path from the main token JSON path.
// It replaces the .json suffix with .cookie, or appends .cookie if missing.
func CookieSnapshotPath(mainPath string) string {
	if strings.HasSuffix(mainPath, ".json") {
		return strings.TrimSuffix(mainPath, ".json") + cookieSnapshotExt
	}
	return mainPath + cookieSnapshotExt
}

// IsRegularFile reports whether the given path exists and is a regular file.
func IsRegularFile(path string) bool {
	if path == "" {
		return false
	}
	if st, err := os.Stat(path); err == nil && !st.IsDir() {
		return true
	}
	return false
}

// ReadJSON reads and unmarshals a JSON file into v.
// Returns os.ErrNotExist if the file does not exist.
func ReadJSON(path string, v any) error {
	b, err := os.ReadFile(path)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return os.ErrNotExist
		}
		return err
	}
	if len(b) == 0 {
		return nil
	}
	return json.Unmarshal(b, v)
}

// WriteJSON marshals v as JSON and writes to path, creating parent directories as needed.
func WriteJSON(path string, v any) error {
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return err
	}
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer func() { _ = f.Close() }()
	enc := json.NewEncoder(f)
	return enc.Encode(v)
}

// RemoveFile removes the file if it exists.
func RemoveFile(path string) error {
	if IsRegularFile(path) {
		return os.Remove(path)
	}
	return nil
}

// TryReadCookieSnapshotInto tries to read a cookie snapshot into v using the .cookie suffix.
// Returns (true, nil) when a snapshot was decoded, or (false, nil) when none exists.
func TryReadCookieSnapshotInto(mainPath string, v any) (bool, error) {
	snap := CookieSnapshotPath(mainPath)
	if err := ReadJSON(snap, v); err != nil {
		if err == os.ErrNotExist {
			return false, nil
		}
		return false, err
	}
	return true, nil
}

// WriteCookieSnapshot writes v to the snapshot path derived from mainPath using the .cookie suffix.
func WriteCookieSnapshot(mainPath string, v any) error {
	path := CookieSnapshotPath(mainPath)
	misc.LogSavingCredentials(path)
	if err := WriteJSON(path, v); err != nil {
		return err
	}
	return nil
}

// ReadAuthFilePreferSnapshot returns the first non-empty auth payload preferring snapshots.
func ReadAuthFilePreferSnapshot(path string) ([]byte, error) {
	return ReadAuthFileWithRetry(path, 1, 0)
}

// ReadAuthFileWithRetry attempts to read an auth file multiple times and prefers cookie snapshots.
func ReadAuthFileWithRetry(path string, attempts int, delay time.Duration) ([]byte, error) {
	if attempts < 1 {
		attempts = 1
	}
	read := func(target string) ([]byte, error) {
		var lastErr error
		for i := 0; i < attempts; i++ {
			data, err := os.ReadFile(target)
			if err == nil {
				return data, nil
			}
			lastErr = err
			if i < attempts-1 {
				time.Sleep(delay)
			}
		}
		return nil, lastErr
	}

	candidates := []string{
		CookieSnapshotPath(path),
		path,
	}

	for idx, candidate := range candidates {
		data, err := read(candidate)
		if err == nil {
			return data, nil
		}
		if errors.Is(err, os.ErrNotExist) {
			if idx < len(candidates)-1 {
				continue
			}
		}
		return nil, err
	}

	return nil, os.ErrNotExist
}

// RemoveCookieSnapshots removes the snapshot file if it exists.
func RemoveCookieSnapshots(mainPath string) {
	_ = RemoveFile(CookieSnapshotPath(mainPath))
}

// Hooks provide customization points for snapshot lifecycle operations.
type Hooks[T any] struct {
	// Apply merges snapshot data into the in-memory store during Apply().
	// Defaults to overwriting the store with the snapshot contents.
	Apply func(store *T, snapshot *T)

	// Snapshot prepares the payload to persist during Persist().
	// Defaults to cloning the store value.
	Snapshot func(store *T) *T

	// Merge chooses which data to flush when a snapshot exists.
	// Defaults to using the snapshot payload as-is.
	Merge func(store *T, snapshot *T) *T

	// WriteMain persists the merged payload into the canonical token path.
	// Defaults to WriteJSON.
	WriteMain func(path string, data *T) error
}

// Manager orchestrates cookie snapshot lifecycle for token storages.
type Manager[T any] struct {
	mainPath string
	store    *T
	hooks    Hooks[T]
}

// NewManager constructs a Manager bound to mainPath and store.
func NewManager[T any](mainPath string, store *T, hooks Hooks[T]) *Manager[T] {
	return &Manager[T]{
		mainPath: mainPath,
		store:    store,
		hooks:    hooks,
	}
}

// Apply loads snapshot data into the in-memory store if available.
// Returns true when a snapshot was applied.
func (m *Manager[T]) Apply() (bool, error) {
	if m == nil || m.store == nil || m.mainPath == "" {
		return false, nil
	}
	var snapshot T
	ok, err := TryReadCookieSnapshotInto(m.mainPath, &snapshot)
	if err != nil {
		return false, err
	}
	if !ok {
		return false, nil
	}
	if m.hooks.Apply != nil {
		m.hooks.Apply(m.store, &snapshot)
	} else {
		*m.store = snapshot
	}
	return true, nil
}

// Persist writes the current store state to the snapshot file.
func (m *Manager[T]) Persist() error {
	if m == nil || m.store == nil || m.mainPath == "" {
		return nil
	}
	var payload *T
	if m.hooks.Snapshot != nil {
		payload = m.hooks.Snapshot(m.store)
	} else {
		clone := new(T)
		*clone = *m.store
		payload = clone
	}
	return WriteCookieSnapshot(m.mainPath, payload)
}

// FlushOptions configure Flush behaviour.
type FlushOptions[T any] struct {
	Fallback func() *T
	Mutate   func(*T)
}

// FlushOption mutates FlushOptions.
type FlushOption[T any] func(*FlushOptions[T])

// WithFallback provides fallback payload when no snapshot exists.
func WithFallback[T any](fn func() *T) FlushOption[T] {
	return func(opts *FlushOptions[T]) { opts.Fallback = fn }
}

// WithMutate allows last-minute mutation of the payload before writing main file.
func WithMutate[T any](fn func(*T)) FlushOption[T] {
	return func(opts *FlushOptions[T]) { opts.Mutate = fn }
}

// Flush commits snapshot (or fallback) into the main token file and removes the snapshot.
func (m *Manager[T]) Flush(options ...FlushOption[T]) error {
	if m == nil || m.mainPath == "" {
		return nil
	}
	cfg := FlushOptions[T]{}
	for _, opt := range options {
		if opt != nil {
			opt(&cfg)
		}
	}
	var snapshot T
	ok, err := TryReadCookieSnapshotInto(m.mainPath, &snapshot)
	if err != nil {
		return err
	}
	var payload *T
	if ok {
		if m.hooks.Merge != nil {
			payload = m.hooks.Merge(m.store, &snapshot)
		} else {
			payload = &snapshot
		}
	} else if cfg.Fallback != nil {
		payload = cfg.Fallback()
	} else if m.store != nil {
		payload = m.store
	}
	if payload == nil {
		return RemoveFile(CookieSnapshotPath(m.mainPath))
	}
	if cfg.Mutate != nil {
		cfg.Mutate(payload)
	}
	if m.hooks.WriteMain != nil {
		if err := m.hooks.WriteMain(m.mainPath, payload); err != nil {
			return err
		}
	} else {
		if err := WriteJSON(m.mainPath, payload); err != nil {
			return err
		}
	}
	RemoveCookieSnapshots(m.mainPath)
	return nil
}