zfs/internal/client/client.go

// Package client provides ZFS snapshot backup client functionality.
// It handles creating snapshots and uploading them to a remote server.
package client

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/mistifyio/go-zfs"
	"github.com/pierrec/lz4/v4"
)

// Client handles snapshot backup operations to a remote server.
// It manages creating local ZFS snapshots and transmitting them
// to the backup server via HTTP or SSH.
type Client struct {
	config *Config
}

// New creates a new Client instance with the provided configuration.
func New(config *Config) *Client {
	return &Client{config: config}
}

// CreateSnapshot creates a local ZFS snapshot of the configured dataset.
// The snapshot is named with a timestamp for easy identification.
// Returns the created snapshot dataset or an error.
func (c *Client) CreateSnapshot() (*zfs.Dataset, error) {
	// Get the local dataset
	ds, err := zfs.GetDataset(c.config.LocalDataset)
	if err != nil {
		return nil, fmt.Errorf("failed to get dataset: %v", err)
	}

	// Generate snapshot name with timestamp
	timestamp := time.Now().Format("2006-01-02_15:04:05")
	snapshotName := fmt.Sprintf("backup_%s", timestamp)

	// Create the snapshot
	snapshot, err := ds.Snapshot(snapshotName, false)
	if err != nil {
		return nil, fmt.Errorf("failed to create snapshot: %v", err)
	}

	fmt.Printf("✓ Created local snapshot: %s@%s\n", c.config.LocalDataset, snapshotName)
	return snapshot, nil
}

// GetSnapshotSize returns the used size of a snapshot in bytes.
func (c *Client) GetSnapshotSize(snapshot *zfs.Dataset) int64 {
	return int64(snapshot.Used)
}

// SendSnapshot sends a snapshot to the backup server.
// It first requests upload authorization, then streams the snapshot
// using the appropriate method (S3 or ZFS receive).
func (c *Client) SendSnapshot(snapshot *zfs.Dataset) error {
	estimatedSize := c.GetSnapshotSize(snapshot)

	// Request upload authorization from server
	uploadReq := map[string]interface{}{
		"client_id":      c.config.ClientID,
		"api_key":        c.config.APIKey,
		"dataset_name":   c.config.LocalDataset,
		"timestamp":      time.Now().Format(time.RFC3339),
		"compressed":     c.config.Compress,
		"estimated_size": estimatedSize,
	}

	reqBody, _ := json.Marshal(uploadReq)
	resp, err := http.Post(c.config.ServerURL+"/upload", "application/json", bytes.NewBuffer(reqBody))
	if err != nil {
		return fmt.Errorf("failed to request upload: %v", err)
	}
	defer resp.Body.Close()

	// Parse server response
	var uploadResp struct {
		Success      bool   `json:"success"`
		Message      string `json:"message"`
		UploadURL    string `json:"upload_url"`
		UploadMethod string `json:"upload_method"`
		StorageKey   string `json:"storage_key"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&uploadResp); err != nil {
		return fmt.Errorf("failed to decode response: %v", err)
	}

	if !uploadResp.Success {
		return fmt.Errorf("upload not authorized: %s", uploadResp.Message)
	}

	fmt.Printf("→ Upload authorized\n")
	fmt.Printf("  Method: %s\n", uploadResp.UploadMethod)
	fmt.Printf("  Storage key: %s\n", uploadResp.StorageKey)

	// Choose upload method based on server response
	if uploadResp.UploadMethod == "s3" {
		return c.streamToS3(snapshot, uploadResp.UploadURL, uploadResp.StorageKey)
	}
	return c.sendViaZFS(snapshot, uploadResp.StorageKey)
}

// streamToS3 streams a ZFS snapshot to S3 storage via HTTP.
// The snapshot is optionally compressed with LZ4 before transmission.
func (c *Client) streamToS3(snapshot *zfs.Dataset, uploadURL, storageKey string) error {
	fmt.Printf("→ Streaming snapshot to S3...\n")

	// Create ZFS send command
	cmd := exec.Command("zfs", "send", snapshot.Name)
	zfsOut, err := cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("failed to create pipe: %v", err)
	}

	if err := cmd.Start(); err != nil {
		return fmt.Errorf("failed to start zfs send: %v", err)
	}

	var reader io.Reader = zfsOut

	// Apply LZ4 compression if enabled
	if c.config.Compress {
		fmt.Printf("  Compressing with LZ4...\n")
		pr, pw := io.Pipe()
		lz4Writer := lz4.NewWriter(pw)
		lz4Writer.Apply(lz4.BlockSizeOption(lz4.BlockSize(4 * 1024 * 1024))) // 4MB blocks for better performance

		go func() {
			// Copy zfs output to LZ4 writer
			io.Copy(lz4Writer, zfsOut)
			// Close LZ4 writer first to flush, then close pipe
			lz4Writer.Close()
			pw.Close()
		}()

		reader = pr
	}

	// Create HTTP request
	req, err := http.NewRequest("POST", c.config.ServerURL+uploadURL, reader)
	if err != nil {
		return fmt.Errorf("failed to create request: %v", err)
	}

	// Set required headers
	req.Header.Set("X-API-Key", c.config.APIKey)
	req.Header.Set("X-Storage-Key", storageKey)
	req.Header.Set("X-Dataset-Name", c.config.LocalDataset)
	req.Header.Set("X-Compressed", fmt.Sprintf("%v", c.config.Compress))
	req.Header.Set("Content-Type", "application/octet-stream")

	// Send request with no timeout for large uploads
	client := &http.Client{
		Timeout: 0,
	}

	httpResp, err := client.Do(req)
	if err != nil {
		cmd.Process.Kill()
		return fmt.Errorf("failed to upload: %v", err)
	}
	defer httpResp.Body.Close()

	if httpResp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(httpResp.Body)
		return fmt.Errorf("upload failed with status %d: %s", httpResp.StatusCode, body)
	}

	if err := cmd.Wait(); err != nil {
		return fmt.Errorf("zfs send failed: %v", err)
	}

	// Parse response
	var result struct {
		Success bool   `json:"success"`
		Message string `json:"message"`
		Size    int64  `json:"size"`
	}

	if err := json.NewDecoder(httpResp.Body).Decode(&result); err != nil {
		return fmt.Errorf("failed to decode response: %v", err)
	}

	if !result.Success {
		return fmt.Errorf("upload failed: %s", result.Message)
	}

	fmt.Printf("✓ Snapshot uploaded successfully!\n")
	fmt.Printf("  Size: %.2f MB\n", float64(result.Size)/(1024*1024))

	return nil
}

// sendViaZFS sends a snapshot via traditional ZFS send/receive over SSH.
// This method is used when the server uses local ZFS storage.
func (c *Client) sendViaZFS(snapshot *zfs.Dataset, receivePath string) error {
	fmt.Printf("→ Sending via ZFS send/receive...\n")

	// Extract server host from URL
	serverHost := c.config.ServerURL
	if len(serverHost) > 7 && strings.HasPrefix(serverHost, "http://") {
		serverHost = serverHost[7:]
	} else if len(serverHost) > 8 && strings.HasPrefix(serverHost, "https://") {
		serverHost = serverHost[8:]
	}

	// Remove port if present
	if idx := strings.LastIndex(serverHost, ":"); idx > 0 {
		serverHost = serverHost[:idx]
	}

	// Execute ZFS send over SSH
	cmd := exec.Command("sh", "-c",
		fmt.Sprintf("zfs send %s | ssh %s 'zfs recv -F %s'",
			snapshot.Name, serverHost, receivePath))

	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	if err := cmd.Run(); err != nil {
		return fmt.Errorf("failed to send snapshot: %v", err)
	}

	fmt.Printf("✓ Snapshot sent successfully!\n")
	return nil
}

// GetStatus retrieves and displays the client's backup status from the server.
// Shows storage usage, quota, and snapshot count.
func (c *Client) GetStatus() error {
	url := fmt.Sprintf("%s/status?client_id=%s&api_key=%s",
		c.config.ServerURL, c.config.ClientID, c.config.APIKey)

	resp, err := http.Get(url)
	if err != nil {
		return fmt.Errorf("failed to get status: %v", err)
	}
	defer resp.Body.Close()

	var status struct {
		Success        bool    `json:"success"`
		TotalSnapshots int     `json:"total_snapshots"`
		UsedBytes      int64   `json:"used_bytes"`
		MaxBytes       int64   `json:"max_bytes"`
		PercentUsed    float64 `json:"percent_used"`
		StorageType    string  `json:"storage_type"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&status); err != nil {
		return fmt.Errorf("failed to decode status: %v", err)
	}

	if !status.Success {
		return fmt.Errorf("status check failed")
	}

	fmt.Printf("\n=== Server Status ===\n")
	fmt.Printf("Storage Type: %s\n", status.StorageType)
	fmt.Printf("Total Snapshots: %d\n", status.TotalSnapshots)
	fmt.Printf("Used: %.2f GB / %.2f GB (%.1f%%)\n",
		float64(status.UsedBytes)/(1024*1024*1024),
		float64(status.MaxBytes)/(1024*1024*1024),
		status.PercentUsed)

	return nil
}

// RequestRotation asks the server to rotate old snapshots.
// This deletes the oldest snapshots to free up space.
func (c *Client) RequestRotation() error {
	reqBody, _ := json.Marshal(map[string]string{
		"client_id": c.config.ClientID,
		"api_key":   c.config.APIKey,
	})

	resp, err := http.Post(c.config.ServerURL+"/rotate", "application/json", bytes.NewBuffer(reqBody))
	if err != nil {
		return fmt.Errorf("failed to request rotation: %v", err)
	}
	defer resp.Body.Close()

	var rotateResp struct {
		Success        bool  `json:"success"`
		DeletedCount   int   `json:"deleted_count"`
		ReclaimedBytes int64 `json:"reclaimed_bytes"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&rotateResp); err != nil {
		return fmt.Errorf("failed to decode response: %v", err)
	}

	if !rotateResp.Success {
		return fmt.Errorf("rotation failed")
	}

	fmt.Printf("✓ Rotation complete\n")
	fmt.Printf("  Deleted: %d snapshots\n", rotateResp.DeletedCount)
	fmt.Printf("  Freed: %.2f GB\n", float64(rotateResp.ReclaimedBytes)/(1024*1024*1024))

	return nil
}

// ServerRotationPolicy represents the rotation policy response from the server
type ServerRotationPolicy struct {
	Success        bool            `json:"success"`
	Message        string          `json:"message"`
	RotationPolicy *SnapshotPolicy `json:"rotation_policy"`
	ServerManaged  bool            `json:"server_managed"`
}

// GetRotationPolicy fetches the rotation policy from the server.
// If the server has a policy configured for this client, it must be used.
// Returns the policy and whether it's server-managed (mandatory).
func (c *Client) GetRotationPolicy() (*ServerRotationPolicy, error) {
	url := fmt.Sprintf("%s/rotation-policy?client_id=%s&api_key=%s",
		c.config.ServerURL, c.config.ClientID, c.config.APIKey)

	resp, err := http.Get(url)
	if err != nil {
		return nil, fmt.Errorf("failed to get rotation policy: %v", err)
	}
	defer resp.Body.Close()

	var policyResp ServerRotationPolicy
	if err := json.NewDecoder(resp.Body).Decode(&policyResp); err != nil {
		return nil, fmt.Errorf("failed to decode response: %v", err)
	}

	if !policyResp.Success {
		return nil, fmt.Errorf("failed to get rotation policy: %s", policyResp.Message)
	}

	return &policyResp, nil
}

// ChangePassword changes the client's API key on the server.
// Requires the current API key for authentication and the new key.
func (c *Client) ChangePassword(newAPIKey string) error {
	reqBody, _ := json.Marshal(map[string]string{
		"client_id":   c.config.ClientID,
		"current_key": c.config.APIKey,
		"new_key":     newAPIKey,
	})

	resp, err := http.Post(c.config.ServerURL+"/client/change-password", "application/json", bytes.NewBuffer(reqBody))
	if err != nil {
		return fmt.Errorf("failed to change password: %v", err)
	}
	defer resp.Body.Close()

	var result struct {
		Success bool   `json:"success"`
		Message string `json:"message"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return fmt.Errorf("failed to decode response: %v", err)
	}

	if !result.Success {
		return fmt.Errorf("failed to change password: %s", result.Message)
	}

	// Update local config with new key
	c.config.APIKey = newAPIKey

	fmt.Printf("✓ Password changed successfully\n")
	return nil
}