ssh-netbox-wrapper/internal/cache/cache.go

package cache

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"golang.org/x/sys/unix"

	"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/config"
	"git.zb-server.de/Sebi/ssh-netbox-wrapper/internal/storage"
)

type Entry struct {
	Name     string    `json:"name"`
	IP       string    `json:"ip"`
	Kind     string    `json:"kind"`
	Tags     []string  `json:"tags,omitempty"`
	CachedAt time.Time `json:"cached_at"`
	LastUsed time.Time `json:"last_used,omitempty"`
	// Static marks a host from the config's hosts: list. It never expires;
	// it is stored only to remember LastUsed for the recently-used list.
	Static bool `json:"static,omitempty"`
}

type changeKind uint8

const (
	upsertChange changeKind = iota
	deleteChange
	usageChange
)

type Cache struct {
	mu           sync.RWMutex
	dirty        map[string]changeKind
	cleared      bool
	entries      map[string]Entry
	path         string
	ttl          time.Duration
	refreshStamp string // path to last_refresh timestamp file
	static       []config.StaticHost
}

type diskFormat struct {
	Entries []Entry `json:"entries"`
}

func New(path string, ttlSeconds int) *Cache {
	stamp := ""
	if path != "" {
		stamp = filepath.Join(filepath.Dir(path), "last_refresh")
	}
	return &Cache{
		entries:      make(map[string]Entry),
		dirty:        make(map[string]changeKind),
		path:         path,
		ttl:          time.Duration(ttlSeconds) * time.Second,
		refreshStamp: stamp,
	}
}

// NeedsRefresh reports whether the last full cache refresh (via `cache refresh`)
// is older than d, or has never happened. Always returns false when no path is set.
func (c *Cache) NeedsRefresh(d time.Duration) bool {
	if c.refreshStamp == "" {
		return false
	}
	data, err := os.ReadFile(c.refreshStamp)
	if err != nil {
		return true // file missing → never refreshed
	}
	var t time.Time
	if err := t.UnmarshalText(data); err != nil {
		return true
	}
	return time.Since(t) >= d
}

// SetRefreshed records the current time as the last successful full refresh.
func (c *Cache) SetRefreshed() error {
	if c.refreshStamp == "" {
		return nil
	}
	data, err := time.Now().MarshalText()
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(c.refreshStamp), 0o755); err != nil {
		return err
	}
	return os.WriteFile(c.refreshStamp, data, 0o644)
}

func (c *Cache) Load() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	data, err := os.ReadFile(c.path)
	if os.IsNotExist(err) {
		c.applyStatic()
		return nil
	}
	if err != nil {
		return err
	}

	var df diskFormat
	if err := json.Unmarshal(data, &df); err != nil {
		return err
	}

	c.dirty = make(map[string]changeKind)
	c.cleared = false
	c.entries = make(map[string]Entry, len(df.Entries))
	for _, e := range df.Entries {
		c.entries[e.Name] = e
	}
	c.applyStatic()
	return nil
}

// Save merges only local changes under a process-wide file lock. Atomic rename
// alone prevents torn JSON but cannot prevent lost updates from another process.
func (c *Cache) Save() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if err := os.MkdirAll(filepath.Dir(c.path), 0o755); err != nil {
		return err
	}
	lock, err := os.OpenFile(c.path+".lock", os.O_CREATE|os.O_RDWR, 0o600)
	if err != nil {
		return err
	}
	defer lock.Close()
	if err := unix.Flock(int(lock.Fd()), unix.LOCK_EX); err != nil {
		return err
	}
	defer unix.Flock(int(lock.Fd()), unix.LOCK_UN)

	merged := make(map[string]Entry)
	if !c.cleared {
		data, err := os.ReadFile(c.path)
		if err != nil && !os.IsNotExist(err) {
			return err
		}
		if err == nil {
			var df diskFormat
			if err := json.Unmarshal(data, &df); err != nil {
				return err
			}
			for _, e := range df.Entries {
				merged[e.Name] = e
			}
		}
	}
	for name, change := range c.dirty {
		if change == deleteChange {
			delete(merged, name)
			continue
		}
		e := c.entries[name]
		old, exists := merged[name]
		if change == usageChange && !exists {
			// A connection finishing after a refresh or clear must not
			// restore an entry that the other process removed.
			continue
		}
		if exists {
			lastUsed := e.LastUsed
			if change == usageChange || old.CachedAt.After(e.CachedAt) {
				e = old
			}
			if old.LastUsed.After(lastUsed) {
				lastUsed = old.LastUsed
			}
			e.LastUsed = lastUsed
		}
		merged[name] = e
	}
	df := diskFormat{Entries: make([]Entry, 0, len(merged))}
	for _, e := range merged {
		df.Entries = append(df.Entries, e)
	}
	sort.Slice(df.Entries, func(i, j int) bool { return df.Entries[i].Name < df.Entries[j].Name })
	data, err := json.MarshalIndent(df, "", "  ")
	if err != nil {
		return err
	}
	if err := storage.WriteFile(c.path, data, 0o600); err != nil {
		return err
	}
	// Static hosts come from the config, not the file: keep them in memory.
	for name, e := range c.entries {
		if e.Static {
			if old, ok := merged[name]; ok && old.LastUsed.After(e.LastUsed) {
				e.LastUsed = old.LastUsed
			}
			merged[name] = e
		}
	}
	c.entries = merged
	c.dirty = make(map[string]changeKind)
	c.cleared = false
	return nil
}

func (c *Cache) Upsert(e Entry) {
	e.CachedAt = time.Now()
	c.mu.Lock()
	if old, ok := c.entries[e.Name]; ok && e.LastUsed.IsZero() {
		e.LastUsed = old.LastUsed
	}
	c.entries[e.Name] = e
	c.dirty[e.Name] = upsertChange
	c.mu.Unlock()
}

// Delete removes an entry when a refresh can no longer resolve it safely.
func (c *Cache) Delete(name string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	delete(c.entries, name)
	c.dirty[name] = deleteChange
}

// MarkUsed records the current time as LastUsed for the named entry.
// It is a no-op if the entry does not exist.
func (c *Cache) MarkUsed(name string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if e, ok := c.entries[name]; ok {
		e.LastUsed = time.Now()
		c.entries[name] = e
		if e.Static {
			// A static host may not be on disk yet; usage alone would not add it.
			c.dirty[name] = upsertChange
		} else if _, changed := c.dirty[name]; !changed {
			c.dirty[name] = usageChange
		}
	}
}

// RecentlyUsed returns the n most recently used entries, sorted by LastUsed desc.
// Entries that have never been used (LastUsed zero) are excluded.
// If n <= 0, all used entries are returned.
func (c *Cache) RecentlyUsed(n int) []Entry {
	c.mu.RLock()
	defer c.mu.RUnlock()

	var used []Entry
	for _, e := range c.entries {
		if !e.LastUsed.IsZero() {
			used = append(used, e)
		}
	}
	sort.Slice(used, func(i, j int) bool {
		return used[i].LastUsed.After(used[j].LastUsed)
	})
	if n > 0 && len(used) > n {
		used = used[:n]
	}
	return used
}

// Search returns all entries whose name starts with prefix (case-insensitive).
// TTL is intentionally ignored — this is used for shell completion.
func (c *Cache) Search(prefix string) []Entry {
	c.mu.RLock()
	defer c.mu.RUnlock()

	prefix = strings.ToLower(prefix)
	var out []Entry
	for name, e := range c.entries {
		_, hostname, qualified := strings.Cut(strings.ToLower(name), "/")
		if strings.HasPrefix(strings.ToLower(name), prefix) || (qualified && !strings.Contains(prefix, "/") && strings.HasPrefix(hostname, prefix)) {
			out = append(out, e)
		}
	}
	return out
}

// GetByShortcut rejects collisions instead of choosing a random map entry.
func (c *Cache) GetByShortcut(shortcut string, normalize func(string) string) (entry Entry, found bool, fresh bool, err error) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	norm := normalize(shortcut)
	for _, e := range c.entries {
		if normalize(e.Name) == norm {
			if found {
				return Entry{}, false, false, fmt.Errorf("ambiguous shortcut %q; use a full hostname", shortcut)
			}
			entry, found = e, true
		}
	}
	return entry, found, found && (entry.Static || c.ttl > 0 && time.Since(entry.CachedAt) < c.ttl), nil
}

// Get returns an entry and reports whether it is still within the TTL.
func (c *Cache) Get(name string) (entry Entry, fresh bool) {
	c.mu.RLock()
	e, ok := c.entries[name]
	c.mu.RUnlock()

	if !ok {
		return Entry{}, false
	}
	if e.Static {
		return e, true
	}
	if c.ttl == 0 {
		return e, false
	}
	return e, time.Since(e.CachedAt) < c.ttl
}

// Clear removes all NetBox entries; static hosts come from the config and stay.
func (c *Cache) Clear() {
	c.mu.Lock()
	static := make(map[string]Entry)
	for name, e := range c.entries {
		if e.Static {
			static[name] = e
		}
	}
	c.entries = static
	c.dirty = make(map[string]changeKind)
	c.cleared = true
	c.mu.Unlock()
}

// SetStatic adds the config's static hosts. They replace cached entries of the
// same name, keep their recorded LastUsed, and static entries that are no
// longer configured are removed on the next Save.
func (c *Cache) SetStatic(hosts []config.StaticHost) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.static = hosts
	c.applyStatic()
}

// WithStatic sets the static hosts now and after every Load.
func (c *Cache) WithStatic(hosts []config.StaticHost) *Cache {
	c.SetStatic(hosts)
	return c
}

func (c *Cache) applyStatic() {
	hosts := c.static
	configured := make(map[string]bool, len(hosts))
	for _, host := range hosts {
		e := Entry{Name: host.Name, IP: host.Address, Kind: config.StaticKind, Tags: host.Tags, CachedAt: time.Now(), Static: true}
		if old, ok := c.entries[host.Name]; ok && old.Static {
			e.LastUsed = old.LastUsed
		}
		c.entries[host.Name] = e
		configured[host.Name] = true
	}
	for name, e := range c.entries {
		if e.Static && !configured[name] {
			delete(c.entries, name)
			c.dirty[name] = deleteChange
		}
	}
}

// GetStatic finds a static host by name or by shortcut. Static hosts are
// checked before NetBox, so they also work when NetBox is unreachable.
func (c *Cache) GetStatic(name string, normalize func(string) string) (Entry, bool, error) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	var match Entry
	found := false
	for key, e := range c.entries {
		if e.Static && strings.EqualFold(key, name) {
			return e, true, nil
		}
	}
	if normalize == nil {
		return Entry{}, false, nil
	}
	short := normalize(name)
	for _, e := range c.entries {
		if e.Static && normalize(e.Name) == short {
			if found {
				return Entry{}, false, fmt.Errorf("ambiguous shortcut %q; use a full hostname", name)
			}
			match, found = e, true
		}
	}
	return match, found, nil
}

func (c *Cache) All() []Entry {
	c.mu.RLock()
	defer c.mu.RUnlock()
	out := make([]Entry, 0, len(c.entries))
	for _, e := range c.entries {
		out = append(out, e)
	}
	return out
}