package backend

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"

	"github.com/mem9-ai/drive9/pkg/datastore"
	"github.com/mem9-ai/drive9/pkg/s3client"
)

func uploadFSLayerDirectParts(t *testing.T, b *Dat9Backend, plan *UploadPlanV2, payload []byte) []CompletePart {
	t.Helper()
	ctx := context.Background()
	partNumbers := make([]int, plan.TotalParts)
	for i := range partNumbers {
		partNumbers[i] = i + 1
	}
	urls, err := b.PresignParts(ctx, plan.UploadID, entriesFromInts(partNumbers))
	if err != nil {
		t.Fatalf("presign layer upload parts: %v", err)
	}
	upload, err := b.GetUpload(ctx, plan.UploadID)
	if err != nil {
		t.Fatalf("get layer upload: %v", err)
	}
	localS3, ok := b.S3().(*s3client.LocalS3Client)
	if !ok {
		t.Fatalf("S3 client = %T, want *s3client.LocalS3Client", b.S3())
	}
	completed := make([]CompletePart, len(urls))
	for i, part := range urls {
		start := int64(part.Number-1) * upload.PartSize
		end := start + part.Size
		etag, err := localS3.UploadPart(ctx, upload.S3UploadID, part.Number, bytes.NewReader(payload[start:end]))
		if err != nil {
			t.Fatalf("upload layer part %d: %v", part.Number, err)
		}
		completed[i] = CompletePart{Number: part.Number, ETag: etag}
	}
	return completed
}

func readLocalPresignedURL(t *testing.T, localS3 *s3client.LocalS3Client, presignedURL string) []byte {
	t.Helper()
	parsed, err := url.Parse(presignedURL)
	if err != nil {
		t.Fatalf("parse presigned URL: %v", err)
	}
	server := httptest.NewServer(localS3.Handler())
	t.Cleanup(server.Close)
	resp, err := http.Get(server.URL + strings.TrimPrefix(parsed.EscapedPath(), "/s3")) //nolint:gosec // local test server
	if err != nil {
		t.Fatalf("GET presigned layer object: %v", err)
	}
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		t.Fatalf("GET presigned layer object status=%d body=%q", resp.StatusCode, body)
	}
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("read presigned layer object: %v", err)
	}
	return body
}

func TestFSLayerDirectUploadPublishesOnlyAfterS3Completion(t *testing.T) {
	b := newTestBackendWithS3(t)
	ctx := context.Background()
	quotaStore := newFakeMetaQuotaStore()
	b.SetMetaQuotaStore(ctx, "tenant-direct-layer", quotaStore)
	const (
		layerID = "layer-direct-complete"
		path    = "/repo/large.bin"
	)
	if err := b.Store().CreateFSLayer(ctx, &datastore.FSLayer{LayerID: layerID, BaseRootPath: "/repo"}); err != nil {
		t.Fatalf("create layer: %v", err)
	}
	payload := make([]byte, s3client.PartSize+257)
	for i := range payload {
		payload[i] = byte(i % 251)
	}

	plan, err := b.InitiateFSLayerUploadV2(ctx, layerID, path, int64(len(payload)), 0, 0o640)
	if err != nil {
		t.Fatalf("initiate layer upload: %v", err)
	}
	if !strings.HasPrefix(plan.Key, "layers/"+layerID+"/") {
		t.Fatalf("layer object key = %q", plan.Key)
	}
	if _, err := b.Store().Stat(ctx, path); !errors.Is(err, datastore.ErrNotFound) {
		t.Fatalf("base path exists before completion: %v", err)
	}
	if _, err := b.Store().GetFSLayerEntry(ctx, layerID, path); !errors.Is(err, datastore.ErrNotFound) {
		t.Fatalf("layer entry exists before completion: %v", err)
	}
	before := &datastore.FSLayerCheckpoint{CheckpointID: "before-direct-complete", LayerID: layerID}
	if err := b.Store().CreateFSLayerCheckpoint(ctx, before); err != nil {
		t.Fatalf("checkpoint before completion: %v", err)
	}
	if before.DurableSeq != 0 {
		t.Fatalf("checkpoint before completion durable_seq=%d, want 0", before.DurableSeq)
	}

	parts := uploadFSLayerDirectParts(t, b, plan, payload)
	digest := sha256.Sum256(payload)
	entry, published, err := b.ConfirmFSLayerUploadV2WithChecksumResult(ctx, plan.UploadID, parts, hex.EncodeToString(digest[:]))
	if err != nil {
		t.Fatalf("complete layer upload: %v", err)
	}
	if !published {
		t.Fatal("first completion did not report a newly published layer entry")
	}
	if entry.StorageType != string(datastore.StorageS3) || entry.StorageRef != plan.Key {
		t.Fatalf("published entry storage = %q/%q, want s3/%q", entry.StorageType, entry.StorageRef, plan.Key)
	}
	if entry.SizeBytes != int64(len(payload)) || entry.ChecksumSHA256 != hex.EncodeToString(digest[:]) || entry.Mode != 0o640 {
		t.Fatalf("published entry metadata = %+v", entry)
	}
	if _, err := b.Store().Stat(ctx, path); !errors.Is(err, datastore.ErrNotFound) {
		t.Fatalf("base path changed after layer completion: %v", err)
	}
	log, err := b.Store().ListFSLayerEntryLog(ctx, layerID)
	if err != nil {
		t.Fatalf("list layer entry log: %v", err)
	}
	if len(log) != 1 || log[0].EntrySeq != entry.EntrySeq {
		t.Fatalf("layer entry log = %+v, want one published entry", log)
	}
	quotaStore.mu.Lock()
	if len(quotaStore.mutations) != 1 || quotaStore.mutations[0].typ != "fs_layer_store" {
		quotaStore.mu.Unlock()
		t.Fatalf("quota mutations = %+v, want one fs_layer_store handoff", quotaStore.mutations)
	}
	quotaStore.mu.Unlock()
	retried, published, err := b.ConfirmFSLayerUploadV2WithChecksumResult(ctx, plan.UploadID, parts, hex.EncodeToString(digest[:]))
	if err != nil {
		t.Fatalf("retry completed layer upload: %v", err)
	}
	if published {
		t.Fatal("retry reported the existing layer entry as newly published")
	}
	if retried.EntrySeq != entry.EntrySeq || retried.StorageRef != entry.StorageRef {
		t.Fatalf("retried entry = %+v, want original seq=%d ref=%q", retried, entry.EntrySeq, entry.StorageRef)
	}
	log, err = b.Store().ListFSLayerEntryLog(ctx, layerID)
	if err != nil {
		t.Fatalf("list layer entry log after retry: %v", err)
	}
	if len(log) != 1 {
		t.Fatalf("layer entry log after retry has %d entries, want 1", len(log))
	}
	if _, err := b.ConfirmFSLayerUploadV2WithChecksum(ctx, plan.UploadID, parts, strings.Repeat("0", sha256.Size*2)); !errors.Is(err, datastore.ErrUploadConflict) {
		t.Fatalf("retry with mismatched checksum error=%v, want ErrUploadConflict", err)
	}
	after := &datastore.FSLayerCheckpoint{CheckpointID: "after-direct-complete", LayerID: layerID}
	if err := b.Store().CreateFSLayerCheckpoint(ctx, after); err != nil {
		t.Fatalf("checkpoint after completion: %v", err)
	}
	if after.DurableSeq != entry.EntrySeq {
		t.Fatalf("checkpoint after completion durable_seq=%d, want %d", after.DurableSeq, entry.EntrySeq)
	}

	presignedURL, err := b.PresignFSLayerEntryData(ctx, entry)
	if err != nil {
		t.Fatalf("presign layer read: %v", err)
	}
	got := readLocalPresignedURL(t, b.S3().(*s3client.LocalS3Client), presignedURL)
	if !bytes.Equal(got, payload) {
		t.Fatalf("presigned layer payload size=%d, want %d", len(got), len(payload))
	}
}

func TestInitiateFSLayerUploadV2RejectsAppendLogBeforeMultipart(t *testing.T) {
	ctx := context.Background()
	b, _ := newAppendLogTestBackend(t)
	if _, err := b.AppendLog(ctx, "/repo/wal", bytes.NewReader([]byte("seed")), 4, 0, 0); err != nil {
		t.Fatalf("create append-log: %v", err)
	}
	if err := b.Store().CreateFSLayer(ctx, &datastore.FSLayer{LayerID: "layer-append-log", BaseRootPath: "/repo"}); err != nil {
		t.Fatalf("create layer: %v", err)
	}

	if _, err := b.InitiateFSLayerUploadV2(ctx, "layer-append-log", "/repo/wal", 1, 0, 0o644); !errors.Is(err, ErrAppendLogUnsupported) {
		t.Fatalf("initiate append-log layer upload error = %v, want ErrAppendLogUnsupported", err)
	}
	if _, err := b.Store().GetFSLayerUploadByPath(ctx, "layer-append-log", "/repo/wal"); !errors.Is(err, datastore.ErrNotFound) {
		t.Fatalf("append-log layer upload persisted: %v", err)
	}
}

func TestFSLayerDirectUploadAbortLeavesNoPublishedOrBaseFile(t *testing.T) {
	b := newTestBackendWithS3(t)
	ctx := context.Background()
	const (
		layerID = "layer-direct-abort"
		path    = "/repo/aborted.bin"
	)
	if err := b.Store().CreateFSLayer(ctx, &datastore.FSLayer{LayerID: layerID, BaseRootPath: "/repo"}); err != nil {
		t.Fatalf("create layer: %v", err)
	}
	plan, err := b.InitiateFSLayerUploadV2(ctx, layerID, path, s3client.PartSize, 0, 0o644)
	if err != nil {
		t.Fatalf("initiate layer upload: %v", err)
	}
	if _, err := b.PresignPart(ctx, plan.UploadID, 1, nil); err != nil {
		t.Fatalf("presign layer part: %v", err)
	}
	if err := b.AbortUploadV2(ctx, plan.UploadID); err != nil {
		t.Fatalf("abort layer upload: %v", err)
	}
	if _, err := b.Store().GetFSLayerEntry(ctx, layerID, path); !errors.Is(err, datastore.ErrNotFound) {
		t.Fatalf("layer entry after abort: %v", err)
	}
	if _, err := b.Store().Stat(ctx, path); !errors.Is(err, datastore.ErrNotFound) {
		t.Fatalf("base path after layer abort: %v", err)
	}
	upload, err := b.GetUpload(ctx, plan.UploadID)
	if err != nil {
		t.Fatalf("get aborted upload: %v", err)
	}
	if upload.Status != datastore.UploadAborted {
		t.Fatalf("upload status=%q, want %q", upload.Status, datastore.UploadAborted)
	}
}

func TestFSLayerDirectUploadCannotPublishAfterLayerSealed(t *testing.T) {
	b := newTestBackendWithS3(t)
	ctx := context.Background()
	const (
		layerID = "layer-direct-sealed"
		path    = "/repo/sealed.bin"
	)
	if err := b.Store().CreateFSLayer(ctx, &datastore.FSLayer{LayerID: layerID, BaseRootPath: "/repo"}); err != nil {
		t.Fatalf("create layer: %v", err)
	}
	payload := bytes.Repeat([]byte("s"), 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)
	if err := b.Store().SetFSLayerState(ctx, layerID, datastore.FSLayerStateSealed); err != nil {
		t.Fatalf("seal layer: %v", err)
	}
	if _, err := b.ConfirmFSLayerUploadV2WithChecksum(ctx, plan.UploadID, parts, ""); !errors.Is(err, datastore.ErrFSLayerStateConflict) {
		t.Fatalf("complete sealed layer error=%v, want ErrFSLayerStateConflict", err)
	}
	if _, err := b.Store().GetFSLayerEntry(ctx, layerID, path); !errors.Is(err, datastore.ErrNotFound) {
		t.Fatalf("sealed layer entry was published: %v", err)
	}
}