HostKeeper/pkg/storage/json_storage.go

package storage

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

	"github.com/google/uuid"

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

// JSONStorage implements Storage interface using JSON files
type JSONStorage struct {
	dataDir  string
	password string // master password for encryption (empty = no encryption)
	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
}

// SetPassword sets the master password for encryption/decryption
func (s *JSONStorage) SetPassword(password string) {
	s.password = password
}

// GetPassword returns the current master password
func (s *JSONStorage) GetPassword() string {
	return s.password
}

// IsEncrypted returns whether encryption is enabled
func (s *JSONStorage) IsEncrypted() bool {
	return s.password != ""
}

// IsDataEncrypted checks if the data files are actually encrypted
func (s *JSONStorage) IsDataEncrypted() bool {
	path := filepath.Join(s.dataDir, "hosts.json")
	data, err := os.ReadFile(path)
	if err != nil {
		return false
	}
	return crypto.IsEncrypted(string(data))
}

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()

	if host.ID == "" {
		host.ID = uuid.New().String()
	}
	if host.CreatedAt.IsZero() {
		host.CreatedAt = time.Now()
	}
	host.UpdatedAt = time.Now()

	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)
	}

	// Encrypt if password is set
	if s.password != "" {
		encrypted, err := crypto.Encrypt(bytes, s.password)
		if err != nil {
			return fmt.Errorf("failed to encrypt hosts: %w", err)
		}
		bytes = []byte(encrypted)
	}

	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()

	if keyPair.ID == "" {
		keyPair.ID = uuid.New().String()
	}
	if keyPair.CreatedAt.IsZero() {
		keyPair.CreatedAt = time.Now()
	}
	keyPair.UpdatedAt = time.Now()

	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)
	}

	// Encrypt if password is set
	if s.password != "" {
		encrypted, err := crypto.Encrypt(bytes, s.password)
		if err != nil {
			return fmt.Errorf("failed to encrypt key pairs: %w", err)
		}
		bytes = []byte(encrypted)
	}

	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()

	if snippet.ID == "" {
		snippet.ID = uuid.New().String()
	}
	if snippet.CreatedAt.IsZero() {
		snippet.CreatedAt = time.Now()
	}
	snippet.UpdatedAt = time.Now()

	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)
	}

	// Encrypt if password is set
	if s.password != "" {
		encrypted, err := crypto.Encrypt(bytes, s.password)
		if err != nil {
			return fmt.Errorf("failed to encrypt snippets: %w", err)
		}
		bytes = []byte(encrypted)
	}

	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
	}

	// Decrypt if password is set and data looks encrypted
	if s.password != "" && crypto.IsEncrypted(string(data)) {
		decrypted, err := crypto.Decrypt(string(data), s.password)
		if err != nil {
			return fmt.Errorf("decryption failed: %w", err)
		}
		data = decrypted
	}

	return json.Unmarshal(data, v)
}