package main

import (
	"archive/tar"
	"compress/gzip"
	"context"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"strings"
)

// backupArchive writes the files in root to path through a private temporary
// file and a rename. Only directories and regular files enter the archive.
func backupArchive(ctx context.Context, root, path string) (size int64, err error) {
	file, err := os.CreateTemp(filepath.Dir(path), ".onechat-backup-*")
	if err != nil {
		return 0, fmt.Errorf("create archive: %w", err)
	}

	defer func() {
		err = errors.Join(err, file.Close())
		removeErr := os.Remove(file.Name())
		if !errors.Is(removeErr, os.ErrNotExist) {
			err = errors.Join(err, removeErr)
		}
	}()

	compressed := gzip.NewWriter(file)
	archive := tar.NewWriter(compressed)
	err = filepath.WalkDir(root, func(path string, entry fs.DirEntry, err error) error {
		if err != nil {
			return fmt.Errorf("walk backup: %w", err)
		}

		if ctx.Err() != nil {
			return ctx.Err()
		}

		if path == root {
			return nil
		}

		name, err := filepath.Rel(root, path)
		if err != nil {
			return fmt.Errorf("name archive entry: %w", err)
		}

		info, err := entry.Info()
		if err != nil {
			return fmt.Errorf("read backup entry: %w", err)
		}

		if !info.IsDir() && !info.Mode().IsRegular() {
			return fmt.Errorf("unsupported backup entry %q", name)
		}

		header, err := tar.FileInfoHeader(info, "")
		if err != nil {
			return fmt.Errorf("make archive header: %w", err)
		}

		header.Name = filepath.ToSlash(name)
		header.Uid, header.Gid = 0, 0
		header.Uname, header.Gname = "", ""
		err = archive.WriteHeader(header)
		if err != nil {
			return fmt.Errorf("write archive header: %w", err)
		}

		if info.IsDir() {
			return nil
		}

		err = backupCopyFile(ctx, archive, path, info.Size())
		if err != nil {
			return fmt.Errorf("archive file: %w", err)
		}

		return nil
	})
	err = errors.Join(err, archive.Close(), compressed.Close())
	if err != nil {
		return 0, fmt.Errorf("write archive: %w", err)
	}

	err = file.Sync()
	if err != nil {
		return 0, fmt.Errorf("sync archive: %w", err)
	}

	info, err := file.Stat()
	if err != nil {
		return 0, fmt.Errorf("read archive size: %w", err)
	}

	err = os.Rename(file.Name(), path)
	if err != nil {
		return 0, fmt.Errorf("install archive: %w", err)
	}

	return info.Size(), nil
}

// backupExtract restores regular files and directories in an empty staging
// directory. Links, duplicate entries, unsafe names, and special modes fail.
// Directory modes apply last so a read-only directory can hold files.
func backupExtract(ctx context.Context, path, root string) (err error) {
	file, err := os.Open(path)
	if err != nil {
		return fmt.Errorf("open archive: %w", err)
	}

	defer func() {
		err = errors.Join(err, file.Close())
	}()

	compressed, err := gzip.NewReader(file)
	if err != nil {
		return fmt.Errorf("open compressed archive: %w", err)
	}

	defer func() {
		err = errors.Join(err, compressed.Close())
	}()

	archive := tar.NewReader(BackupReader{
		ctx:    ctx,
		reader: compressed,
	})
	seen := make(map[string]bool)

	var directories []*tar.Header

	for {
		var header *tar.Header

		header, err = archive.Next()
		if errors.Is(err, io.EOF) {
			break
		}

		if err != nil {
			return fmt.Errorf("read archive entry: %w", err)
		}

		name := strings.TrimSuffix(header.Name, "/")
		if !filepath.IsLocal(name) || strings.Contains(name, "\\") || filepath.Clean(name) != name || backupExcluded(name) {
			return fmt.Errorf("unsafe archive path %q", header.Name)
		}

		if seen[name] {
			return fmt.Errorf("duplicate archive path %q", name)
		}

		seen[name] = true
		if header.Mode & ^int64(0o777) != 0 {
			return fmt.Errorf("unsafe archive mode for %q", name)
		}

		target := filepath.Join(root, name)
		switch header.Typeflag {
		case tar.TypeDir:
			err = os.MkdirAll(target, 0o700)
			directories = append(directories, header)
		case tar.TypeReg:
			err = os.MkdirAll(filepath.Dir(target), 0o700)
			if err == nil {
				err = backupExtractFile(archive, target, os.FileMode(header.Mode))
			}
		default:
			return fmt.Errorf("unsupported archive entry %q", name)
		}

		if err != nil {
			return fmt.Errorf("restore archive entry %q: %w", name, err)
		}
	}

	// Read through the gzip trailer to check its checksum.
	_, err = io.Copy(io.Discard, compressed)
	if err != nil {
		return fmt.Errorf("check archive checksum: %w", err)
	}

	for i := len(directories) - 1; i >= 0; i-- {
		header := directories[i]
		err = os.Chmod(filepath.Join(root, header.Name), os.FileMode(header.Mode))
		if err != nil {
			return fmt.Errorf("restore directory mode: %w", err)
		}
	}

	return nil
}

// BackupReader checks cancellation between reads of an archive or file.
type BackupReader struct {
	ctx    context.Context
	reader io.Reader
}

func (r BackupReader) Read(data []byte) (int, error) {
	if r.ctx.Err() != nil {
		return 0, r.ctx.Err()
	}

	count, err := r.reader.Read(data)
	if err != nil {
		return count, fmt.Errorf("read backup data: %w", err)
	}

	return count, nil
}

// backupCopyFile copies the size fixed at the snapshot of path. Growth after
// that point is not part of the copy, and a shorter file is an error.
func backupCopyFile(ctx context.Context, out io.Writer, path string, size int64) (err error) {
	file, err := os.Open(path)
	if err != nil {
		return fmt.Errorf("open backup file: %w", err)
	}

	defer func() {
		err = errors.Join(err, file.Close())
	}()

	_, err = io.CopyN(out, BackupReader{
		ctx:    ctx,
		reader: file,
	}, size)
	if err != nil {
		return fmt.Errorf("copy backup file: %w", err)
	}

	return nil
}

func backupExtractFile(in io.Reader, path string, mode os.FileMode) (err error) {
	file, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
	if err != nil {
		return fmt.Errorf("create restored file: %w", err)
	}

	defer func() {
		err = errors.Join(err, file.Close())
	}()

	_, err = io.Copy(file, in)
	if err != nil {
		return fmt.Errorf("write restored file: %w", err)
	}

	err = file.Chmod(mode)
	if err != nil {
		return fmt.Errorf("restore file mode: %w", err)
	}

	return nil
}

// backupExcluded keeps runtime locks, sockets, probe files, and import copies
// out of the snapshot.
func backupExcluded(name string) bool {
	first, _, _ := strings.Cut(filepath.ToSlash(name), "/")
	if first == socketFileName || first == tokenLockFileName || first == "probe" || strings.HasPrefix(first, "agent-backup-") {
		return true
	}

	if filepath.Base(name) == "lock" || (first == imageDirName && strings.HasPrefix(filepath.Base(name), ".image-")) {
		return true
	}

	return false
}