// Licensed to Apache Software Foundation (ASF) under one or more contributor
// license agreements. See the NOTICE file distributed with
// this work for additional information regarding copyright
// ownership. Apache Software Foundation (ASF) licenses this file to you under
// the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.

package backup

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/spf13/cobra"
	"go.uber.org/multierr"

	commonv1 "github.com/apache/skywalking-banyandb/api/proto/banyandb/common/v1"
	"github.com/apache/skywalking-banyandb/banyand/backup/snapshot"
	"github.com/apache/skywalking-banyandb/banyand/internal/storage"
	"github.com/apache/skywalking-banyandb/pkg/config"
	"github.com/apache/skywalking-banyandb/pkg/fs/remote"
	remoteconfig "github.com/apache/skywalking-banyandb/pkg/fs/remote/config"
	"github.com/apache/skywalking-banyandb/pkg/logger"
	banyandbpath "github.com/apache/skywalking-banyandb/pkg/path"
	"github.com/apache/skywalking-banyandb/pkg/version"
)

// NewRestoreCommand creates a new restore command.
func NewRestoreCommand() *cobra.Command {
	logging := logger.Logging{}
	rootCmd := &cobra.Command{
		DisableAutoGenTag: true,
		Version:           version.Build(),
		Short:             "Restore BanyanDB data from remote storage",
		PersistentPreRunE: func(cmd *cobra.Command, _ []string) error {
			if err := config.Load("logging", cmd.Flags()); err != nil {
				return err
			}

			return logger.Init(logging)
		},
	}
	rootCmd.Flags().StringVar(&logging.Env, "logging-env", "prod", "the logging environment")
	rootCmd.Flags().StringVar(&logging.Level, "logging-level", "info", "the root level of logging")
	rootCmd.AddCommand(newRunCommand())
	rootCmd.AddCommand(NewTimeDirCommand())
	return rootCmd
}

func newRunCommand() *cobra.Command {
	var (
		source       string
		streamRoot   string
		measureRoot  string
		propertyRoot string
		traceRoot    string
		schemaRoot   string
		fsConfig     remoteconfig.FsConfig
	)
	// Initialize nested structs to avoid nil pointer during flag binding
	fsConfig.S3 = &remoteconfig.S3Config{}
	fsConfig.Azure = &remoteconfig.AzureConfig{}
	fsConfig.GCP = &remoteconfig.GCPConfig{}

	cmd := &cobra.Command{
		Use:   "run",
		Short: "Restore BanyanDB data from remote storage",
		RunE: func(_ *cobra.Command, _ []string) error {
			if streamRoot == "" && measureRoot == "" && propertyRoot == "" && traceRoot == "" && schemaRoot == "" {
				return errors.New("at least one of stream-root-path, measure-root-path, property-root-path, trace-root-path or schema-root-path is required")
			}
			if source == "" {
				return errors.New("source is required")
			}

			cleanPaths := func(paths ...*string) error {
				for _, p := range paths {
					if *p != "" {
						var err error
						if *p, err = banyandbpath.Get(*p); err != nil {
							return err
						}
					}
				}
				return nil
			}

			if err := cleanPaths(&streamRoot, &measureRoot, &propertyRoot, &traceRoot); err != nil {
				return err
			}

			fs, err := newFS(source, &fsConfig)
			if err != nil {
				return err
			}
			defer fs.Close()

			catalogs := []catalogEntry{
				{rootPath: streamRoot, catalogName: snapshot.CatalogName(commonv1.Catalog_CATALOG_STREAM)},
				{rootPath: measureRoot, catalogName: snapshot.CatalogName(commonv1.Catalog_CATALOG_MEASURE)},
				{rootPath: propertyRoot, catalogName: snapshot.CatalogName(commonv1.Catalog_CATALOG_PROPERTY)},
				{rootPath: traceRoot, catalogName: snapshot.CatalogName(commonv1.Catalog_CATALOG_TRACE)},
				{rootPath: schemaRoot, catalogName: snapshot.SchemaPropertyCatalogName},
			}
			var errs error
			for _, c := range catalogs {
				if c.rootPath == "" {
					continue
				}
				if err := restoreFromTimeDir(fs, c.rootPath, c.catalogName); err != nil {
					errs = multierr.Append(errs, err)
				}
			}
			return errs
		},
	}
	cmd.Flags().StringVar(&source, "source", "", "Source URL (e.g., file:///backups)")
	cmd.Flags().StringVar(&streamRoot, "stream-root-path", "/tmp", "Root directory for stream catalog")
	cmd.Flags().StringVar(&measureRoot, "measure-root-path", "/tmp", "Root directory for measure catalog")
	cmd.Flags().StringVar(&propertyRoot, "property-root-path", "/tmp", "Root directory for property catalog")
	cmd.Flags().StringVar(&traceRoot, "trace-root-path", "/tmp", "Root directory for trace catalog")
	cmd.Flags().StringVar(&schemaRoot, "schema-root-path", "/tmp", "Root directory for schema property catalog")
	cmd.Flags().StringVar(&fsConfig.S3.S3ConfigFilePath, "s3-config-file", "", "Path to the s3 configuration file")
	cmd.Flags().StringVar(&fsConfig.S3.S3CredentialFilePath, "s3-credential-file", "", "Path to the s3 credential file")
	cmd.Flags().StringVar(&fsConfig.S3.S3ProfileName, "s3-profile", "", "S3 profile name")
	// Azure flags
	cmd.Flags().StringVar(&fsConfig.Azure.AzureAccountName, "azure-account-name", "", "Azure storage account name")
	cmd.Flags().StringVar(&fsConfig.Azure.AzureAccountKey, "azure-account-key", "", "Azure storage account key")
	cmd.Flags().StringVar(&fsConfig.Azure.AzureSASToken, "azure-sas-token", "", "Azure SAS token (alternative to account key)")
	cmd.Flags().StringVar(&fsConfig.Azure.AzureEndpoint, "azure-endpoint", "", "Azure blob service endpoint (override)")
	// GCP flags
	cmd.Flags().StringVar(&fsConfig.GCP.GCPServiceAccountFile, "gcp-service-account-file", "", "Path to the GCP service account JSON file")

	return cmd
}

type catalogEntry struct {
	rootPath    string
	catalogName string
}

func restoreFromTimeDir(fs remote.FS, rootPath, catalogName string) error {
	timeDirPath := filepath.Join(rootPath, catalogName, "time-dir")
	data, err := os.ReadFile(timeDirPath)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			logger.Infof("no %s time-dir file found, skip it: %s", catalogName, timeDirPath)
			return nil
		}
		return err
	}
	timeDir := strings.TrimSpace(string(data))
	if err = restoreByName(fs, timeDir, rootPath, catalogName); err != nil {
		return fmt.Errorf("%s restore failed: %w", catalogName, err)
	}
	logger.Infof("delete %s time-dir file", catalogName)
	_ = os.Remove(timeDirPath)
	return nil
}

func restoreCatalog(fs remote.FS, timeDir, rootPath string, catalog commonv1.Catalog) error {
	return restoreByName(fs, timeDir, rootPath, snapshot.CatalogName(catalog))
}

func restoreByName(fs remote.FS, timeDir, rootPath, catalogName string) error {
	remotePrefix := filepath.Join(timeDir, catalogName, "/")

	remoteFiles, err := fs.List(context.Background(), remotePrefix)
	if err != nil {
		return fmt.Errorf("failed to list remote files: %w", err)
	}

	localDir := filepath.Join(rootPath, catalogName, storage.DataDir)
	if err = os.MkdirAll(localDir, storage.DirPerm); err != nil {
		return fmt.Errorf("failed to create local directory %s: %w", localDir, err)
	}

	logger.Infof("Restoring %s to %s from %s, remote total %d files", catalogName, localDir, remotePrefix, len(remoteFiles))

	remoteRelSet := make(map[string]bool)
	for _, remoteFile := range remoteFiles {
		relPath, relPathErr := validatedRemoteRelPath(timeDir, catalogName, remoteFile)
		if relPathErr != nil {
			return relPathErr
		}
		remoteRelSet[path.Join(catalogName, relPath)] = true
	}

	localFiles, err := getAllFiles(localDir)
	if err != nil {
		return fmt.Errorf("failed to list local files: %w", err)
	}

	for _, localRelPath := range localFiles {
		localRelPathWithCatalog := path.Join(catalogName, filepath.ToSlash(localRelPath))
		if !remoteRelSet[localRelPathWithCatalog] {
			localPath := filepath.Join(localDir, localRelPath)
			logger.Infof("found local file: %s not exist in the remote storage, so delete it", localRelPathWithCatalog)
			if err := os.Remove(localPath); err != nil {
				return fmt.Errorf("failed to remove local file %s: %w", localPath, err)
			}
			cleanEmptyDirs(filepath.Dir(localPath), localDir)
		}
	}

	for _, remoteFile := range remoteFiles {
		relPath, err := validatedRemoteRelPath(timeDir, catalogName, remoteFile)
		if err != nil {
			return err
		}
		localPath := filepath.Join(rootPath, catalogName, storage.DataDir, relPath)

		if !contains(localFiles, relPath) {
			if err := os.MkdirAll(filepath.Dir(localPath), storage.DirPerm); err != nil {
				return fmt.Errorf("failed to create directory for %s: %w", localPath, err)
			}

			if err := downloadFile(context.Background(), fs, remoteFile, localPath); err != nil {
				return fmt.Errorf("failed to download %s: %w", remoteFile, err)
			}
			logger.Infof("Downloaded %s to %s", remoteFile, localPath)
		} else {
			logger.Infof("Skipping %s because it already exists", remoteFile)
		}
	}

	return nil
}

func validatedRemoteRelPath(timeDir, catalogName, remoteFile string) (string, error) {
	remotePath := filepath.ToSlash(remoteFile)
	if path.IsAbs(remotePath) || banyandbpath.HasVolumeName(remotePath) {
		return "", fmt.Errorf("remote file %q escapes backup prefix", remoteFile)
	}
	prefix := path.Clean(path.Join(filepath.ToSlash(timeDir), filepath.ToSlash(catalogName)))
	cleanRemotePath := path.Clean(remotePath)
	if banyandbpath.HasVolumeName(cleanRemotePath) || cleanRemotePath == "." || prefix == "." {
		return "", fmt.Errorf("remote file %q escapes backup prefix", remoteFile)
	}
	prefixWithSlash := prefix + "/"
	if !strings.HasPrefix(cleanRemotePath, prefixWithSlash) {
		return "", fmt.Errorf("remote file %q escapes backup prefix", remoteFile)
	}
	relPath := strings.TrimPrefix(cleanRemotePath, prefixWithSlash)
	return relPath, nil
}

func cleanEmptyDirs(dir, stopDir string) {
	for {
		if dir == stopDir || dir == "." {
			break
		}
		files, err := getAllFiles(dir)
		if err != nil || len(files) > 0 {
			break
		}
		logger.Infof("Detected empty file directory: %s, removing it", dir)
		_ = os.Remove(dir)
		dir = filepath.Dir(dir)
	}
}

func downloadFile(ctx context.Context, fs remote.FS, remotePath, localPath string) error {
	reader, err := fs.Download(ctx, remotePath)
	if err != nil {
		return err
	}
	defer reader.Close()

	file, err := os.Create(localPath)
	if err != nil {
		return err
	}
	defer file.Close()

	if _, err := io.Copy(file, reader); err != nil {
		return err
	}

	return nil
}
