package datastore

import (
	"context"
	"database/sql"
	"errors"
	"testing"
	"time"

	"github.com/mem9-ai/drive9/pkg/semantic"
)

func TestInsertFileTx(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()

	if err := s.InTx(context.Background(), func(tx *sql.Tx) error {
		return s.InsertFileTx(tx, &File{
			FileID:      "f1",
			StorageType: StorageDB9,
			StorageRef:  "/blobs/f1",
			ContentType: "text/plain",
			SizeBytes:   12,
			Revision:    1,
			Status:      StatusConfirmed,
			CreatedAt:   now,
			ConfirmedAt: &now,
		})
	}); err != nil {
		t.Fatal(err)
	}

	got, err := s.GetFile(context.Background(), "f1")
	if err != nil {
		t.Fatal(err)
	}
	if got.StorageRef != "/blobs/f1" || got.Revision != 1 {
		t.Fatalf("unexpected inserted file: %+v", got)
	}
}

func TestInsertFileClearsKeyIDForNonKMSMode(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	if err := s.InsertFile(context.Background(), &File{
		FileID:                 "f1",
		StorageType:            StorageS3,
		StorageRef:             "/blobs/f1",
		StorageEncryptionMode:  StorageEncryptionSSES3,
		StorageEncryptionKeyID: "stale-key",
		SizeBytes:              10,
		Revision:               1,
		Status:                 StatusConfirmed,
		CreatedAt:              now,
		ConfirmedAt:            &now,
	}); err != nil {
		t.Fatal(err)
	}

	got, err := s.GetFile(context.Background(), "f1")
	if err != nil {
		t.Fatal(err)
	}
	if got.StorageEncryptionMode != StorageEncryptionSSES3 {
		t.Fatalf("storage_encryption_mode=%q, want %q", got.StorageEncryptionMode, StorageEncryptionSSES3)
	}
	if got.StorageEncryptionKeyID != "" {
		t.Fatalf("storage_encryption_key_id=%q, want empty", got.StorageEncryptionKeyID)
	}
}

func TestInsertUploadClearsKeyIDForNonKMSMode(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	if err := s.InsertUpload(context.Background(), &Upload{
		UploadID:               "u1",
		FileID:                 "f1",
		TargetPath:             "/upload.bin",
		S3UploadID:             "s3-upload",
		S3Key:                  "blobs/f1",
		StorageEncryptionMode:  StorageEncryptionNone,
		StorageEncryptionKeyID: "stale-key",
		TotalSize:              10,
		PartSize:               10,
		PartsTotal:             1,
		Status:                 UploadUploading,
		CreatedAt:              now,
		UpdatedAt:              now,
		ExpiresAt:              now.Add(time.Hour),
	}); err != nil {
		t.Fatal(err)
	}

	got, err := s.GetUpload(context.Background(), "u1")
	if err != nil {
		t.Fatal(err)
	}
	if got.StorageEncryptionMode != StorageEncryptionNone {
		t.Fatalf("storage_encryption_mode=%q, want %q", got.StorageEncryptionMode, StorageEncryptionNone)
	}
	if got.StorageEncryptionKeyID != "" {
		t.Fatalf("storage_encryption_key_id=%q, want empty", got.StorageEncryptionKeyID)
	}
}

func TestFSLayerUploadTargetRoundTripAndIsolation(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	baseRevision := int64(17)

	newLayerUpload := func(uploadID, layerID string) *Upload {
		return &Upload{
			UploadID:         uploadID,
			FileID:           uploadID + "-object",
			TargetKind:       UploadTargetFSLayer,
			LayerID:          layerID,
			TargetPath:       "/same.bin",
			BaseInodeID:      "base-inode-1",
			ExpectedRevision: &baseRevision,
			Mode:             0o640,
			S3UploadID:       "s3-" + uploadID,
			S3Key:            "layers/" + layerID + "/" + uploadID,
			TotalSize:        8,
			PartSize:         8,
			PartsTotal:       1,
			Status:           UploadUploading,
			CreatedAt:        now,
			UpdatedAt:        now,
			ExpiresAt:        now.Add(time.Hour),
		}
	}

	first := newLayerUpload("layer-upload-a", "layer-a")
	if err := s.InsertUpload(ctx, first); err != nil {
		t.Fatalf("insert first layer upload: %v", err)
	}
	second := newLayerUpload("layer-upload-b", "layer-b")
	if err := s.InsertUpload(ctx, second); err != nil {
		t.Fatalf("same path in another layer must not conflict: %v", err)
	}
	base := &Upload{
		UploadID:   "base-upload",
		FileID:     "base-file",
		TargetKind: UploadTargetFile,
		TargetPath: "/same.bin",
		S3UploadID: "s3-base",
		S3Key:      "blobs/base-file",
		TotalSize:  8,
		PartSize:   8,
		PartsTotal: 1,
		Status:     UploadUploading,
		CreatedAt:  now,
		UpdatedAt:  now,
		ExpiresAt:  now.Add(time.Hour),
	}
	if err := s.InsertUpload(ctx, base); err != nil {
		t.Fatalf("base upload must not conflict with layer path: %v", err)
	}
	if err := s.InTx(ctx, func(tx *sql.Tx) error {
		reserved, err := s.ActiveUploadReservedBytesTx(tx)
		if err != nil {
			return err
		}
		if reserved != base.TotalSize {
			t.Fatalf("core reserved bytes = %d, want %d", reserved, base.TotalSize)
		}
		return nil
	}); err != nil {
		t.Fatalf("read core reserved bytes: %v", err)
	}

	duplicate := newLayerUpload("layer-upload-a-duplicate", "layer-a")
	if err := s.InsertUpload(ctx, duplicate); !errors.Is(err, ErrUploadConflict) {
		t.Fatalf("duplicate layer/path error = %v, want ErrUploadConflict", err)
	}

	got, err := s.GetFSLayerUploadByPath(ctx, "layer-a", "/same.bin")
	if err != nil {
		t.Fatalf("get layer upload: %v", err)
	}
	if got.UploadID != first.UploadID || got.TargetKind != UploadTargetFSLayer || got.LayerID != "layer-a" ||
		got.BaseInodeID != "base-inode-1" || got.ExpectedRevision == nil || *got.ExpectedRevision != 17 || got.Mode != 0o640 {
		t.Fatalf("layer upload round trip = %+v", got)
	}

	gotBase, err := s.GetUploadByPath(ctx, "/same.bin")
	if err != nil {
		t.Fatalf("get base upload: %v", err)
	}
	if gotBase.UploadID != base.UploadID || gotBase.TargetKind != UploadTargetFile || gotBase.LayerID != "" {
		t.Fatalf("base upload round trip = %+v", gotBase)
	}
}

func TestTransitionUploadStatusMapsActivePathDuplicateToConflict(t *testing.T) {
	s := newTestStore(t)
	ctx := context.Background()
	now := time.Now().UTC()
	newUpload := func(id string) *Upload {
		return &Upload{
			UploadID: id, FileID: id + "-object", TargetKind: UploadTargetFSLayer,
			LayerID: "layer-transition-conflict", TargetPath: "/same.bin",
			S3UploadID: "s3-" + id, S3Key: "layers/layer-transition-conflict/" + id,
			TotalSize: 8, PartSize: 8, PartsTotal: 1, Status: UploadInitiated,
			CreatedAt: now, UpdatedAt: now, ExpiresAt: now.Add(time.Hour),
		}
	}
	if err := s.InsertUpload(ctx, newUpload("transition-a")); err != nil {
		t.Fatalf("insert first upload: %v", err)
	}
	if err := s.InsertUpload(ctx, newUpload("transition-b")); err != nil {
		t.Fatalf("insert second initiated upload: %v", err)
	}
	if err := s.TransitionUploadStatus(ctx, "transition-a", UploadInitiated, UploadUploading); err != nil {
		t.Fatalf("transition first upload: %v", err)
	}
	if err := s.TransitionUploadStatus(ctx, "transition-b", UploadInitiated, UploadUploading); !errors.Is(err, ErrUploadConflict) {
		t.Fatalf("transition duplicate upload error=%v, want ErrUploadConflict", err)
	}
}

func TestUpdateFileContentTxClearsEmbeddingState(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	if err := s.InsertFile(context.Background(), &File{
		FileID:      "f1",
		StorageType: StorageDB9,
		StorageRef:  "/blobs/f1",
		ContentType: "text/plain",
		SizeBytes:   10,
		Revision:    1,
		Status:      StatusConfirmed,
		CreatedAt:   now,
		ConfirmedAt: &now,
	}); err != nil {
		t.Fatal(err)
	}
	setEmbeddingRevision(t, s, "f1", 9)

	var newRev int64
	if err := s.InTx(context.Background(), func(tx *sql.Tx) error {
		var err error
		newRev, err = s.UpdateFileContentTx(tx, "f1", StorageDB9, "/blobs/f1-v2", "text/markdown", "sum-2", "new text", []byte("blob-2"), 22, "")
		return err
	}); err != nil {
		t.Fatal(err)
	}
	if newRev != 2 {
		t.Fatalf("new revision=%d, want 2", newRev)
	}

	got, err := s.GetFile(context.Background(), "f1")
	if err != nil {
		t.Fatal(err)
	}
	if got.Revision != 2 || got.StorageRef != "/blobs/f1-v2" || got.ContentText != "new text" {
		t.Fatalf("unexpected updated file: %+v", got)
	}
	if got.EmbeddingRevision != nil {
		t.Fatalf("embedding revision should be cleared, got %v", *got.EmbeddingRevision)
	}
}

func TestUpdateFileContentTxRejectsAppendLogLayout(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	if err := s.InsertFile(context.Background(), &File{
		FileID:        "append-log-file",
		StorageType:   StorageS3,
		ContentLayout: ContentLayoutAppendLog,
		StorageRef:    "append-log/ns-a/file-a/generation-a",
		SizeBytes:     4,
		Revision:      1,
		Status:        StatusConfirmed,
		CreatedAt:     now,
		ConfirmedAt:   &now,
	}); err != nil {
		t.Fatal(err)
	}
	err := s.InTx(context.Background(), func(tx *sql.Tx) error {
		_, err := s.UpdateFileContentTx(tx, "append-log-file", StorageS3, "blobs/replacement", "text/plain", "", "", nil, 4, "")
		return err
	})
	if !errors.Is(err, ErrContentLayoutConflict) {
		t.Fatalf("UpdateFileContentTx error = %v, want ErrContentLayoutConflict", err)
	}
	got, err := s.GetFile(context.Background(), "append-log-file")
	if err != nil {
		t.Fatal(err)
	}
	if got.ContentLayout != ContentLayoutAppendLog || got.StorageRef != "append-log/ns-a/file-a/generation-a" {
		t.Fatalf("append-log file changed after rejected generic update: %+v", got)
	}
}

func TestUpdateFileContentTxAllowsIdenticalSingleContent(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	data := []byte("same bytes")
	if err := s.InsertFile(context.Background(), &File{
		FileID:      "single-file",
		StorageType: StorageDB9,
		StorageRef:  "inline",
		ContentBlob: data,
		ContentType: "text/plain",
		SizeBytes:   int64(len(data)),
		Revision:    1,
		Status:      StatusConfirmed,
		CreatedAt:   now,
		ConfirmedAt: &now,
	}); err != nil {
		t.Fatal(err)
	}
	var revision int64
	if err := s.InTx(context.Background(), func(tx *sql.Tx) error {
		var err error
		revision, err = s.UpdateFileContentTx(tx, "single-file", StorageDB9, "inline", "text/plain", "", "", data, int64(len(data)), "")
		return err
	}); err != nil {
		t.Fatalf("identical single update: %v", err)
	}
	if revision != 2 {
		t.Fatalf("identical single revision = %d, want 2", revision)
	}
}

func TestUpdateFileStorageEncryptionTx(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	if err := s.InsertFile(context.Background(), &File{
		FileID:                "f1",
		StorageType:           StorageS3,
		StorageRef:            "/blobs/f1",
		StorageEncryptionMode: StorageEncryptionLegacy,
		SizeBytes:             10,
		Revision:              1,
		Status:                StatusConfirmed,
		CreatedAt:             now,
		ConfirmedAt:           &now,
	}); err != nil {
		t.Fatal(err)
	}

	if err := s.InTx(context.Background(), func(tx *sql.Tx) error {
		return s.UpdateFileStorageEncryptionTx(tx, "f1", StorageEncryptionSSEKMS, "arn:aws:kms:ap-southeast-1:123456789012:key/test")
	}); err != nil {
		t.Fatal(err)
	}

	got, err := s.GetFile(context.Background(), "f1")
	if err != nil {
		t.Fatal(err)
	}
	if got.StorageEncryptionMode != StorageEncryptionSSEKMS {
		t.Fatalf("storage_encryption_mode=%q, want %q", got.StorageEncryptionMode, StorageEncryptionSSEKMS)
	}
	if got.StorageEncryptionKeyID != "arn:aws:kms:ap-southeast-1:123456789012:key/test" {
		t.Fatalf("storage_encryption_key_id=%q", got.StorageEncryptionKeyID)
	}
}

func TestUpdateFileStorageEncryptionTxSameValues(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	if err := s.InsertFile(context.Background(), &File{
		FileID:                "f1",
		StorageType:           StorageDB9,
		StorageRef:            "/blobs/f1",
		StorageEncryptionMode: StorageEncryptionNone,
		SizeBytes:             10,
		Revision:              1,
		Status:                StatusConfirmed,
		CreatedAt:             now,
		ConfirmedAt:           &now,
	}); err != nil {
		t.Fatal(err)
	}

	if err := s.InTx(context.Background(), func(tx *sql.Tx) error {
		return s.UpdateFileStorageEncryptionTx(tx, "f1", StorageEncryptionNone, "")
	}); err != nil {
		t.Fatal(err)
	}
}

func TestUpdateFileStorageEncryptionTxClearsKeyIDForNonKMSMode(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	if err := s.InsertFile(context.Background(), &File{
		FileID:                 "f1",
		StorageType:            StorageS3,
		StorageRef:             "/blobs/f1",
		StorageEncryptionMode:  StorageEncryptionSSEKMS,
		StorageEncryptionKeyID: "arn:aws:kms:ap-southeast-1:123456789012:key/test",
		SizeBytes:              10,
		Revision:               1,
		Status:                 StatusConfirmed,
		CreatedAt:              now,
		ConfirmedAt:            &now,
	}); err != nil {
		t.Fatal(err)
	}

	if err := s.InTx(context.Background(), func(tx *sql.Tx) error {
		return s.UpdateFileStorageEncryptionTx(tx, "f1", StorageEncryptionNone, "stale-key")
	}); err != nil {
		t.Fatal(err)
	}

	got, err := s.GetFile(context.Background(), "f1")
	if err != nil {
		t.Fatal(err)
	}
	if got.StorageEncryptionMode != StorageEncryptionNone {
		t.Fatalf("storage_encryption_mode=%q, want %q", got.StorageEncryptionMode, StorageEncryptionNone)
	}
	if got.StorageEncryptionKeyID != "" {
		t.Fatalf("storage_encryption_key_id=%q, want empty", got.StorageEncryptionKeyID)
	}
}

func TestUpdateFileStorageEncryptionTxMissingFile(t *testing.T) {
	s := newTestStore(t)
	err := s.InTx(context.Background(), func(tx *sql.Tx) error {
		return s.UpdateFileStorageEncryptionTx(tx, "missing", StorageEncryptionNone, "")
	})
	if !errors.Is(err, ErrNotFound) {
		t.Fatalf("err=%v, want %v", err, ErrNotFound)
	}
}

func TestClearFileEmbeddingStateTx(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	if err := s.InsertFile(context.Background(), &File{
		FileID:      "f1",
		StorageType: StorageDB9,
		StorageRef:  "/blobs/f1",
		Revision:    1,
		Status:      StatusConfirmed,
		CreatedAt:   now,
		ConfirmedAt: &now,
	}); err != nil {
		t.Fatal(err)
	}
	setEmbeddingRevision(t, s, "f1", 12)

	if err := s.InTx(context.Background(), func(tx *sql.Tx) error {
		return s.ClearFileEmbeddingStateTx(tx, "f1")
	}); err != nil {
		t.Fatal(err)
	}

	got, err := s.GetFile(context.Background(), "f1")
	if err != nil {
		t.Fatal(err)
	}
	if got.EmbeddingRevision != nil {
		t.Fatalf("embedding revision should be nil after clear, got %v", *got.EmbeddingRevision)
	}
}

func TestClearFileEmbeddingStateTxAlreadyClear(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	if err := s.InsertFile(context.Background(), &File{
		FileID:      "f-clear",
		StorageType: StorageDB9,
		StorageRef:  "/blobs/f-clear",
		Revision:    1,
		Status:      StatusConfirmed,
		CreatedAt:   now,
		ConfirmedAt: &now,
	}); err != nil {
		t.Fatal(err)
	}
	if err := s.InTx(context.Background(), func(tx *sql.Tx) error {
		return s.ClearFileEmbeddingStateTx(tx, "f-clear")
	}); err != nil {
		t.Fatalf("clear already-empty embedding state: %v", err)
	}
}

func TestUpdateFileContentAutoEmbeddingTxPreservesEmbeddingState(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	if err := s.InsertFile(context.Background(), &File{
		FileID:      "f1",
		StorageType: StorageDB9,
		StorageRef:  "/blobs/f1",
		ContentType: "text/plain",
		SizeBytes:   10,
		Revision:    1,
		Status:      StatusConfirmed,
		CreatedAt:   now,
		ConfirmedAt: &now,
	}); err != nil {
		t.Fatal(err)
	}
	setEmbeddingRevision(t, s, "f1", 9)

	var newRev int64
	if err := s.InTx(context.Background(), func(tx *sql.Tx) error {
		var err error
		newRev, err = s.UpdateFileContentAutoEmbeddingTx(tx, "f1", StorageDB9, "/blobs/f1-v2", "text/markdown", "sum-2", "new text", []byte("blob-2"), 22, "")
		return err
	}); err != nil {
		t.Fatal(err)
	}
	if newRev != 2 {
		t.Fatalf("new revision=%d, want 2", newRev)
	}

	got, err := s.GetFile(context.Background(), "f1")
	if err != nil {
		t.Fatal(err)
	}
	if got.Revision != 2 || got.StorageRef != "/blobs/f1-v2" || got.ContentText != "new text" {
		t.Fatalf("unexpected updated file: %+v", got)
	}
	if got.EmbeddingRevision == nil || *got.EmbeddingRevision != 9 {
		t.Fatalf("embedding revision should be preserved, got %v", got.EmbeddingRevision)
	}
}

func TestEnqueueSemanticTaskTxRollsBackWithOuterTransaction(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	task := newSemanticTask("task-tx", "file-1", 1, now, now)

	err := s.InTx(context.Background(), func(tx *sql.Tx) error {
		created, err := s.EnqueueSemanticTaskTx(tx, task)
		if err != nil {
			return err
		}
		if !created {
			t.Fatalf("expected transactional enqueue to create task")
		}
		return context.Canceled
	})
	if err != context.Canceled {
		t.Fatalf("rollback error=%v, want %v", err, context.Canceled)
	}
	if count := countSemanticTasks(t, s); count != 0 {
		t.Fatalf("semantic task count after rollback=%d, want 0", count)
	}

	if err := s.InTx(context.Background(), func(tx *sql.Tx) error {
		created, err := s.EnqueueSemanticTaskTx(tx, &semantic.Task{
			TaskID:          "task-commit",
			TaskType:        semantic.TaskTypeEmbed,
			ResourceID:      "file-1",
			ResourceVersion: 2,
			Status:          semantic.TaskQueued,
			MaxAttempts:     3,
			AvailableAt:     now,
			CreatedAt:       now,
			UpdatedAt:       now,
		})
		if err != nil {
			return err
		}
		if !created {
			t.Fatalf("expected committed transactional enqueue to create task")
		}
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	if count := countSemanticTasks(t, s); count != 1 {
		t.Fatalf("semantic task count after commit=%d, want 1", count)
	}
}

func TestUpdateFileContentTxWithDescription(t *testing.T) {
	s := newTestStore(t)
	now := time.Now().UTC()
	if err := s.InsertFile(context.Background(), &File{
		FileID:      "f1",
		StorageType: StorageDB9,
		StorageRef:  "/blobs/f1",
		Revision:    1,
		Status:      StatusConfirmed,
		CreatedAt:   now,
		ConfirmedAt: &now,
	}); err != nil {
		t.Fatal(err)
	}

	var newRev int64
	if err := s.InTx(context.Background(), func(tx *sql.Tx) error {
		var err error
		newRev, err = s.UpdateFileContentTx(tx, "f1", StorageDB9, "/blobs/f1-v2", "text/markdown", "sum-2", "new text", []byte("blob-2"), 22, "my description")
		return err
	}); err != nil {
		t.Fatal(err)
	}
	if newRev != 2 {
		t.Fatalf("new revision=%d, want 2", newRev)
	}

	got, err := s.GetFile(context.Background(), "f1")
	if err != nil {
		t.Fatal(err)
	}
	if got.Description != "my description" {
		t.Fatalf("description=%q, want %q", got.Description, "my description")
	}
	if got.DescriptionEmbeddingRevision != nil {
		t.Fatalf("description_embedding_revision should be cleared, got %v", got.DescriptionEmbeddingRevision)
	}
}