backup/cmd/restore.go

package cmd

import (
	"compress/gzip"
	"fmt"
	"io"
	"os"
	"runtime"
	"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 gzip archive to a file or disk",
	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
}

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

func initialRestoreModel(src, dst string) (*restoreModel, error) {
	fileInfo, err := os.Stat(src)
	if err != nil {
		return nil, err
	}

	p := progress.New(
		progress.WithDefaultGradient(),
		progress.WithWidth(40),
	)

	return &restoreModel{
		progress:   p,
		source:     src,
		target:     dst,
		totalBytes: fileInfo.Size(),
		startTime:  time.Now(),
		progressCh: make(chan int64, 100),
	}, nil
}

func (m *restoreModel) Init() tea.Cmd {
	return tea.Batch(
		m.startRestore,
		m.restoreProgressListener,
	)
}

func (m *restoreModel) startRestore() tea.Msg {
	go func() {
		srcFile, err := os.Open(m.source)
		if err != nil {
			m.progressCh <- -1
			return
		}
		defer srcFile.Close()

		// Создаем кастомный reader для отслеживания прогресса
		progressReader := &restoreProgressReader{
			reader:     srcFile,
			total:      m.totalBytes,
			progressCh: m.progressCh,
		}

		// Создаем gzip reader на основе кастомного reader'а
		gzipReader, err := gzip.NewReader(progressReader)
		if err != nil {
			m.progressCh <- -1
			return
		}
		defer gzipReader.Close()

		// Определяем, является ли цель блочным устройством
		var dstFile *os.File
		if isBlockDevice(m.target) {
			// Открываем блочное устройство для записи (без усечения!)
			dstFile, err = os.OpenFile(m.target, os.O_WRONLY|os.O_SYNC, 0)
			if err != nil {
				m.progressCh <- -1
				return
			}
		} else {
			// Обычный файл: создаем/усекаем
			dstFile, err = os.Create(m.target)
			if err != nil {
				m.progressCh <- -1
				return
			}
		}
		defer dstFile.Close()

		// Копируем распакованные данные в целевой файл/устройство
		_, err = io.Copy(dstFile, gzipReader)
		if err != nil {
			m.progressCh <- -1
			return
		}

		// Синхронизируем, чтобы все данные точно записались
		if err := dstFile.Sync(); err != nil {
			m.progressCh <- -1
			return
		}

		m.progressCh <- -2 // Сигнал завершения
	}()

	return nil
}

// restoreProgressReader отслеживает прогресс чтения архива
type restoreProgressReader struct {
	reader     io.Reader
	total      int64
	progressCh chan int64
	read       int64
}

func (r *restoreProgressReader) Read(p []byte) (int, error) {
	n, err := r.reader.Read(p)
	if n > 0 {
		r.read += int64(n)
		r.progressCh <- r.read
	}
	return n, err
}

func (m *restoreModel) restoreProgressListener() tea.Msg {
	select {
	case n := <-m.progressCh:
		switch {
		case n == -1: // Ошибка
			return restoreErrorMsg{fmt.Errorf("restoration error")}
		case n == -2: // Завершение
			return restoreDoneMsg{}
		default: // Прогресс
			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 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
		// Устанавливаем 100% и завершаем
		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)
		return fmt.Sprintf("\n✅ Restoration complete!\n\n"+
			"File: %s → %s\n"+
			"Time: %v\n\n",
			m.source, m.target,
			duration.Round(time.Millisecond))
	}

	title := fmt.Sprintf("Restoring %s → %s", 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",
	)
}

// isBlockDevice проверяет, является ли путь блочным устройством (только для Linux)
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 (required)")
	restoreCmd.Flags().StringVarP(&restoreTarget, "target", "t", "", "Target file or disk (e.g., /dev/sdb) (required)")
	restoreCmd.MarkFlagRequired("source")
	restoreCmd.MarkFlagRequired("target")
}