backup/cmd/restore.go

package cmd

import (
	"archive/tar"
	"compress/gzip"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/charmbracelet/bubbles/progress"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/spf13/cobra"
)

var (
	restoreSource string
	restoreTarget string
)

var restoreCmd = &cobra.Command{
	Use:   "restore",
	Short: "Restore a backup",
	Long:  "Restore a backup from a local archive, P2P stream, or HTTP server.",
	Run: func(cmd *cobra.Command, args []string) {
		model, err := initialRestoreModel(restoreSource, restoreTarget)
		if err != nil {
			fmt.Println("Error:", err)
			os.Exit(1)
		}

		p := tea.NewProgram(model, tea.WithAltScreen())
		if _, err := p.Run(); err != nil {
			fmt.Println("Error:", err)
			os.Exit(1)
		}
	},
}

type restoreModel struct {
	progress       progress.Model
	source         string
	target         string
	totalBytes     int64
	processed      int64
	done           bool
	err            error
	startTime      time.Time
	progressCh     chan int64
	isDirectory    bool
	httpSourceInfo *ServerTargetInfo
	isP2PSource    bool
	statusMessage  string
}

// restoreSourceReadyMsg is sent when the source is ready to be read.
// It carries the reader, the total size, and whether it's a directory backup.
type restoreSourceReadyMsg struct {
	reader     io.ReadCloser
	totalBytes int64
	isDir      bool
}

type restoreProgressMsg int64
type restoreDoneMsg struct{}
type restoreErrorMsg struct{ err error }

func initialRestoreModel(src, dst string) (*restoreModel, error) {
	p := progress.New(
		progress.WithDefaultGradient(),
		progress.WithWidth(40),
	)

	var totalBytes int64 = 1 // Placeholder, will be updated.
	var isDir, isP2P bool
	targetPath := dst
	status := "Initializing..."

	httpInfo, isHttp := parseServerTarget(src)

	if isHttp {
		isDir = strings.HasSuffix(httpInfo.Filename, ".tar.gz")
		status = fmt.Sprintf("Connecting to %s...", httpInfo.Address)
		src = httpInfo.URL
	} else if isNetworkAddress(src) {
		isP2P = true
		status = fmt.Sprintf("Listening on %s...", src)
	} else {
		fileInfo, err := os.Stat(src)
		if err != nil {
			return nil, fmt.Errorf("failed to stat source %s: %w", src, err)
		}
		totalBytes = fileInfo.Size()
		isDir = strings.HasSuffix(src, ".tar.gz")
	}

	if isDir && targetPath == "" {
		// For directory restores, we can default to the current directory.
		// We do this check here, but also again after getting the P2P type.
		var err error
		targetPath, err = os.Getwd()
		if err != nil {
			return nil, fmt.Errorf("failed to get current directory: %w", err)
		}
	} else if !isDir && targetPath == "" && !isP2P {
		// Target is mandatory for file/disk restores (except P2P, where we can't know the type yet).
		return nil, fmt.Errorf("target flag -t is required for this restoration")
	}

	return &restoreModel{
		progress:       p,
		source:         src,
		target:         targetPath,
		totalBytes:     totalBytes,
		startTime:      time.Now(),
		progressCh:     make(chan int64, 100),
		isDirectory:    isDir,
		httpSourceInfo: httpInfo,
		isP2PSource:    isP2P,
		statusMessage:  status,
	}, nil
}

func (m *restoreModel) Init() tea.Cmd {
	return m.fetchSourceCmd
}

func (m *restoreModel) fetchSourceCmd() tea.Msg {
	if m.httpSourceInfo != nil {
		// New HTTP Mode
		client := &http.Client{}
		resp, err := client.Get(m.source)
		if err != nil {
			return restoreErrorMsg{fmt.Errorf("failed to connect to server: %w", err)}
		}
		if resp.StatusCode != http.StatusOK {
			bodyBytes, _ := io.ReadAll(resp.Body)
			resp.Body.Close()
			return restoreErrorMsg{fmt.Errorf("server returned non-200 status: %s\n%s", resp.Status, string(bodyBytes))}
		}
		return restoreSourceReadyMsg{reader: resp.Body, totalBytes: resp.ContentLength, isDir: m.isDirectory}

	} else if m.isP2PSource {
		// Old P2P Mode
		listener, err := net.Listen("tcp", m.source)
		if err != nil {
			return restoreErrorMsg{fmt.Errorf("failed to listen on %s: %w", m.source, err)}
		}
		// Accept is blocking, so it must be in a command.
		conn, err := listener.Accept()
		listener.Close() // Close listener after one connection.
		if err != nil {
			return restoreErrorMsg{fmt.Errorf("failed to accept connection: %w", err)}
		}
		var size int64
		if err := binary.Read(conn, binary.BigEndian, &size); err != nil {
			conn.Close()
			return restoreErrorMsg{fmt.Errorf("failed to read backup size: %w", err)}
		}
		typeByte := make([]byte, 1)
		if _, err := io.ReadFull(conn, typeByte); err != nil {
			conn.Close()
			return restoreErrorMsg{fmt.Errorf("failed to read backup type: %w", err)}
		}
		return restoreSourceReadyMsg{reader: conn, totalBytes: size, isDir: typeByte[0] == 1}
	}

	// Local File Mode
	fileInfo, err := os.Stat(m.source)
	if err != nil {
		return restoreErrorMsg{err}
	}
	file, err := os.Open(m.source)
	if err != nil {
		return restoreErrorMsg{err}
	}
	return restoreSourceReadyMsg{reader: file, totalBytes: fileInfo.Size(), isDir: m.isDirectory}
}

func (m *restoreModel) processRestoreCmd(r io.ReadCloser) tea.Cmd {
	return func() tea.Msg {
		defer r.Close()
		progressReader := &progressTracker{
			Reader:     r,
			progressCh: m.progressCh,
		}

		var restoreErr error
		if m.isDirectory {
			restoreErr = m.restoreDirectory(progressReader)
		} else {
			restoreErr = m.restoreFileOrDisk(progressReader)
		}

		if restoreErr != nil {
			m.err = restoreErr
			m.progressCh <- -1
		} else {
			m.progressCh <- -2 // Done signal
		}
		return nil
	}
}

func (m *restoreModel) restoreFileOrDisk(r io.Reader) error {
	gzipReader, err := gzip.NewReader(r)
	if err != nil {
		return err
	}
	defer gzipReader.Close()

	var dstFile *os.File
	if isBlockDevice(m.target) {
		dstFile, err = os.OpenFile(m.target, os.O_WRONLY|os.O_SYNC, 0)
	} else {
		if err := os.MkdirAll(filepath.Dir(m.target), 0755); err != nil {
			return err
		}
		dstFile, err = os.Create(m.target)
	}
	if err != nil {
		return err
	}
	defer dstFile.Close()

	_, err = io.Copy(dstFile, gzipReader)
	if err != nil && err != io.EOF {
		return err
	}
	return dstFile.Sync()
}

func (m *restoreModel) restoreDirectory(r io.Reader) error {
	gzipReader, err := gzip.NewReader(r)
	if err != nil {
		return err
	}
	defer gzipReader.Close()

	tarReader := tar.NewReader(gzipReader)

	for {
		header, err := tarReader.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		targetPath := filepath.Join(m.target, header.Name)

		switch header.Typeflag {
		case tar.TypeDir:
			if err := os.MkdirAll(targetPath, os.FileMode(header.Mode)); err != nil {
				return err
			}
		case tar.TypeReg:
			if err := os.MkdirAll(filepath.Dir(targetPath), 0755); err != nil {
				return err
			}
			outFile, err := os.OpenFile(targetPath, os.O_CREATE|os.O_WRONLY, os.FileMode(header.Mode))
			if err != nil {
				return err
			}
			_, copyErr := io.Copy(outFile, tarReader)
			closeErr := outFile.Close()
			if copyErr != nil {
				return copyErr
			}
			if closeErr != nil {
				return closeErr
			}
		}
	}
	return nil
}

func (m *restoreModel) restoreProgressListener() tea.Msg {
	n := <-m.progressCh
	if n < 0 {
		if n == -1 {
			if m.err == nil {
				m.err = fmt.Errorf("unknown restoration error")
			}
			return restoreErrorMsg{m.err}
		}
		return restoreDoneMsg{}
	}
	m.processed = n
	return restoreProgressMsg(n)
}

func (m *restoreModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case tea.KeyMsg:
		if msg.Type == tea.KeyCtrlC {
			return m, tea.Quit
		}
		return m, nil

	case tea.WindowSizeMsg:
		m.progress.Width = msg.Width - 4
		return m, nil

	case restoreSourceReadyMsg:
		m.totalBytes = msg.totalBytes
		if m.totalBytes <= 0 {
			m.totalBytes = 1
		}
		m.isDirectory = msg.isDir // Update directory type, especially for P2P.

		// Final check for target directory now that we know the type for sure.
		if m.isDirectory && m.target == "" {
			wd, err := os.Getwd()
			if err != nil {
				return m, func() tea.Msg { return restoreErrorMsg{err} }
			}
			m.target = wd
		} else if !m.isDirectory && m.target == "" {
			err := fmt.Errorf("target flag -t is required for file/disk restoration")
			return m, func() tea.Msg { return restoreErrorMsg{err} }
		}

		m.statusMessage = "Receiving data..."
		return m, tea.Batch(
			m.processRestoreCmd(msg.reader),
			m.restoreProgressListener,
		)

	case restoreProgressMsg:
		progressVal := float64(msg) / float64(m.totalBytes)
		if progressVal > 1.0 {
			progressVal = 1.0
		}
		cmd := m.progress.SetPercent(progressVal)
		return m, tea.Batch(cmd, m.restoreProgressListener)

	case restoreDoneMsg:
		m.done = true
		cmd := m.progress.SetPercent(1.0)
		return m, tea.Batch(cmd, tea.Quit)

	case restoreErrorMsg:
		m.err = msg.err
		return m, tea.Quit
	}
	return m, nil
}

func (m *restoreModel) View() string {
	if m.err != nil {
		return fmt.Sprintf("\n❌ Error: %v\n\n", m.err)
	}

	if m.done {
		duration := time.Since(m.startTime)
		restoreType := "File/Disk"
		if m.isDirectory {
			restoreType = "Directory"
		}
		if m.httpSourceInfo != nil || m.isP2PSource {
			restoreType += " network"
		}
		return fmt.Sprintf("\n✅ %s restoration complete!\n\n"+
			"Source: %s\n"+
			"Target: %s\n"+
			"Time: %v\n\n",
			restoreType, m.source, m.target,
			duration.Round(time.Millisecond))
	}

	if m.processed == 0 {
		return m.statusMessage + "\n\nPress Ctrl+C to cancel"
	}

	operation := "Restoring"
	if m.httpSourceInfo != nil {
		operation = "Downloading"
	} else if m.isP2PSource {
		operation = "Receiving"
	}

	title := fmt.Sprintf("%s from %s → %s", operation, m.source, m.target)
	progressVal := float64(m.processed) / float64(m.totalBytes)
	progressView := m.progress.ViewAs(progressVal)
	stats := fmt.Sprintf("%s / %s (%.1f%%)", formatBytes(m.processed), formatBytes(m.totalBytes), progressVal*100)

	return lipgloss.JoinVertical(lipgloss.Left,
		title,
		"",
		progressView,
		"",
		stats,
		"\nPress Ctrl+C to cancel",
	)
}

func isBlockDevice(path string) bool {
	if runtime.GOOS != "linux" {
		return false
	}
	fileInfo, err := os.Stat(path)
	if err != nil {
		return false
	}
	stat, ok := fileInfo.Sys().(*syscall.Stat_t)
	if !ok {
		return false
	}
	return (stat.Mode & syscall.S_IFMT) == syscall.S_IFBLK
}

func init() {
	rootCmd.AddCommand(restoreCmd)
	restoreCmd.Flags().StringVarP(&restoreSource, "source", "s", "", "Source archive, P2P address (e.g., 0.0.0.0:8080), or HTTP URL (e.g., 127.0.0.1:8080/backup.gz) (required)")
	restoreCmd.Flags().StringVarP(&restoreTarget, "target", "t", "", "Target file, disk, or directory. Required for non-directory restores.")
	restoreCmd.MarkFlagRequired("source")
}