HostKeeper/pkg/storage/json_storage.go

package storage

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"

	"git.tukangketik.id/swanadiva/hostkeeper/internal/models"
)

// JSONStorage implements Storage interface using JSON files
type JSONStorage struct {
	dataDir string
	mu      sync.RWMutex
}

// NewJSONStorage creates a new JSON storage instance
func NewJSONStorage(dataDir string) (*JSONStorage, error) {
	if err := os.MkdirAll(dataDir, 0700); err != nil {
		return nil, fmt.Errorf("failed to create data directory: %w", err)
	}

	s := &JSONStorage{dataDir: dataDir}

	if err := s.ensureDataFiles(); err != nil {
		return nil, fmt.Errorf("failed to initialize data files: %w", err)
	}

	return s, nil
}

func (s *JSONStorage) ensureDataFiles() error {
	files := map[string]string{
		"hosts.json":    "hosts",
		"keys.json":     "key_pairs",
		"snippets.json": "snippets",
	}

	for file, key := range files {
		path := filepath.Join(s.dataDir, file)
		if _, err := os.Stat(path); os.IsNotExist(err) {
			data := map[string]interface{}{key: []interface{}{}}
			dataBytes, _ := json.MarshalIndent(data, "", "  ")
			if err := os.WriteFile(path, dataBytes, 0600); err != nil {
				return err
			}
		}
	}
	return nil
}

// ============ Host Operations ============

func (s *JSONStorage) getHostsPath() string {
	return filepath.Join(s.dataDir, "hosts.json")
}

func (s *JSONStorage) ListHosts(ctx context.Context) ([]*models.Host, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.listHostsInternal()
}

// listHostsInternal reads hosts WITHOUT locking (caller must hold lock)
func (s *JSONStorage) listHostsInternal() ([]*models.Host, error) {
	var data struct {
		Hosts []*models.Host `json:"hosts"`
	}

	if err := s.readJSON(s.getHostsPath(), &data); err != nil {
		return nil, fmt.Errorf("failed to read hosts: %w", err)
	}

	return data.Hosts, nil
}

func (s *JSONStorage) GetHost(ctx context.Context, id string) (*models.Host, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	hosts, err := s.listHostsInternal()
	if err != nil {
		return nil, err
	}

	for _, host := range hosts {
		if host.ID == id {
			return host, nil
		}
	}

	return nil, fmt.Errorf("host not found: %s", id)
}

func (s *JSONStorage) SaveHost(ctx context.Context, host *models.Host) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	hosts, err := s.listHostsInternal()
	if err != nil {
		return err
	}

	found := false
	for i, h := range hosts {
		if h.ID == host.ID {
			hosts[i] = host
			found = true
			break
		}
	}

	if !found {
		hosts = append(hosts, host)
	}

	return s.replaceHosts(hosts)
}

func (s *JSONStorage) DeleteHost(ctx context.Context, id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	hosts, err := s.listHostsInternal()
	if err != nil {
		return err
	}

	for i, host := range hosts {
		if host.ID == id {
			hosts = append(hosts[:i], hosts[i+1:]...)
			return s.replaceHosts(hosts)
		}
	}

	return fmt.Errorf("host not found: %s", id)
}

func (s *JSONStorage) replaceHosts(hosts []*models.Host) error {
	var data struct {
		Hosts []*models.Host `json:"hosts"`
	}
	data.Hosts = hosts

	bytes, err := json.MarshalIndent(data, "", "  ")
	if err != nil {
		return fmt.Errorf("failed to marshal hosts: %w", err)
	}

	return os.WriteFile(s.getHostsPath(), bytes, 0600)
}

// ============ KeyPair Operations ============

func (s *JSONStorage) getKeysPath() string {
	return filepath.Join(s.dataDir, "keys.json")
}

func (s *JSONStorage) ListKeyPairs(ctx context.Context) ([]*models.KeyPair, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.listKeyPairsInternal()
}

// listKeyPairsInternal reads key pairs WITHOUT locking (caller must hold lock)
func (s *JSONStorage) listKeyPairsInternal() ([]*models.KeyPair, error) {
	var data struct {
		KeyPairs []*models.KeyPair `json:"key_pairs"`
	}

	if err := s.readJSON(s.getKeysPath(), &data); err != nil {
		return nil, fmt.Errorf("failed to read key pairs: %w", err)
	}

	return data.KeyPairs, nil
}

func (s *JSONStorage) GetKeyPair(ctx context.Context, id string) (*models.KeyPair, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	keys, err := s.listKeyPairsInternal()
	if err != nil {
		return nil, err
	}

	for _, key := range keys {
		if key.ID == id {
			return key, nil
		}
	}

	return nil, fmt.Errorf("key pair not found: %s", id)
}

func (s *JSONStorage) SaveKeyPair(ctx context.Context, keyPair *models.KeyPair) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	keys, err := s.listKeyPairsInternal()
	if err != nil {
		return err
	}

	found := false
	for i, k := range keys {
		if k.ID == keyPair.ID {
			keys[i] = keyPair
			found = true
			break
		}
	}

	if !found {
		keys = append(keys, keyPair)
	}

	return s.replaceKeyPairs(keys)
}

func (s *JSONStorage) DeleteKeyPair(ctx context.Context, id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	keys, err := s.listKeyPairsInternal()
	if err != nil {
		return err
	}

	for i, key := range keys {
		if key.ID == id {
			keys = append(keys[:i], keys[i+1:]...)
			return s.replaceKeyPairs(keys)
		}
	}

	return fmt.Errorf("key pair not found: %s", id)
}

func (s *JSONStorage) replaceKeyPairs(keys []*models.KeyPair) error {
	var data struct {
		KeyPairs []*models.KeyPair `json:"key_pairs"`
	}
	data.KeyPairs = keys

	bytes, err := json.MarshalIndent(data, "", "  ")
	if err != nil {
		return fmt.Errorf("failed to marshal key pairs: %w", err)
	}

	return os.WriteFile(s.getKeysPath(), bytes, 0600)
}

// ============ Snippet Operations ============

func (s *JSONStorage) getSnippetsPath() string {
	return filepath.Join(s.dataDir, "snippets.json")
}

func (s *JSONStorage) ListSnippets(ctx context.Context) ([]*models.Snippet, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.listSnippetsInternal()
}

// listSnippetsInternal reads snippets WITHOUT locking (caller must hold lock)
func (s *JSONStorage) listSnippetsInternal() ([]*models.Snippet, error) {
	var data struct {
		Snippets []*models.Snippet `json:"snippets"`
	}

	if err := s.readJSON(s.getSnippetsPath(), &data); err != nil {
		return nil, fmt.Errorf("failed to read snippets: %w", err)
	}

	return data.Snippets, nil
}

func (s *JSONStorage) GetSnippet(ctx context.Context, id string) (*models.Snippet, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	snippets, err := s.listSnippetsInternal()
	if err != nil {
		return nil, err
	}

	for _, snippet := range snippets {
		if snippet.ID == id {
			return snippet, nil
		}
	}

	return nil, fmt.Errorf("snippet not found: %s", id)
}

func (s *JSONStorage) SaveSnippet(ctx context.Context, snippet *models.Snippet) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	snippets, err := s.listSnippetsInternal()
	if err != nil {
		return err
	}

	found := false
	for i, sn := range snippets {
		if sn.ID == snippet.ID {
			snippets[i] = snippet
			found = true
			break
		}
	}

	if !found {
		snippets = append(snippets, snippet)
	}

	return s.replaceSnippets(snippets)
}

func (s *JSONStorage) DeleteSnippet(ctx context.Context, id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	snippets, err := s.listSnippetsInternal()
	if err != nil {
		return err
	}

	for i, snippet := range snippets {
		if snippet.ID == id {
			snippets = append(snippets[:i], snippets[i+1:]...)
			return s.replaceSnippets(snippets)
		}
	}

	return fmt.Errorf("snippet not found: %s", id)
}

func (s *JSONStorage) replaceSnippets(snippets []*models.Snippet) error {
	var data struct {
		Snippets []*models.Snippet `json:"snippets"`
	}
	data.Snippets = snippets

	bytes, err := json.MarshalIndent(data, "", "  ")
	if err != nil {
		return fmt.Errorf("failed to marshal snippets: %w", err)
	}

	return os.WriteFile(s.getSnippetsPath(), bytes, 0600)
}

// ============ Export/Import ============

func (s *JSONStorage) ExportData(ctx context.Context) (*ExportData, error) {
	hosts, err := s.ListHosts(ctx)
	if err != nil {
		return nil, err
	}

	keys, err := s.ListKeyPairs(ctx)
	if err != nil {
		return nil, err
	}

	snippets, err := s.ListSnippets(ctx)
	if err != nil {
		return nil, err
	}

	return &ExportData{
		Hosts:    hosts,
		KeyPairs: keys,
		Snippets: snippets,
	}, nil
}

func (s *JSONStorage) ImportData(ctx context.Context, data *ExportData, strategy MergeStrategy) error {
	switch strategy {
	case MergeStrategyReplace:
		if err := s.replaceHosts(data.Hosts); err != nil {
			return err
		}
		if err := s.replaceKeyPairs(data.KeyPairs); err != nil {
			return err
		}
		if err := s.replaceSnippets(data.Snippets); err != nil {
			return err
		}

	case MergeStrategyMerge:
		if err := s.mergeHosts(ctx, data.Hosts); err != nil {
			return err
		}
		if err := s.mergeKeyPairs(ctx, data.KeyPairs); err != nil {
			return err
		}
		if err := s.mergeSnippets(ctx, data.Snippets); err != nil {
			return err
		}

	default:
		return fmt.Errorf("unknown merge strategy: %s", strategy)
	}

	return nil
}

func (s *JSONStorage) mergeHosts(ctx context.Context, newHosts []*models.Host) error {
	existingHosts, err := s.ListHosts(ctx)
	if err != nil {
		return err
	}

	existingIDs := make(map[string]bool)
	for _, host := range existingHosts {
		existingIDs[host.ID] = true
	}

	for _, newHost := range newHosts {
		if !existingIDs[newHost.ID] {
			existingHosts = append(existingHosts, newHost)
		}
	}

	return s.replaceHosts(existingHosts)
}

func (s *JSONStorage) mergeKeyPairs(ctx context.Context, newKeys []*models.KeyPair) error {
	existingKeys, err := s.ListKeyPairs(ctx)
	if err != nil {
		return err
	}

	existingIDs := make(map[string]bool)
	for _, key := range existingKeys {
		existingIDs[key.ID] = true
	}

	for _, newKey := range newKeys {
		if !existingIDs[newKey.ID] {
			existingKeys = append(existingKeys, newKey)
		}
	}

	return s.replaceKeyPairs(existingKeys)
}

func (s *JSONStorage) mergeSnippets(ctx context.Context, newSnippets []*models.Snippet) error {
	existingSnippets, err := s.ListSnippets(ctx)
	if err != nil {
		return err
	}

	existingIDs := make(map[string]bool)
	for _, snippet := range existingSnippets {
		existingIDs[snippet.ID] = true
	}

	for _, newSnippet := range newSnippets {
		if !existingIDs[newSnippet.ID] {
			existingSnippets = append(existingSnippets, newSnippet)
		}
	}

	return s.replaceSnippets(existingSnippets)
}

// ============ Helpers ============

func (s *JSONStorage) readJSON(path string, v interface{}) error {
	data, err := os.ReadFile(path)
	if err != nil {
		return err
	}
	return json.Unmarshal(data, v)
}