// Copyright 2026 tamakyi. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package cmd

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/smithy-go"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/urfave/cli/v2"

	"github.com/tamakyi/TamaBox/internal/conf"
	"github.com/tamakyi/TamaBox/internal/storage"
)

const maxImageManifestBytes = 64 << 20 // 64 MiB，足以容纳大量对象且避免恶意清单耗尽内存。

// RestoreImages 从 backup 生成的 images-manifest.json 和 images/ 恢复 R2 对象。
var RestoreImages = &cli.Command{
	Name:   "restore-images",
	Usage:  "Validate and restore backed-up images to R2",
	Action: runRestoreImages,
	Flags: []cli.Flag{
		&cli.StringFlag{
			Name:     "input",
			Required: true,
			Usage:    "Backup directory or images-manifest.json path",
		},
		&cli.BoolFlag{
			Name:  "dry-run",
			Usage: "Validate local files and remote targets without writing objects",
		},
		&cli.BoolFlag{
			Name:  "yes",
			Usage: "Confirm actual remote writes",
		},
		&cli.BoolFlag{
			Name:  "overwrite",
			Usage: "Allow replacing objects that already exist",
		},
		&cli.StringFlag{
			Name:  "target-prefix",
			Value: "",
			Usage: "Prefix prepended to every restored key (for isolated drills)",
		},
	},
}

type restoreImagesOptions struct {
	Input        string
	DryRun       bool
	Yes          bool
	Overwrite    bool
	TargetPrefix string
}

// restoreObjectStore 的 PutObject 是带重试的高层操作，便于离线 fake 精确断言零写入。
type restoreObjectStore interface {
	HeadObject(context.Context, *s3.HeadObjectInput, ...func(*s3.Options)) (*s3.HeadObjectOutput, error)
	PutObject(context.Context, *s3.PutObjectInput) error
}

type r2RestoreObjectStore struct {
	client *s3.Client
}

func (s *r2RestoreObjectStore) HeadObject(ctx context.Context, input *s3.HeadObjectInput, options ...func(*s3.Options)) (*s3.HeadObjectOutput, error) {
	return s.client.HeadObject(ctx, input, options...)
}

func (s *r2RestoreObjectStore) PutObject(ctx context.Context, input *s3.PutObjectInput) error {
	// storage.PutObject 提供每次独立 60 秒超时、最多 3 次重试、Body 回绕和可取消退避。
	return storage.PutObject(ctx, s.client, input)
}

type restorePlannedObject struct {
	entry      imageManifestEntry
	sourcePath string
	targetKey  string
}

func runRestoreImages(ctx *cli.Context) error {
	opts := restoreImagesOptions{
		Input:        ctx.String("input"),
		DryRun:       ctx.Bool("dry-run"),
		Yes:          ctx.Bool("yes"),
		Overwrite:    ctx.Bool("overwrite"),
		TargetPrefix: ctx.String("target-prefix"),
	}
	if opts.Input == "" {
		return errors.New("--input must not be empty")
	}
	if !opts.DryRun && !opts.Yes {
		return errors.New("actual restore writes require --yes (or use --dry-run)")
	}
	if err := conf.Init(); err != nil {
		return errors.Wrap(err, "load configuration")
	}

	client, err := backupS3Client(ctx.Context)
	if err != nil {
		return errors.Wrap(err, "get R2 client for image restore")
	}
	return restoreImagesWithStore(ctx.Context, &r2RestoreObjectStore{client: client}, conf.Upload.ImageBucket, opts)
}

func restoreImagesWithStore(ctx context.Context, store restoreObjectStore, bucket string, opts restoreImagesOptions) error {
	if opts.Input == "" {
		return errors.New("restore input must not be empty")
	}
	if !opts.DryRun && !opts.Yes {
		return errors.New("actual restore writes require --yes (or use --dry-run)")
	}

	manifestPath, imageDir, err := resolveRestoreInput(opts.Input)
	if err != nil {
		return err
	}
	entries, err := readImageManifest(manifestPath)
	if err != nil {
		return err
	}

	// 关键安全顺序：先验证所有 manifest 条目和全部本地文件，再做任何 Head/Put。
	// 因此重复 key、非法 hash/size、缺文件、篡改或 symlink 都不会造成部分上传。
	objects := make([]restorePlannedObject, 0, len(entries))
	for _, entry := range entries {
		targetKey, err := restoreTargetKey(opts.TargetPrefix, entry.Key)
		if err != nil {
			return err
		}
		sourcePath, err := localImagePath(imageDir, entry.Key)
		if err != nil {
			return err
		}
		if err := validateRestoreLocalFile(imageDir, sourcePath, entry); err != nil {
			return err
		}
		objects = append(objects, restorePlannedObject{
			entry:      entry,
			sourcePath: sourcePath,
			targetKey:  targetKey,
		})
	}

	// 默认先 Head 所有目标。任何目标已存在都会在 Put 循环开始前失败，保证零 Put。
	if !opts.Overwrite {
		for _, object := range objects {
			headCtx, cancel := context.WithTimeout(ctx, imageObjectAttemptTimeout)
			_, err := store.HeadObject(headCtx, &s3.HeadObjectInput{
				Bucket: aws.String(bucket),
				Key:    aws.String(object.targetKey),
			})
			cancel()
			if err == nil {
				return errors.Errorf("restore target %q already exists; no objects were uploaded (use --overwrite to replace)", object.targetKey)
			}
			if !isS3ObjectNotFound(err) {
				return errors.Wrapf(err, "check restore target %q; no objects were uploaded", object.targetKey)
			}
		}
	}

	if opts.DryRun {
		logrus.WithFields(logrus.Fields{
			"objects": len(objects),
			"prefix":  opts.TargetPrefix,
		}).Info("图片恢复演练完成，未写入对象")
		return nil
	}

	succeeded := 0
	for _, object := range objects {
		file, err := openRestoreFileForPut(object.sourcePath, object.entry.Size)
		if err != nil {
			return errors.Wrapf(err, "open validated restore file for %q after %d successful uploads; restore is not transactional", object.entry.Key, succeeded)
		}
		input := &s3.PutObjectInput{
			Bucket:        aws.String(bucket),
			Key:           aws.String(object.targetKey),
			Body:          file,
			ContentLength: aws.Int64(object.entry.Size),
			ContentType:   nonEmptyStringPointer(object.entry.ContentType),
			CacheControl:  nonEmptyStringPointer(object.entry.CacheControl),
		}
		putErr := store.PutObject(ctx, input)
		closeErr := file.Close()
		if putErr != nil {
			return errors.Wrapf(putErr, "restore object %q to %q failed after %d successful uploads; restore is not transactional", object.entry.Key, object.targetKey, succeeded)
		}
		if closeErr != nil {
			return errors.Wrapf(closeErr, "close restored file %q (remote object was uploaded); restore is not transactional", object.sourcePath)
		}
		succeeded++
	}

	logrus.WithFields(logrus.Fields{
		"objects": succeeded,
		"prefix":  opts.TargetPrefix,
	}).Info("图片恢复完成")
	return nil
}

func resolveRestoreInput(input string) (manifestPath, imageDir string, err error) {
	info, err := os.Lstat(input)
	if err != nil {
		return "", "", errors.Wrapf(err, "inspect restore input %q", input)
	}
	if info.Mode()&os.ModeSymlink != 0 {
		return "", "", errors.Errorf("restore input %q must not be a symlink", input)
	}

	manifestPath = input
	if info.IsDir() {
		manifestPath = filepath.Join(input, "images-manifest.json")
	} else if !info.Mode().IsRegular() {
		return "", "", errors.Errorf("restore input %q is neither a directory nor a regular manifest file", input)
	}
	return manifestPath, filepath.Join(filepath.Dir(manifestPath), "images"), nil
}

func readImageManifest(manifestPath string) ([]imageManifestEntry, error) {
	info, err := os.Lstat(manifestPath)
	if err != nil {
		return nil, errors.Wrapf(err, "inspect image manifest %q", manifestPath)
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
		return nil, errors.Errorf("image manifest %q must be a regular non-symlink file", manifestPath)
	}
	if info.Size() > maxImageManifestBytes {
		return nil, errors.Errorf("image manifest %q exceeds the %d-byte limit", manifestPath, maxImageManifestBytes)
	}

	file, err := os.Open(manifestPath)
	if err != nil {
		return nil, errors.Wrapf(err, "open image manifest %q", manifestPath)
	}
	data, readErr := io.ReadAll(io.LimitReader(file, maxImageManifestBytes+1))
	closeErr := file.Close()
	if readErr != nil {
		return nil, errors.Wrap(readErr, "read image manifest")
	}
	if closeErr != nil {
		return nil, errors.Wrap(closeErr, "close image manifest")
	}
	if len(data) > maxImageManifestBytes {
		return nil, errors.Errorf("image manifest %q exceeds the %d-byte limit", manifestPath, maxImageManifestBytes)
	}

	decoder := json.NewDecoder(bytes.NewReader(data))
	var entries []imageManifestEntry
	if err := decoder.Decode(&entries); err != nil {
		return nil, errors.Wrap(err, "decode image manifest")
	}
	if entries == nil {
		return nil, errors.New("image manifest must be a JSON array, not null")
	}
	var trailing json.RawMessage
	if err := decoder.Decode(&trailing); err != io.EOF {
		if err == nil {
			return nil, errors.New("image manifest contains trailing JSON")
		}
		return nil, errors.Wrap(err, "image manifest contains trailing data")
	}

	seen := make(map[string]struct{}, len(entries))
	for index := range entries {
		entry := &entries[index]
		if err := validateObjectKey(entry.Key); err != nil {
			return nil, errors.Wrapf(err, "invalid manifest entry %d", index)
		}
		if _, exists := seen[entry.Key]; exists {
			return nil, errors.Errorf("image manifest contains duplicate key %q", entry.Key)
		}
		seen[entry.Key] = struct{}{}
		if entry.Size < 0 {
			return nil, errors.Errorf("manifest entry %q has negative size %d", entry.Key, entry.Size)
		}
		decodedHash, err := hex.DecodeString(entry.SHA256)
		if err != nil || len(decodedHash) != sha256.Size {
			return nil, errors.Errorf("manifest entry %q has invalid SHA-256 %q", entry.Key, entry.SHA256)
		}
		entry.SHA256 = strings.ToLower(entry.SHA256)
	}
	return entries, nil
}

func restoreTargetKey(prefix, key string) (string, error) {
	if prefix != "" && !strings.HasSuffix(prefix, "/") {
		prefix += "/"
	}
	target := prefix + key
	if err := validateObjectKey(target); err != nil {
		return "", errors.Wrapf(err, "invalid restore target for source key %q", key)
	}
	return target, nil
}

func validateRestoreLocalFile(imageDir, sourcePath string, entry imageManifestEntry) error {
	root, err := filepath.Abs(imageDir)
	if err != nil {
		return err
	}
	if err := validateRestoreDirectoryChain(root, filepath.Dir(sourcePath)); err != nil {
		return err
	}
	info, err := os.Lstat(sourcePath)
	if err != nil {
		return errors.Wrapf(err, "inspect local image for %q", entry.Key)
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
		return errors.Errorf("local image for %q must be a regular non-symlink file", entry.Key)
	}
	if info.Size() != entry.Size {
		return errors.Errorf("local image size mismatch for %q: manifest=%d actual=%d", entry.Key, entry.Size, info.Size())
	}

	file, err := os.Open(sourcePath)
	if err != nil {
		return errors.Wrapf(err, "open local image for %q", entry.Key)
	}
	openedInfo, statErr := file.Stat()
	if statErr != nil {
		_ = file.Close()
		return errors.Wrapf(statErr, "stat opened local image for %q", entry.Key)
	}
	if !openedInfo.Mode().IsRegular() || openedInfo.Size() != entry.Size {
		_ = file.Close()
		return errors.Errorf("opened local image for %q is not the expected regular file", entry.Key)
	}
	hash := sha256.New()
	_, hashErr := io.Copy(hash, file)
	closeErr := file.Close()
	if hashErr != nil {
		return errors.Wrapf(hashErr, "hash local image for %q", entry.Key)
	}
	if closeErr != nil {
		return errors.Wrapf(closeErr, "close local image for %q", entry.Key)
	}
	actualHash := hex.EncodeToString(hash.Sum(nil))
	if actualHash != entry.SHA256 {
		return errors.Errorf("local image SHA-256 mismatch for %q: manifest=%s actual=%s", entry.Key, entry.SHA256, actualHash)
	}
	return nil
}

func validateRestoreDirectoryChain(root, parent string) error {
	rootInfo, err := os.Lstat(root)
	if err != nil {
		return errors.Wrapf(err, "inspect restore images directory %q", root)
	}
	if rootInfo.Mode()&os.ModeSymlink != 0 || !rootInfo.IsDir() {
		return errors.Errorf("restore images path %q must be a real directory", root)
	}
	rel, err := filepath.Rel(root, parent)
	if err != nil || !pathWithinDirectory(root, parent) {
		return errors.Errorf("restore file parent %q escapes images directory", parent)
	}
	if rel == "." {
		return nil
	}
	current := root
	for _, segment := range strings.Split(rel, string(filepath.Separator)) {
		current = filepath.Join(current, segment)
		info, err := os.Lstat(current)
		if err != nil {
			return errors.Wrapf(err, "inspect restore parent directory %q", current)
		}
		if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
			return errors.Errorf("restore parent %q must be a real directory", current)
		}
	}
	return nil
}

func openRestoreFileForPut(sourcePath string, expectedSize int64) (*os.File, error) {
	info, err := os.Lstat(sourcePath)
	if err != nil {
		return nil, err
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() || info.Size() != expectedSize {
		return nil, fmt.Errorf("local file %q changed after validation", sourcePath)
	}
	file, err := os.Open(sourcePath)
	if err != nil {
		return nil, err
	}
	openedInfo, err := file.Stat()
	if err != nil {
		_ = file.Close()
		return nil, err
	}
	if !openedInfo.Mode().IsRegular() || openedInfo.Size() != expectedSize {
		_ = file.Close()
		return nil, fmt.Errorf("local file %q changed after validation", sourcePath)
	}
	return file, nil
}

func isS3ObjectNotFound(err error) bool {
	if err == nil {
		return false
	}
	var statusErr interface{ HTTPStatusCode() int }
	if errors.As(err, &statusErr) && statusErr.HTTPStatusCode() == 404 {
		return true
	}
	var apiErr smithy.APIError
	if errors.As(err, &apiErr) {
		switch apiErr.ErrorCode() {
		case "NotFound", "NoSuchKey", "404":
			return true
		}
	}
	return false
}

func nonEmptyStringPointer(value string) *string {
	if value == "" {
		return nil
	}
	return aws.String(value)
}