package backend

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

	"github.com/mem9-ai/drive9/pkg/backend/internal/patchcopy"
	"github.com/mem9-ai/drive9/pkg/datastore"
	"github.com/mem9-ai/drive9/pkg/logger"
	"github.com/mem9-ai/drive9/pkg/pathutil"
	"github.com/mem9-ai/drive9/pkg/s3client"
	"go.uber.org/zap"
)

// PatchPlan is returned by InitiatePatchUpload. It tells the client which parts
// to upload (dirty) and provides presigned Range-GET URLs so the client can
// download the original data for those parts, merge local modifications, and
// upload the result. Unchanged parts are copied server-side via S3
// UploadPartCopy — no data flows through the client for those.
type PatchPlan struct {
	UploadID    string             `json:"upload_id"`
	PartSize    int64              `json:"part_size"`
	UploadParts []*PatchUploadPart `json:"upload_parts"` // parts the client must upload
	CopiedParts []int              `json:"copied_parts"` // part numbers copied server-side
}

// AppendPlan is returned by InitiateAppendUpload. It is a specialized patch
// plan for appending bytes to an existing large S3-backed file.
type AppendPlan struct {
	BaseSize int64 `json:"base_size"`
	PatchPlan
}

// PatchUploadPart describes one dirty part that the client needs to upload.
type PatchUploadPart struct {
	Number      int               `json:"number"`
	URL         string            `json:"url"`               // presigned PUT URL for uploading
	Size        int64             `json:"size"`              // expected part size
	Headers     map[string]string `json:"headers,omitempty"` // required headers for the PUT
	ExpiresAt   time.Time         `json:"expires_at"`
	ReadURL     string            `json:"read_url,omitempty"`     // presigned GET URL to download original part data (empty for parts beyond original file)
	ReadHeaders map[string]string `json:"read_headers,omitempty"` // required headers for the GET (e.g. Range, signed into the presigned URL)
}

const patchPartCopyConcurrency = 8

// InitiatePatchUpload creates a multipart upload for modifying an existing
// large file. Only the dirty parts are uploaded by the client; unchanged parts
// are copied server-side within S3 using UploadPartCopy.
//
// Parameters:
//   - path: the file to patch
//   - newSize: the total size of the file after patching
//   - dirtyParts: 1-based part numbers that the client has modified
func (b *Dat9Backend) InitiatePatchUpload(ctx context.Context, path string, newSize int64, dirtyParts []int, clientPartSize int64) (*PatchPlan, error) {
	return b.InitiatePatchUploadIfRevision(ctx, path, newSize, dirtyParts, clientPartSize, -1)
}

// InitiateAppendUpload starts an incremental append upload for an existing
// large S3-backed file.
func (b *Dat9Backend) InitiateAppendUpload(ctx context.Context, path string, appendSize int64, clientPartSize int64) (*AppendPlan, error) {
	return b.InitiateAppendUploadIfRevision(ctx, path, appendSize, clientPartSize, -1)
}

// InitiateAppendUploadIfRevision starts an incremental append upload with
// optional CAS semantics.
func (b *Dat9Backend) InitiateAppendUploadIfRevision(ctx context.Context, path string, appendSize int64, clientPartSize int64, expectedRevision int64) (*AppendPlan, error) {
	if appendSize <= 0 {
		return nil, fmt.Errorf("append_size must be positive")
	}

	path, err := pathutil.Canonicalize(path)
	if err != nil {
		return nil, fmt.Errorf("canonicalize append path %q: %w", path, err)
	}
	if err := rejectRootFileNodePath(path); err != nil {
		return nil, err
	}
	if b.s3 == nil {
		return nil, ErrS3NotConfigured
	}

	nf, err := b.store.Stat(ctx, path)
	if err != nil {
		return nil, fmt.Errorf("stat append target %q: %w", path, err)
	}
	if nf.File == nil || nf.File.StorageType != datastore.StorageS3 {
		return nil, fmt.Errorf("%w: %s", ErrNotS3Stored, path)
	}
	if nf.File.ContentLayout == datastore.ContentLayoutAppendLog {
		return nil, ErrAppendLogUnsupported
	}
	if expectedRevision == 0 {
		return nil, datastore.ErrRevisionConflict
	}
	if expectedRevision > 0 && nf.File.Revision != expectedRevision {
		return nil, datastore.ErrRevisionConflict
	}

	baseSize := nf.File.SizeBytes
	newSize := baseSize + appendSize
	if newSize < baseSize {
		return nil, fmt.Errorf("append_size overflows file size")
	}

	partSize := clientPartSize
	if partSize < s3client.MinPartSize {
		partSize = s3client.CalcAdaptivePartSize(baseSize)
	}
	if len(s3client.CalcParts(newSize, partSize)) > MaxMultipartParts {
		partSize = s3client.CalcAdaptivePartSize(newSize)
	}

	dirtyParts := appendDirtyPartNumbers(baseSize, newSize, partSize)
	plan, err := b.InitiatePatchUploadIfRevision(ctx, path, newSize, dirtyParts, partSize, expectedRevision)
	if err != nil {
		return nil, fmt.Errorf("initiate append patch upload for %q: %w", path, err)
	}
	return &AppendPlan{
		BaseSize:  baseSize,
		PatchPlan: *plan,
	}, nil
}

func appendDirtyPartNumbers(baseSize int64, newSize int64, partSize int64) []int {
	if newSize <= baseSize || partSize <= 0 {
		return nil
	}

	newParts := s3client.CalcParts(newSize, partSize)
	basePartCount := 0
	if baseSize > 0 {
		basePartCount = len(s3client.CalcParts(baseSize, partSize))
	}

	dirty := make([]int, 0, len(newParts)-basePartCount+1)
	if baseSize > 0 && baseSize%partSize != 0 {
		dirty = append(dirty, basePartCount)
	}
	start := basePartCount + 1
	for pn := start; pn <= len(newParts); pn++ {
		dirty = append(dirty, pn)
	}
	if len(dirty) == 0 && len(newParts) > 0 {
		dirty = append(dirty, len(newParts))
	}
	return dirty
}

// InitiatePatchUploadIfRevision starts a patch upload with optional CAS semantics.
func (b *Dat9Backend) InitiatePatchUploadIfRevision(ctx context.Context, path string, newSize int64, dirtyParts []int, clientPartSize int64, expectedRevision int64) (*PatchPlan, error) {
	start := time.Now()
	if err := b.ensureUploadSizeAllowed(newSize); err != nil {
		b.recordTenantOperation("backend", "patch_upload", "error", time.Since(start))
		return nil, err
	}
	if err := b.ensureFileSizeQuota(ctx, newSize); err != nil {
		b.recordTenantOperation("backend", "patch_upload", "error", time.Since(start))
		return nil, err
	}

	path, err := pathutil.Canonicalize(path)
	if err != nil {
		b.recordTenantOperation("backend", "patch_upload", "error", time.Since(start))
		return nil, err
	}
	if err := rejectRootFileNodePath(path); err != nil {
		b.recordTenantOperation("backend", "patch_upload", "error", time.Since(start))
		return nil, err
	}
	if b.s3 == nil {
		b.recordTenantOperation("backend", "patch_upload", "error", time.Since(start))
		return nil, ErrS3NotConfigured
	}

	// Look up existing file to get its S3 key
	nf, err := b.store.Stat(ctx, path)
	if err != nil {
		b.recordTenantOperation("backend", "patch_upload", "error", time.Since(start))
		return nil, fmt.Errorf("stat existing file: %w", err)
	}
	if nf.File == nil || nf.File.StorageType != datastore.StorageS3 {
		b.recordTenantOperation("backend", "patch_upload", "error", time.Since(start))
		return nil, fmt.Errorf("%w: %s", ErrNotS3Stored, path)
	}
	if nf.File.ContentLayout == datastore.ContentLayoutAppendLog {
		b.recordTenantOperation("backend", "patch_upload", "error", time.Since(start))
		return nil, ErrAppendLogUnsupported
	}
	if expectedRevision == 0 {
		b.recordTenantOperation("backend", "patch_upload", "conflict", time.Since(start))
		return nil, datastore.ErrRevisionConflict
	}
	if expectedRevision > 0 && nf.File.Revision != expectedRevision {
		b.recordTenantOperation("backend", "patch_upload", "conflict", time.Since(start))
		return nil, datastore.ErrRevisionConflict
	}

	sourceKey := nf.File.StorageRef
	origSize := nf.File.SizeBytes

	// Use client-provided part size if valid (>= MinPartSize); otherwise
	// fall back to fixed 8 MiB for backward compatibility with old clients
	// that compute dirty_parts at fixed 8 MiB boundaries.
	partSize := clientPartSize
	if partSize < s3client.MinPartSize {
		partSize = s3client.PartSize
	}
	if partSize > s3client.MaxPartSize {
		b.recordTenantOperation("backend", "patch_upload", "error", time.Since(start))
		return nil, fmt.Errorf("part_size %d exceeds S3 per-part limit of %d", partSize, s3client.MaxPartSize)
	}
	newParts := s3client.CalcParts(newSize, partSize)
	if len(newParts) > MaxMultipartParts {
		b.recordTenantOperation("backend", "patch_upload", "error", time.Since(start))
		return nil, fmt.Errorf("part_size %d produces %d parts for size %d, exceeds S3 limit of %d", partSize, len(newParts), newSize, MaxMultipartParts)
	}

	// Create new S3 multipart upload (new key — old object stays until confirm)
	fileID := b.genID()
	newS3Key := "blobs/" + fileID
	encOpts, encMode, encKeyID := b.s3WriteEncryption(newS3Key)

	// Use ChecksumAlgoNone: patch parts are assembled client-side after
	// presigning, so the checksum is not known at presign time.  Declaring
	// SHA-256 here would force every UploadPart to carry the header, but
	// the presigned URL does not sign it — causing S3 403/400.  See #555.
	mpu, err := b.s3.CreateMultipartUpload(ctx, newS3Key, s3client.ChecksumAlgoNone, encOpts)
	if err != nil {
		logger.Error(ctx, "backend_patch_upload_create_mpu_failed", zap.String("tenant_id", b.tenantID), zap.String("path", path), zap.Error(err))
		b.recordTenantOperation("backend", "patch_upload", "error", time.Since(start))
		return nil, fmt.Errorf("create multipart upload: %w", err)
	}

	// Build dirty set for O(1) lookup
	dirtySet := make(map[int]bool, len(dirtyParts))
	for _, p := range dirtyParts {
		dirtySet[p] = true
	}

	// How many parts did the original file have?
	origPartCount := 0
	if origSize > 0 {
		origPartCount = len(s3client.CalcParts(origSize, partSize))
	}

	plan := &PatchPlan{
		UploadID: "", // set below after DB insert
		PartSize: partSize,
	}
	copyTasks := make([]patchcopy.Task, 0, len(newParts))

	// Process each part
	for _, p := range newParts {
		if !dirtySet[p.Number] && p.Number <= origPartCount {
			// Unchanged part within original file range → server-side copy
			partStart := int64(p.Number-1) * partSize
			partEnd := partStart + p.Size - 1
			// Clamp to original file size (last part may be smaller)
			if partEnd >= origSize {
				partEnd = origSize - 1
			}
			copyTasks = append(copyTasks, patchcopy.Task{
				PartNumber: p.Number,
				StartByte:  partStart,
				EndByte:    partEnd,
			})
		} else {
			// Dirty part or new part beyond original → client must upload
			u, err := b.s3.PresignUploadPart(ctx, newS3Key, mpu.UploadID, p.Number, p.Size, s3client.ChecksumAlgoNone, "", s3client.UploadTTL)
			if err != nil {
				_ = patchcopy.Abort(ctx, b.s3, newS3Key, mpu.UploadID, postS3UploadFinalizeTimeout)
				logger.Error(ctx, "backend_patch_upload_presign_failed", zap.String("tenant_id", b.tenantID), zap.String("path", path), zap.Int("part", p.Number), zap.Error(err))
				b.recordTenantOperation("backend", "patch_upload", "error", time.Since(start))
				return nil, fmt.Errorf("presign part %d: %w", p.Number, err)
			}

			pup := &PatchUploadPart{
				Number:    u.Number,
				URL:       u.URL,
				Size:      u.Size,
				Headers:   u.Headers,
				ExpiresAt: u.ExpiresAt,
			}

			// If this part overlaps with the original file, provide a read URL
			// so the client can download the original data for merging.
			if p.Number <= origPartCount {
				partStart := int64(p.Number-1) * partSize
				partEnd := partStart + p.Size - 1
				if partEnd >= origSize {
					partEnd = origSize - 1
				}
				readURL, err := b.s3.PresignGetObjectRange(ctx, sourceKey, partStart, partEnd, s3client.DownloadTTL)
				if err != nil {
					logger.Warn(ctx, "backend_patch_upload_presign_read_failed", zap.String("tenant_id", b.tenantID), zap.String("path", path), zap.Int("part", p.Number), zap.Error(err))
					// Non-fatal: client can still upload the full part without merging
				} else {
					pup.ReadURL = readURL
					// AWS presigned URLs with Range bake the header into the
					// signature — the client MUST send the matching Range header.
					pup.ReadHeaders = map[string]string{
						"Range": fmt.Sprintf("bytes=%d-%d", partStart, partEnd),
					}
				}
			}

			plan.UploadParts = append(plan.UploadParts, pup)
		}
	}
	if err := patchcopy.CopyOrAbort(
		ctx,
		b.s3,
		newS3Key,
		mpu.UploadID,
		sourceKey,
		copyTasks,
		patchPartCopyConcurrency,
		postS3UploadFinalizeTimeout,
	); err != nil {
		var partErr *patchcopy.PartError
		if errors.As(err, &partErr) {
			logger.Error(ctx, "backend_patch_upload_copy_failed", zap.String("tenant_id", b.tenantID), zap.String("path", path), zap.Int("part", partErr.PartNumber), zap.Error(err))
		} else {
			logger.Error(ctx, "backend_patch_upload_copy_failed", zap.String("tenant_id", b.tenantID), zap.String("path", path), zap.Error(err))
		}
		b.recordTenantOperation("backend", "patch_upload", "error", time.Since(start))
		return nil, fmt.Errorf("copy retained parts: %w", err)
	}
	for _, task := range copyTasks {
		plan.CopiedParts = append(plan.CopiedParts, task.PartNumber)
	}

	// Insert DB records (same pattern as InitiateUploadWithChecksums)
	now := time.Now()
	uploadID := b.genID()
	expiresAt := now.Add(24 * time.Hour)

	// Server-reserve-first saga (same as upload initiate).
	reserved, err := b.reserveUploadOnServer(ctx, uploadID, path, newSize, 0)
	if err != nil {
		_ = patchcopy.Abort(ctx, b.s3, newS3Key, mpu.UploadID, postS3UploadFinalizeTimeout)
		b.recordTenantOperation("backend", "patch_upload", "quota_exceeded", time.Since(start))
		return nil, err
	}

	// One active upload per path: supersede rather than 409 so dead sessions
	// can't block the path (see supersedeActiveUpload). Do this only after
	// part-size validation and S3 copy/presign setup have succeeded, so a
	// malformed replacement request does not tear down a healthy session.
	existing, err := b.activeUploadByPath(ctx, path)
	if err != nil {
		if reserved {
			b.abortUploadReservation(ctx, uploadID, newSize)
		}
		_ = patchcopy.Abort(ctx, b.s3, newS3Key, mpu.UploadID, postS3UploadFinalizeTimeout)
		b.recordTenantOperation("backend", "patch_upload", "error", time.Since(start))
		return nil, fmt.Errorf("lookup active upload for %s: %w", path, err)
	}
	if existing != nil {
		if err := b.supersedeActiveUpload(ctx, "patch_upload", existing); err != nil {
			if reserved {
				b.abortUploadReservation(ctx, uploadID, newSize)
			}
			_ = patchcopy.Abort(ctx, b.s3, newS3Key, mpu.UploadID, postS3UploadFinalizeTimeout)
			b.recordTenantOperation("backend", "patch_upload", "error", time.Since(start))
			return nil, fmt.Errorf("supersede active upload for %s: %w", path, err)
		}
	}

	if err := b.store.InTx(ctx, func(tx *sql.Tx) error {
		if !b.UseServerQuota() {
			if err := b.ensureTenantStorageQuotaTx(tx, path, newSize); err != nil {
				return err
			}
		}
		if err := b.store.InsertFileTx(tx, &datastore.File{
			FileID:                 fileID,
			StorageType:            datastore.StorageS3,
			StorageRef:             newS3Key,
			StorageEncryptionMode:  encMode,
			StorageEncryptionKeyID: encKeyID,
			SizeBytes:              newSize,
			Revision:               1,
			Status:                 datastore.StatusPending,
			CreatedAt:              now,
		}); err != nil {
			return err
		}
		return b.store.InsertUploadTx(tx, &datastore.Upload{
			UploadID:               uploadID,
			FileID:                 fileID,
			TargetPath:             path,
			S3UploadID:             mpu.UploadID,
			S3Key:                  newS3Key,
			StorageEncryptionMode:  encMode,
			StorageEncryptionKeyID: encKeyID,
			TotalSize:              newSize,
			PartSize:               partSize,
			PartsTotal:             len(newParts),
			ExpectedRevision:       expectedRevisionPtr(expectedRevision),
			Status:                 datastore.UploadUploading,
			CreatedAt:              now,
			UpdatedAt:              now,
			ExpiresAt:              expiresAt,
		})
	}); err != nil {
		if reserved {
			b.abortUploadReservation(ctx, uploadID, newSize)
		}
		_ = patchcopy.Abort(ctx, b.s3, newS3Key, mpu.UploadID, postS3UploadFinalizeTimeout)
		logger.Error(ctx, "backend_patch_upload_insert_upload_failed", zap.String("tenant_id", b.tenantID), zap.String("path", path), zap.Error(err))
		b.recordTenantOperation("backend", "patch_upload", "error", time.Since(start))
		return nil, err
	}

	plan.UploadID = uploadID
	b.recordTenantOperation("backend", "patch_upload", "ok", time.Since(start))

	logger.Info(ctx, "backend_patch_upload_initiated", zap.String("tenant_id", b.tenantID),
		zap.String("path", path),
		zap.Int64("new_size", newSize),
		zap.Int("total_parts", len(newParts)),
		zap.Int("dirty_parts", len(plan.UploadParts)),
		zap.Int("copied_parts", len(plan.CopiedParts)),
	)

	return plan, nil
}