package main

import (
	"context"

	"github.com/restic/restic/internal/backend"
	"github.com/restic/restic/internal/errors"
	"github.com/restic/restic/internal/restic"
	"github.com/restic/restic/internal/walker"

	"github.com/spf13/cobra"
)

var cmdRepairSnapshots = &cobra.Command{
	Use:   "snapshots [flags] [snapshot ID] [...]",
	Short: "Repair snapshots",
	Long: `
The "repair snapshots" command allows to repair broken snapshots.
It scans the given snapshots and generates new ones where
damaged tress and file contents are removed.
If the broken snapshots are deleted, a prune run will
be able to refit the repository.

The command depends on a good state of the index, so if
there are inaccurancies in the index, make sure to run
"repair index" before!


WARNING:
========
Repairing and deleting broken snapshots causes data loss!
It will remove broken dirs and modify broken files in
the modified snapshots.

If the contents of directories and files are still available,
the better option is to redo a backup which in that case is 
able to "heal" already present snapshots.
Only use this command if you need to recover an old and
broken snapshot!

EXIT STATUS
===========

Exit status is 0 if the command was successful, and non-zero if there was any error.
`,
	DisableAutoGenTag: true,
	RunE: func(cmd *cobra.Command, args []string) error {
		return runRepairSnapshots(cmd.Context(), globalOptions, repairSnapshotOptions, args)
	},
}

// RepairOptions collects all options for the repair command.
type RepairOptions struct {
	DryRun bool
	Forget bool

	restic.SnapshotFilter
}

var repairSnapshotOptions RepairOptions

func init() {
	cmdRepair.AddCommand(cmdRepairSnapshots)
	flags := cmdRepairSnapshots.Flags()

	flags.BoolVarP(&repairSnapshotOptions.DryRun, "dry-run", "n", false, "do not do anything, just print what would be done")
	flags.BoolVarP(&repairSnapshotOptions.Forget, "forget", "", false, "remove original snapshots after creating new ones")

	initMultiSnapshotFilter(flags, &repairSnapshotOptions.SnapshotFilter, true)
}

func runRepairSnapshots(ctx context.Context, gopts GlobalOptions, opts RepairOptions, args []string) error {
	repo, err := OpenRepository(ctx, globalOptions)
	if err != nil {
		return err
	}

	if !opts.DryRun {
		var lock *restic.Lock
		var err error
		lock, ctx, err = lockRepoExclusive(ctx, repo, gopts.RetryLock, gopts.JSON)
		defer unlockRepo(lock)
		if err != nil {
			return err
		}
	} else {
		repo.SetDryRun()
	}

	snapshotLister, err := backend.MemorizeList(ctx, repo.Backend(), restic.SnapshotFile)
	if err != nil {
		return err
	}

	if err := repo.LoadIndex(ctx); err != nil {
		return err
	}

	// Three error cases are checked:
	// - tree is a nil tree (-> will be replaced by an empty tree)
	// - trees which cannot be loaded (-> the tree contents will be removed)
	// - files whose contents are not fully available  (-> file will be modified)
	rewriter := walker.NewTreeRewriter(walker.RewriteOpts{
		RewriteNode: func(node *restic.Node, path string) *restic.Node {
			if node.Type != "file" {
				return node
			}

			ok := true
			var newContent restic.IDs
			var newSize uint64
			// check all contents and remove if not available
			for _, id := range node.Content {
				if size, found := repo.LookupBlobSize(id, restic.DataBlob); !found {
					ok = false
				} else {
					newContent = append(newContent, id)
					newSize += uint64(size)
				}
			}
			if !ok {
				if newSize == 0 {
					Printf("removed defective file '%v'\n", path+node.Name)
					node = nil
				} else {
					Printf("repaired defective file '%v'\n", path+node.Name)
					node.Content = newContent
					node.Size = newSize
				}
			}
			return node
		},
		RewriteFailedTree: func(nodeID restic.ID, path string, _ error) (restic.ID, error) {
			if path == "/" {
				// remove snapshots with invalid root node
				return restic.ID{}, nil
			}
			// If a subtree fails to load, remove it
			Printf("removed defective dir '%v'", path)
			emptyID, err := restic.SaveTree(ctx, repo, &restic.Tree{})
			if err != nil {
				return restic.ID{}, err
			}
			return emptyID, nil
		},
		AllowUnstableSerialization: true,
	})

	changedCount := 0
	for sn := range FindFilteredSnapshots(ctx, snapshotLister, repo, &opts.SnapshotFilter, args) {
		Verbosef("\nsnapshot %s of %v at %s)\n", sn.ID().Str(), sn.Paths, sn.Time)
		changed, err := filterAndReplaceSnapshot(ctx, repo, sn,
			func(ctx context.Context, sn *restic.Snapshot) (restic.ID, error) {
				return rewriter.RewriteTree(ctx, repo, "/", *sn.Tree)
			}, opts.DryRun, opts.Forget, "repaired")
		if err != nil {
			return errors.Fatalf("unable to rewrite snapshot ID %q: %v", sn.ID().Str(), err)
		}
		if changed {
			changedCount++
		}
	}

	Verbosef("\n")
	if changedCount == 0 {
		if !opts.DryRun {
			Verbosef("no snapshots were modified\n")
		} else {
			Verbosef("no snapshots would be modified\n")
		}
	} else {
		if !opts.DryRun {
			Verbosef("modified %v snapshots\n", changedCount)
		} else {
			Verbosef("would modify %v snapshots\n", changedCount)
		}
	}

	return nil
}