//go:build failpoint

package backend

import (
	"bytes"
	"context"
	"errors"
	"strings"
	"testing"

	"github.com/c4pt0r/agfs/agfs-server/pkg/filesystem"
	"github.com/mem9-ai/drive9/pkg/datastore"
	"github.com/mem9-ai/drive9/pkg/s3client"
	"github.com/pingcap/failpoint"
)

const inTxCommitOutcomeUnknownFailpoint = "github.com/mem9-ai/drive9/pkg/datastore/inTxCommitOutcomeUnknown"

func TestDirectWriteRetainsObjectWhenCommitOutcomeIsUnknown(t *testing.T) {
	b := newTestBackend(t)
	rec := &deleteRecordingS3Client{S3Client: b.s3}
	b.s3 = rec
	fake := newFakeMetaQuotaStore()
	b.SetMetaQuotaStore(context.Background(), "tenant-a", fake)
	b.storageNamespaceID = "ns-a"

	if err := failpoint.Enable(inTxCommitOutcomeUnknownFailpoint, `return("injected commit response failure")`); err != nil {
		t.Fatalf("enable failpoint: %v", err)
	}
	t.Cleanup(func() {
		_ = failpoint.Disable(inTxCommitOutcomeUnknownFailpoint)
	})

	payload := make([]byte, b.inlineThreshold+1)
	if _, err := b.Write("/commit-unknown.bin", payload, 0, filesystem.WriteFlagCreate); err == nil {
		t.Fatal("write succeeded despite injected unknown commit outcome")
	}
	if keys := rec.deletedKeys(); len(keys) != 0 {
		t.Fatalf("deleted keys = %#v, want none", keys)
	}

	nf, err := b.Store().Stat(context.Background(), "/commit-unknown.bin")
	if err != nil {
		t.Fatalf("committed file is not readable: %v", err)
	}
	if nf.File == nil || nf.File.StorageRef == "" {
		t.Fatalf("committed file has no object reference: %+v", nf.File)
	}

	fake.mu.Lock()
	defer fake.mu.Unlock()
	if len(fake.objectGCCandidates) != 1 {
		t.Fatalf("object gc candidates = %d, want 1: %+v", len(fake.objectGCCandidates), fake.objectGCCandidates)
	}
	if fake.objectGCCandidates[0].StorageRef != nf.File.StorageRef {
		t.Fatalf("candidate storage ref = %q, want %q", fake.objectGCCandidates[0].StorageRef, nf.File.StorageRef)
	}
}

func TestAppendLogRewriteRetainsCandidatesWhenCommitOutcomeIsUnknown(t *testing.T) {
	b := newTestBackend(t)
	local, ok := b.s3.(*s3client.LocalS3Client)
	if !ok {
		t.Fatalf("S3 client = %T, want local client", b.s3)
	}
	fake := newFakeMetaQuotaStore()
	b.appendLogS3 = local
	b.SetMetaQuotaStore(context.Background(), "tenant-a", fake)
	b.storageNamespaceID = "ns-a"
	ctx := context.Background()
	created, err := b.AppendLog(ctx, "/commit-unknown-wal", bytes.NewReader([]byte("seed")), 4, 0, 0)
	if err != nil {
		t.Fatalf("create append-log: %v", err)
	}
	// Arm only after quota preflight, so the injected response failure belongs
	// to the transaction that publishes the replacement generation.
	b.appendLogS3 = &appendLogPutHookClient{
		AppendLogClient: local,
		beforePut: func() {
			if err := failpoint.Enable(inTxCommitOutcomeUnknownFailpoint, `return("injected commit response failure")`); err != nil {
				t.Fatalf("enable failpoint: %v", err)
			}
		},
	}
	t.Cleanup(func() { _ = failpoint.Disable(inTxCommitOutcomeUnknownFailpoint) })

	if _, err := b.RewriteAppendLog(ctx, "/commit-unknown-wal", bytes.NewReader([]byte("reset")), 5, created.Revision); !errors.Is(err, datastore.ErrCommitOutcomeUnknown) {
		t.Fatalf("rewrite error = %v, want ErrCommitOutcomeUnknown", err)
	}
	nf, err := b.Store().Stat(ctx, "/commit-unknown-wal")
	if err != nil {
		t.Fatalf("read committed append-log after unknown outcome: %v", err)
	}
	if nf.File == nil || nf.File.Revision != 2 || nf.File.SizeBytes != 5 {
		t.Fatalf("committed append-log metadata = %+v, want revision=2 size=5", nf.File)
	}
	fake.mu.Lock()
	defer fake.mu.Unlock()
	if fake.retryFailedWriteCalls != 1 {
		t.Fatalf("failed-write candidate retries = %d, want 1", fake.retryFailedWriteCalls)
	}
	if len(fake.objectGCCandidates) != 2 {
		t.Fatalf("object GC candidates = %d, want replacement and old-ref candidates: %+v", len(fake.objectGCCandidates), fake.objectGCCandidates)
	}
	for _, candidate := range fake.objectGCCandidates {
		if candidate.ContentLayout != string(datastore.ContentLayoutAppendLog) {
			t.Fatalf("candidate layout = %q, want append_log", candidate.ContentLayout)
		}
	}
}

func TestAppendLogCreateReleasesCandidateWhenCommitOutcomeIsUnknown(t *testing.T) {
	b := newTestBackend(t)
	local, ok := b.s3.(*s3client.LocalS3Client)
	if !ok {
		t.Fatalf("S3 client = %T, want local client", b.s3)
	}
	fake := newFakeMetaQuotaStore()
	b.SetMetaQuotaStore(context.Background(), "tenant-a", fake)
	b.storageNamespaceID = "ns-a"
	b.appendLogS3 = &appendLogPutHookClient{
		AppendLogClient: local,
		beforePut: func() {
			if err := failpoint.Enable(inTxCommitOutcomeUnknownFailpoint, `return("injected commit response failure")`); err != nil {
				t.Fatalf("enable failpoint: %v", err)
			}
		},
	}
	t.Cleanup(func() { _ = failpoint.Disable(inTxCommitOutcomeUnknownFailpoint) })

	if _, err := b.AppendLog(context.Background(), "/commit-unknown-create", bytes.NewReader([]byte("seed")), 4, 0, 0); !errors.Is(err, datastore.ErrCommitOutcomeUnknown) {
		t.Fatalf("create error = %v, want ErrCommitOutcomeUnknown", err)
	}
	fake.mu.Lock()
	defer fake.mu.Unlock()
	if fake.retryFailedWriteCalls != 1 {
		t.Fatalf("failed-write candidate retries = %d, want 1", fake.retryFailedWriteCalls)
	}
}

func TestAppendLogRebaseReleasesCandidateWhenCommitOutcomeIsUnknown(t *testing.T) {
	b := newTestBackend(t)
	local, ok := b.s3.(*s3client.LocalS3Client)
	if !ok {
		t.Fatalf("S3 client = %T, want local client", b.s3)
	}
	fake := newFakeMetaQuotaStore()
	b.SetMetaQuotaStore(context.Background(), "tenant-a", fake)
	b.storageNamespaceID = "ns-a"
	b.appendLogS3 = local
	ctx := context.Background()
	created, err := b.AppendLog(ctx, "/commit-unknown-rebase", bytes.NewReader([]byte("seed")), 4, 0, 0)
	if err != nil {
		t.Fatalf("create append-log: %v", err)
	}
	local.SetAppendLogFaults(s3client.AppendLogFaults{AppendErr: s3client.ErrAppendOffsetMismatch})
	b.appendLogS3 = &appendLogPutHookClient{
		AppendLogClient: local,
		beforePut: func() {
			if err := failpoint.Enable(inTxCommitOutcomeUnknownFailpoint, `return("injected commit response failure")`); err != nil {
				t.Fatalf("enable failpoint: %v", err)
			}
		},
	}
	t.Cleanup(func() { _ = failpoint.Disable(inTxCommitOutcomeUnknownFailpoint) })

	if _, err := b.AppendLog(ctx, "/commit-unknown-rebase", bytes.NewReader([]byte("tail")), 4, created.Revision, created.Size); !errors.Is(err, datastore.ErrCommitOutcomeUnknown) {
		t.Fatalf("rebase error = %v, want ErrCommitOutcomeUnknown", err)
	}
	fake.mu.Lock()
	defer fake.mu.Unlock()
	if fake.retryFailedWriteCalls != 1 {
		t.Fatalf("failed-write candidate retries = %d, want 1", fake.retryFailedWriteCalls)
	}
}

func TestFSLayerUploadRetainsObjectWhenCommitOutcomeIsUnknown(t *testing.T) {
	b := newTestBackendWithS3(t)
	ctx := context.Background()
	const (
		layerID = "layer-commit-unknown"
		path    = "/repo/commit-unknown.bin"
	)
	if err := b.Store().CreateFSLayer(ctx, &datastore.FSLayer{LayerID: layerID, BaseRootPath: "/repo"}); err != nil {
		t.Fatalf("create layer: %v", err)
	}
	fake := newFakeMetaQuotaStore()
	b.SetMetaQuotaStore(ctx, "tenant-a", fake)
	b.storageNamespaceID = "ns-a"
	payload := bytes.Repeat([]byte("c"), s3client.PartSize)
	plan, err := b.InitiateFSLayerUploadV2(ctx, layerID, path, int64(len(payload)), 0, 0o644)
	if err != nil {
		t.Fatalf("initiate layer upload: %v", err)
	}
	parts := uploadFSLayerDirectParts(t, b, plan, payload)
	rec := &deleteRecordingS3Client{S3Client: b.s3}
	b.s3 = rec

	if err := failpoint.Enable(inTxCommitOutcomeUnknownFailpoint, `return("injected commit response failure")`); err != nil {
		t.Fatalf("enable failpoint: %v", err)
	}
	t.Cleanup(func() {
		_ = failpoint.Disable(inTxCommitOutcomeUnknownFailpoint)
	})

	completed, err := b.ConfirmFSLayerUploadV2WithChecksum(ctx, plan.UploadID, parts, "")
	if err != nil {
		t.Fatalf("complete after confirmed unknown commit: %v", err)
	}
	if completed == nil || completed.StorageRef != plan.Key {
		t.Fatalf("completed entry = %+v, want storage ref %q", completed, plan.Key)
	}
	if keys := rec.deletedKeys(); len(keys) != 0 {
		t.Fatalf("deleted keys = %#v, want none", keys)
	}
	fake.mu.Lock()
	if fake.retryFailedWriteCalls != 0 || len(fake.objectGCCandidates) != 0 {
		t.Fatalf("confirmed commit GC state calls=%d candidates=%+v, want completed cleanup", fake.retryFailedWriteCalls, fake.objectGCCandidates)
	}
	if len(fake.mutations) != 1 || fake.mutations[0].typ != "fs_layer_store" {
		t.Fatalf("quota mutations = %+v, want recovered fs_layer_store handoff", fake.mutations)
	}
	fake.mu.Unlock()
	entry, err := b.Store().GetFSLayerEntry(ctx, layerID, path)
	if err != nil {
		t.Fatalf("committed layer entry is not readable: %v", err)
	}
	if entry.StorageRef != plan.Key {
		t.Fatalf("entry storage ref = %q, want %q", entry.StorageRef, plan.Key)
	}

	if err := failpoint.Disable(inTxCommitOutcomeUnknownFailpoint); err != nil {
		t.Fatalf("disable failpoint: %v", err)
	}
	if _, err := b.ConfirmFSLayerUploadV2WithChecksum(ctx, plan.UploadID, parts, strings.Repeat("0", 64)); !errors.Is(err, datastore.ErrUploadConflict) {
		t.Fatalf("retry with mismatched checksum error=%v, want ErrUploadConflict", err)
	}
	fake.mu.Lock()
	if len(fake.objectGCCandidates) != 0 {
		t.Fatalf("GC candidates after mismatched completed retry = %+v, want none", fake.objectGCCandidates)
	}
	fake.mu.Unlock()
	retried, err := b.ConfirmFSLayerUploadV2WithChecksum(ctx, plan.UploadID, parts, "")
	if err != nil {
		t.Fatalf("retry completed layer upload: %v", err)
	}
	if retried.EntrySeq != entry.EntrySeq || retried.StorageRef != entry.StorageRef {
		t.Fatalf("retried entry = %+v, want seq=%d ref=%q", retried, entry.EntrySeq, entry.StorageRef)
	}
	fake.mu.Lock()
	if len(fake.objectGCCandidates) != 0 {
		t.Fatalf("GC candidates after completed retry = %+v, want none", fake.objectGCCandidates)
	}
	fake.mu.Unlock()
}

func TestMultipartUploadAccountsQuotaWhenCommitOutcomeIsUnknown(t *testing.T) {
	b := newTestBackendWithS3(t)
	ctx := context.Background()
	fake := newFakeMetaQuotaStore()
	b.SetMetaQuotaStore(ctx, "tenant-a", fake)

	payload := bytes.Repeat([]byte("u"), 2<<20)
	plan, err := b.InitiateUpload(ctx, "/commit-unknown-upload.bin", int64(len(payload)))
	if err != nil {
		t.Fatalf("initiate upload: %v", err)
	}
	upload, err := b.GetUpload(ctx, plan.UploadID)
	if err != nil {
		t.Fatalf("get upload: %v", err)
	}
	for _, part := range plan.Parts {
		start := int64(part.Number-1) * plan.PartSize
		end := start + part.Size
		if end > int64(len(payload)) {
			end = int64(len(payload))
		}
		if _, err := b.S3().(*s3client.LocalS3Client).UploadPart(ctx, upload.S3UploadID, part.Number, bytes.NewReader(payload[start:end])); err != nil {
			t.Fatalf("upload part %d: %v", part.Number, err)
		}
	}

	if err := failpoint.Enable(inTxCommitOutcomeUnknownFailpoint, `return("injected commit response failure")`); err != nil {
		t.Fatalf("enable failpoint: %v", err)
	}
	t.Cleanup(func() { _ = failpoint.Disable(inTxCommitOutcomeUnknownFailpoint) })

	if err := b.ConfirmUpload(ctx, plan.UploadID); err != nil {
		t.Fatalf("confirm after committed unknown outcome: %v", err)
	}
	committed, err := b.Store().Stat(ctx, "/commit-unknown-upload.bin")
	if err != nil || committed.File == nil {
		t.Fatalf("committed file = %+v, err=%v", committed, err)
	}
	fake.mu.Lock()
	defer fake.mu.Unlock()
	if len(fake.mutations) != 1 || fake.mutations[0].typ != "upload_complete" {
		t.Fatalf("quota mutations = %+v, want recovered upload_complete handoff", fake.mutations)
	}
}