package backend

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/pingcap/failpoint"

	"github.com/mem9-ai/drive9/pkg/datastore"
	"github.com/mem9-ai/drive9/pkg/logger"
	"github.com/mem9-ai/drive9/pkg/meta"
	"go.uber.org/zap"
)

type fileCreateMutationData struct {
	FileID      string      `json:"file_id"`
	SizeBytes   int64       `json:"size_bytes"`
	IsMedia     bool        `json:"is_media"`
	StorageType StorageType `json:"storage_type,omitempty"`
	LayerOwned  bool        `json:"layer_owned,omitempty"`
}

type fileOverwriteMutationData struct {
	FileID       string      `json:"file_id"`
	OldSizeBytes int64       `json:"old_size_bytes"`
	OldIsMedia   bool        `json:"old_is_media"`
	NewSizeBytes int64       `json:"new_size_bytes"`
	NewIsMedia   bool        `json:"new_is_media"`
	StorageType  StorageType `json:"storage_type,omitempty"`
	LayerOwned   bool        `json:"layer_owned,omitempty"`
}

// fileDeleteMutationData is retained for replaying historical file_delete
// mutation log entries. New file deletes are cleaned up through file_gc_tasks.
type fileDeleteMutationData struct {
	FileID    string `json:"file_id"`
	SizeBytes int64  `json:"size_bytes"`
	IsMedia   bool   `json:"is_media"`
}

type fsLayerStorageMutationData struct {
	LayerID     string      `json:"layer_id"`
	PathHash    string      `json:"path_hash"`
	EntrySeq    int64       `json:"entry_seq"`
	SizeBytes   int64       `json:"size_bytes"`
	StorageType StorageType `json:"storage_type"`
}

type llmCostMutationData struct {
	TaskType       string `json:"task_type"`
	TaskID         string `json:"task_id"`
	CostMillicents int64  `json:"cost_millicents"`
	RawUnits       int64  `json:"raw_units"`
	RawUnitType    string `json:"raw_unit_type"`
}

const postCommitQuotaMutationTimeout = 30 * time.Second

// quotaCounterDeltas accumulates the tenant_quota_usage counter adjustments
// produced by one mutation apply. Applies RETURN their deltas instead of
// executing the per-counter Incr*Tx statements inline; the transaction owner
// (dispatcher batch runner, per-item applyQuotaMutation, or replayOne)
// aggregates them per tenant and flushes them at the END of the transaction
// via IncrQuotaUsageCountersTx — "hot row last" — so the single per-tenant
// tenant_quota_usage row is locked only just before commit instead of from
// the first mutation apply to commit, and N mutations for one tenant
// collapse into one UPDATE.
type quotaCounterDeltas struct {
	storageBytes   int64
	dbStorageBytes int64
	s3StorageBytes int64
	fileCount      int64
	mediaFileCount int64
	reservedBytes  int64
}

func (d *quotaCounterDeltas) add(o quotaCounterDeltas) {
	d.storageBytes += o.storageBytes
	d.dbStorageBytes += o.dbStorageBytes
	d.s3StorageBytes += o.s3StorageBytes
	d.fileCount += o.fileCount
	d.mediaFileCount += o.mediaFileCount
	d.reservedBytes += o.reservedBytes
}

// applyQuotaCounterDeltasTx flushes one tenant's aggregated counter deltas at
// the end of a mutation transaction. IncrQuotaUsageCountersTx is a no-op on
// all-zero deltas, so callers may invoke this unconditionally.
func applyQuotaCounterDeltasTx(store MetaQuotaStore, tx *sql.Tx, tenantID string, deltas quotaCounterDeltas) error {
	return store.IncrQuotaUsageCountersTx(tx, tenantID,
		deltas.storageBytes, deltas.dbStorageBytes, deltas.s3StorageBytes,
		deltas.fileCount, deltas.mediaFileCount, deltas.reservedBytes)
}

// PostCommitQuotaMutationError reports that the user-visible file mutation has
// already committed, but the central quota handoff failed afterward.
type PostCommitQuotaMutationError struct {
	Op  string
	Err error
}

func (e *PostCommitQuotaMutationError) Error() string {
	return fmt.Sprintf("%s after commit: %v", e.Op, e.Err)
}

func (e *PostCommitQuotaMutationError) Unwrap() error { return e.Err }

func postCommitQuotaMutationContext() (context.Context, context.CancelFunc) {
	return context.WithTimeout(backgroundWithTrace(), postCommitQuotaMutationTimeout)
}

func postCommitQuotaMutationError(op string, err error) error {
	if err == nil {
		return nil
	}
	return &PostCommitQuotaMutationError{Op: op, Err: err}
}

func isQuotaMediaContentType(contentType string) bool {
	return strings.HasPrefix(contentType, "image/") || strings.HasPrefix(contentType, "audio/")
}

// logQuotaMutation durably records a mutation in the central quota_mutation_log.
// This is the only runtime quota accounting handoff after tenant quota_outbox
// removal. Callers must surface an error instead of silently dropping the
// mutation. The current write paths still call this after the tenant DB commit,
// so a process crash between tenant commit and this insert remains a known
// reconciliation window.
func (b *Dat9Backend) logQuotaMutation(ctx context.Context, mutationType string, payload any) (int64, error) {
	if b.metaStore == nil || b.tenantID == "" {
		return 0, nil
	}
	data, err := json.Marshal(payload)
	if err != nil {
		logger.Warn(ctx, "central_quota_mutation_marshal_failed",
			zap.String("tenant_id", b.tenantID),
			zap.String("mutation_type", mutationType),
			zap.Error(err))
		return 0, err
	}
	logID, err := b.metaStore.InsertMutationLog(ctx, &MutationLogView{
		TenantID:     b.tenantID,
		MutationType: mutationType,
		MutationData: data,
	})
	if err != nil {
		logger.Warn(ctx, "central_quota_mutation_log_insert_failed",
			zap.String("tenant_id", b.tenantID),
			zap.String("mutation_type", mutationType),
			zap.Error(err))
		b.recordTenantOperation("central_quota", mutationType, "log_error", time.Duration(0))
		return 0, err
	}
	return logID, nil
}

// applyQuotaMutation applies a previously-logged mutation and marks it as
// applied. This is the async-safe half of the split mutation path — if it
// fails or never runs, MutationReplayWorker picks up the pending log entry.
// The apply returns its tenant_quota_usage counter deltas instead of
// executing them inline; they are flushed in one statement right before the
// mark ("hot row last"), keeping the semantics identical from the caller's
// perspective while shrinking the hot-row lock window to just before commit.
func (b *Dat9Backend) applyQuotaMutation(ctx context.Context, mutationType string, logID int64, pending quotaPendingDeltas, apply func(context.Context, *sql.Tx) (quotaCounterDeltas, error)) {
	if err := b.metaStore.InTx(ctx, func(tx *sql.Tx) error {
		deltas, err := apply(ctx, tx)
		if err != nil {
			return err
		}
		if err := applyQuotaCounterDeltasTx(b.metaStore, tx, b.tenantID, deltas); err != nil {
			return err
		}
		return b.metaStore.MarkMutationAppliedTx(tx, logID)
	}); err != nil {
		logger.Warn(ctx, "central_quota_mutation_apply_failed",
			zap.String("tenant_id", b.tenantID),
			zap.String("mutation_type", mutationType),
			zap.Int64("log_id", logID),
			zap.Error(err))
		b.recordTenantOperation("central_quota", mutationType, "pending", time.Duration(0))
		return
	}
	if b.quotaUsageCache != nil {
		b.quotaUsageCache.invalidate()
	}
	b.clearPendingCentralMutationDeltas(pending.storageDelta, pending.fileDelta, pending.mediaDelta)
	b.recordTenantOperation("central_quota", mutationType, "ok", time.Duration(0))
}

// logAndEnqueueMutation atomically logs a mutation and enqueues its apply
// function under mutationMu. This ensures that within a single backend
// instance, durable log_id order and dispatcher enqueue order are identical,
// preventing reordering between concurrent same-tenant writes on this
// process. The mutex scope is kept minimal: just the log insert (~1ms) +
// shard channel send (non-blocking while the 4096-slot shard buffer has
// capacity; a full buffer applies backpressure).
//
// Cross-instance ordering: in a multi-pod deployment, each pod has its own
// mutationMu and dispatcher. Two pods can apply mutations for the same
// tenant in different log_id order. This is a pre-existing condition — the
// the old synchronous log+apply path also had no cross-pod ordering.
// UpsertFileMetaTx is last-writer-wins; MutationReplayWorker replays
// pending (unapplied) entries in (tenant_id, id) order, which handles
// crash recovery. Cross-pod last-writer divergence on file_meta requires
// operational reconciliation from tenant file metadata.
func (b *Dat9Backend) logAndEnqueueMutation(ctx context.Context, mutationType string, payload any, pending quotaPendingDeltas, apply func(context.Context, *sql.Tx) (quotaCounterDeltas, error)) error {
	if b.metaStore == nil || b.tenantID == "" {
		return nil
	}
	start := time.Now()

	b.mutationMu.Lock()
	logID, err := b.logQuotaMutation(ctx, mutationType, payload)
	if err != nil {
		b.mutationMu.Unlock()
		return err
	}
	b.addPendingCentralMutationDeltas(pending.storageDelta, pending.fileDelta, pending.mediaDelta)
	b.enqueueMutationItem(ctx, quotaMutationItem{
		backend:        b,
		store:          b.metaStore,
		tenantID:       b.tenantID,
		tidbCloudOrgID: b.tidbCloudOrgID,
		mutationType:   mutationType,
		logID:          logID,
		pending:        pending,
		apply:          apply,
	})
	b.mutationMu.Unlock()

	logger.InfoBenchTiming(ctx, "central_quota_mutation_sync_timing",
		zap.String("tenant_id", b.tenantID),
		zap.String("mutation_type", mutationType),
		zap.Int64("log_id", logID),
		zap.Float64("total_ms", backendDurationMs(time.Since(start))))
	return nil
}

// applyCentralFileStateTx applies a create/overwrite shadow-state mutation:
// read the current tenant_file_meta row (FOR UPDATE), upsert it, and RETURN
// the resulting tenant_quota_usage counter deltas for the caller to flush at
// the end of the transaction ("hot row last" — see quotaCounterDeltas).
func applyCentralFileStateTx(store MetaQuotaStore, tx *sql.Tx, tenantID, fileID string, sizeBytes int64, isMedia bool, storageType StorageType) (quotaCounterDeltas, error) {
	return applyCentralFileStateWithOwnershipTx(store, tx, tenantID, fileID, sizeBytes, isMedia, storageType, false)
}

func applyCentralFileStateWithOwnershipTx(store MetaQuotaStore, tx *sql.Tx, tenantID, fileID string, sizeBytes int64, isMedia bool, storageType StorageType, layerOwned bool) (quotaCounterDeltas, error) {
	oldSize := int64(0)
	oldIsMedia := false
	oldExists := false
	old, err := store.GetFileMetaForUpdateTx(tx, tenantID, fileID)
	if err != nil {
		if !errors.Is(err, meta.ErrNotFound) {
			return quotaCounterDeltas{}, err
		}
	} else if old != nil {
		oldExists = true
		oldSize = old.SizeBytes
		oldIsMedia = old.IsMedia
	}
	// No-op fast path: the shadow row already matches the new state, so the
	// upsert would write identical values. Skipping it is safe because
	// UpsertFileMetaTx only maintains size_bytes/is_media/storage_type (no updated_at
	// upkeep) and nothing consumes tenant_file_meta.updated_at for
	// correctness — file GC settlement keys off HasPendingFileMutation plus
	// size/isMedia. This avoids a useless write for the overwhelming
	// majority of mutations (same-size overwrites). The mutation itself is
	// still marked applied by the caller.
	newStorageType := storageTypePtr(storageType)
	if oldExists && oldSize == sizeBytes && oldIsMedia == isMedia && equalStorageTypes(old.StorageType, newStorageType) && old.LayerOwned == layerOwned {
		return quotaCounterDeltas{}, nil
	}
	if err := store.UpsertFileMetaTx(tx, &FileMetaView{
		TenantID:    tenantID,
		FileID:      fileID,
		SizeBytes:   sizeBytes,
		IsMedia:     isMedia,
		StorageType: newStorageType,
		LayerOwned:  layerOwned,
	}); err != nil {
		return quotaCounterDeltas{}, err
	}
	oldStorageBytes := oldSize
	if oldExists && old.LayerOwned {
		oldStorageBytes = 0
	}
	newStorageBytes := sizeBytes
	if layerOwned {
		newStorageBytes = 0
	}
	deltas := quotaCounterDeltas{
		storageBytes:   newStorageBytes - oldStorageBytes,
		mediaFileCount: quotaMediaDelta(oldIsMedia, isMedia),
	}
	if !oldExists {
		deltas.fileCount = 1
	}
	if oldExists && old.StorageType != nil && !old.LayerOwned {
		deltas.addStorageTierBytes(*old.StorageType, -oldSize)
	}
	if newStorageType != nil && !layerOwned {
		deltas.addStorageTierBytes(*newStorageType, sizeBytes)
	}
	return deltas, nil
}

func storageTypePtr(storageType StorageType) *StorageType {
	if storageType != StorageTypeDB9 && storageType != StorageTypeS3 {
		return nil
	}
	value := storageType
	return &value
}

func equalStorageTypes(a, b *StorageType) bool {
	if a == nil || b == nil {
		return a == nil && b == nil
	}
	return *a == *b
}

func (d *quotaCounterDeltas) addStorageTierBytes(storageType StorageType, bytes int64) {
	switch storageType {
	case StorageTypeDB9:
		d.dbStorageBytes += bytes
	case StorageTypeS3:
		d.s3StorageBytes += bytes
	}
}

func applyCentralFileCreateTx(store MetaQuotaStore, tx *sql.Tx, tenantID string, data fileCreateMutationData) (quotaCounterDeltas, error) {
	return applyCentralFileStateWithOwnershipTx(store, tx, tenantID, data.FileID, data.SizeBytes, data.IsMedia, data.StorageType, data.LayerOwned)
}

func applyCentralFileOverwriteTx(store MetaQuotaStore, tx *sql.Tx, tenantID string, data fileOverwriteMutationData) (quotaCounterDeltas, error) {
	return applyCentralFileStateWithOwnershipTx(store, tx, tenantID, data.FileID, data.NewSizeBytes, data.NewIsMedia, data.StorageType, data.LayerOwned)
}

func (b *Dat9Backend) recordCentralFileCreateMutation(ctx context.Context, fileID string, sizeBytes int64, contentType string, storageType StorageType) error {
	return b.recordCentralFileCreateMutationWithLayerOwnership(ctx, fileID, sizeBytes, contentType, storageType, false)
}

func (b *Dat9Backend) recordCentralFileCreateMutationWithLayerOwnership(ctx context.Context, fileID string, sizeBytes int64, contentType string, storageType StorageType, layerOwned bool) error {
	isMedia := isQuotaMediaContentType(contentType)
	data := fileCreateMutationData{
		FileID:      fileID,
		SizeBytes:   sizeBytes,
		IsMedia:     isMedia,
		StorageType: storageType,
		LayerOwned:  layerOwned,
	}
	mediaDelta := int64(0)
	if isMedia {
		mediaDelta = 1
	}
	storageDelta := sizeBytes
	if layerOwned {
		storageDelta = 0
	}
	return b.logAndEnqueueMutation(ctx, "file_create", data, quotaPendingDeltas{
		storageDelta: storageDelta,
		fileDelta:    1,
		mediaDelta:   mediaDelta,
	}, func(applyCtx context.Context, tx *sql.Tx) (quotaCounterDeltas, error) {
		return applyCentralFileCreateTx(b.metaStore, tx, b.tenantID, data)
	})
}

func (b *Dat9Backend) recordCentralFileOverwriteMutation(ctx context.Context, fileID string, oldSize int64, oldContentType string, newSize int64, newContentType string, storageType StorageType) error {
	return b.recordCentralFileOverwriteMutationWithLayerOwnership(ctx, fileID, oldSize, oldContentType, newSize, newContentType, storageType, false)
}

func (b *Dat9Backend) recordCentralFileOverwriteMutationWithLayerOwnership(ctx context.Context, fileID string, oldSize int64, oldContentType string, newSize int64, newContentType string, storageType StorageType, layerOwned bool) error {
	var injected error
	failpoint.InjectCall("centralFileOverwriteQuotaFail", fileID, &injected)
	if injected != nil {
		return injected
	}
	oldIsMedia := isQuotaMediaContentType(oldContentType)
	newIsMedia := isQuotaMediaContentType(newContentType)
	data := fileOverwriteMutationData{
		FileID:       fileID,
		OldSizeBytes: oldSize,
		OldIsMedia:   oldIsMedia,
		NewSizeBytes: newSize,
		NewIsMedia:   newIsMedia,
		StorageType:  storageType,
		LayerOwned:   layerOwned,
	}
	storageDelta := newSize - oldSize
	if layerOwned {
		storageDelta = 0
	}
	return b.logAndEnqueueMutation(ctx, "file_overwrite", data, quotaPendingDeltas{
		storageDelta: storageDelta,
		mediaDelta:   quotaMediaDelta(oldIsMedia, newIsMedia),
	}, func(applyCtx context.Context, tx *sql.Tx) (quotaCounterDeltas, error) {
		return applyCentralFileOverwriteTx(b.metaStore, tx, b.tenantID, data)
	})
}

// RecordFSLayerStorage records one immutable fs-layer payload version in the
// central storage ledger. Metadata-only layer operations own no payload bytes.
// The tenant entry has already committed when this is called, so the durable
// central handoff deliberately outlives cancellation of the request context.
func (b *Dat9Backend) RecordFSLayerStorage(_ context.Context, entry *datastore.FSLayerEntry) error {
	if entry == nil || entry.Op != datastore.FSLayerEntryOpUpsert || entry.Kind != datastore.FSLayerEntryKindFile {
		return nil
	}
	quotaCtx, cancelQuota := postCommitQuotaMutationContext()
	defer cancelQuota()
	storageType := StorageType(entry.StorageType)
	if storageType == "" {
		storageType = StorageTypeDB9
	}
	data := fsLayerStorageMutationData{
		LayerID: entry.LayerID, PathHash: entry.PathHash, EntrySeq: entry.EntrySeq,
		SizeBytes: entry.SizeBytes, StorageType: storageType,
	}
	err := b.logAndEnqueueMutation(quotaCtx, "fs_layer_store", data, quotaPendingDeltas{
		storageDelta: entry.SizeBytes,
	}, func(applyCtx context.Context, tx *sql.Tx) (quotaCounterDeltas, error) {
		return applyFSLayerStorageTx(b.metaStore, tx, b.tenantID, data)
	})
	return postCommitQuotaMutationError("record central quota fs-layer storage", err)
}

// applyFSLayerStorageTx records one immutable layer payload idempotently and
// returns the physical storage/tier deltas. Both the live handoff and replay
// path use this helper so their billing behavior cannot diverge.
func applyFSLayerStorageTx(store MetaQuotaStore, tx *sql.Tx, tenantID string, data fsLayerStorageMutationData) (quotaCounterDeltas, error) {
	inserted, err := store.InsertLayerStorageMetaIfAbsentTx(tx, &LayerStorageMetaView{
		TenantID: tenantID, LayerID: data.LayerID, PathHash: data.PathHash,
		EntrySeq: data.EntrySeq, SizeBytes: data.SizeBytes, StorageType: classifiedStorageType(data.StorageType),
	})
	if err != nil || !inserted {
		return quotaCounterDeltas{}, err
	}
	deltas := quotaCounterDeltas{storageBytes: data.SizeBytes}
	deltas.addStorageTierBytes(data.StorageType, data.SizeBytes)
	return deltas, nil
}

func classifiedStorageType(storageType StorageType) StorageType {
	if storageType == StorageTypeDB9 || storageType == StorageTypeS3 {
		return storageType
	}
	return ""
}

func quotaStorageType(storageType datastore.StorageType) StorageType {
	switch storageType {
	case datastore.StorageDB9:
		return StorageTypeDB9
	case datastore.StorageS3:
		return StorageTypeS3
	default:
		return ""
	}
}

func fsLayerCoreStorageType(storageType datastore.StorageType) StorageType {
	if storageType == "" {
		return StorageTypeDB9
	}
	return quotaStorageType(storageType)
}

func (b *Dat9Backend) syncCentralLLMCostRecord(ctx context.Context, taskType, taskID string, costMillicents, rawUnits int64, rawUnitType string) error {
	return b.logAndEnqueueMutation(ctx, "llm_cost_record", llmCostMutationData{
		TaskType:       taskType,
		TaskID:         taskID,
		CostMillicents: costMillicents,
		RawUnits:       rawUnits,
		RawUnitType:    rawUnitType,
	}, quotaPendingDeltas{}, func(applyCtx context.Context, tx *sql.Tx) (quotaCounterDeltas, error) {
		// LLM cost writes go to tenant_llm_usage and tenant_monthly_llm_cost
		// — different tables than the tenant_quota_usage hot row, and rare —
		// so they stay inline and this apply returns zero counter deltas.
		if err := b.metaStore.InsertCentralLLMUsageTx(tx, &LLMUsageView{
			TenantID:       b.tenantID,
			TaskType:       taskType,
			TaskID:         taskID,
			CostMillicents: costMillicents,
			RawUnits:       rawUnits,
			RawUnitType:    rawUnitType,
		}); err != nil {
			return quotaCounterDeltas{}, err
		}
		if err := b.metaStore.IncrMonthlyLLMCostTx(tx, b.tenantID, costMillicents); err != nil {
			return quotaCounterDeltas{}, err
		}
		return quotaCounterDeltas{}, nil
	})
}