package backend

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

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

const embeddingBackfillResourceID = "embedding-config"

// The enqueue helpers report whether a task row was actually created (the
// store dedupes identical queued tasks), so callers can fire the post-commit
// semantic task notifier only when there is new work to claim.
func (b *Dat9Backend) enqueueEmbedTaskTx(tx *sql.Tx, fileID string, revision int64) (bool, error) {
	now := time.Now().UTC()
	return b.store.EnqueueSemanticTaskTx(tx, newEmbedTask(
		b.genID(), fileID, revision, b.appEmbeddingGeneration,
		b.appEmbeddingConfigRevision, b.appEmbeddingTenantManaged, now,
	))
}

// forceRequeueEmbedTaskTx guarantees a queued embed task for the current
// revision, re-queuing terminal rows in place and resetting rows being
// processed under an active lease (clearing the receipt fails the in-flight
// owner's lease fence). Metadata-only updates (setmeta) change the
// description without bumping the revision, so neither the terminal-row
// dedupe of plain enqueue nor the active-lease exemption of ensure can
// apply: both would leave the new description without an embedding.
func (b *Dat9Backend) forceRequeueEmbedTaskTx(tx *sql.Tx, fileID string, revision int64) (bool, error) {
	now := time.Now().UTC()
	return b.store.ForceRequeueSemanticTaskTx(tx, newEmbedTask(
		b.genID(), fileID, revision, b.appEmbeddingGeneration,
		b.appEmbeddingConfigRevision, b.appEmbeddingTenantManaged, now,
	))
}

// EnqueueEmbeddingBackfillTask creates the one durable tenant-local backfill
// task for the backend's effective config revision. Repeated cold opens are
// deduplicated by semantic_tasks' existing unique key.
func (b *Dat9Backend) EnqueueEmbeddingBackfillTask(ctx context.Context) (bool, error) {
	runtime, ready := b.AppEmbeddingRuntime()
	if !ready {
		return false, nil
	}
	if runtime.ConfigRevision == 0 || runtime.ConfigRevision > uint64(1<<63-1) {
		return false, fmt.Errorf("invalid app embedding config revision %d", runtime.ConfigRevision)
	}
	now := time.Now().UTC()
	payloadJSON, _ := json.Marshal(semantic.EmbedBackfillTaskPayload{
		Generation: runtime.Generation, ConfigRevision: runtime.ConfigRevision,
	})
	created, err := b.store.EnqueueSemanticTask(ctx, &semantic.Task{
		TaskID: b.genID(), TaskType: semantic.TaskTypeEmbedBackfill,
		ResourceID: embeddingBackfillResourceID, ResourceVersion: int64(runtime.ConfigRevision),
		Status: semantic.TaskQueued, MaxAttempts: 5, AvailableAt: now,
		PayloadJSON: payloadJSON, CreatedAt: now, UpdatedAt: now,
	})
	if err != nil {
		return false, err
	}
	if created {
		b.notifyWorkEnqueued(BackendWorkSemantic)
	}
	return created, nil
}

// EnsureEmbeddingBackfillCandidates queues ordinary generation-pinned embed
// tasks for one bounded backfill page.
func (b *Dat9Backend) EnsureEmbeddingBackfillCandidates(ctx context.Context, candidates []datastore.EmbeddingBackfillCandidate) (bool, error) {
	tasks := make([]*semantic.Task, 0, len(candidates))
	now := time.Now().UTC()
	for _, candidate := range candidates {
		tasks = append(tasks, newEmbedTask(
			b.genID(), candidate.InodeID, candidate.Revision, b.appEmbeddingGeneration,
			b.appEmbeddingConfigRevision, b.appEmbeddingTenantManaged, now,
		))
	}
	queued, err := b.store.EnsureSemanticTasksQueued(ctx, tasks)
	if err != nil {
		return false, err
	}
	if queued {
		b.notifyWorkEnqueued(BackendWorkSemantic)
	}
	return queued, nil
}

func (b *Dat9Backend) enqueueImgExtractTaskTx(tx *sql.Tx, fileID string, revision int64, path, contentType string) (bool, error) {
	now := time.Now().UTC()
	task, err := newImgExtractTask(b.genID(), fileID, revision, path, contentType, now)
	if err != nil {
		return false, err
	}
	return b.store.EnqueueSemanticTaskTx(tx, task)
}

func (b *Dat9Backend) enqueueAudioExtractTaskTx(tx *sql.Tx, fileID string, revision int64, path, contentType string, videoFallback bool) (bool, error) {
	now := time.Now().UTC()
	task, err := newAudioExtractTask(b.genID(), fileID, revision, path, contentType, videoFallback, now)
	if err != nil {
		return false, err
	}
	return b.store.EnqueueSemanticTaskTx(tx, task)
}

func (b *Dat9Backend) enqueueVideoExtractTaskTx(tx *sql.Tx, fileID string, revision int64, path, contentType string) (bool, error) {
	now := time.Now().UTC()
	task, err := newVideoExtractTask(b.genID(), fileID, revision, path, contentType, now)
	if err != nil {
		return false, err
	}
	return b.store.EnqueueSemanticTaskTx(tx, task)
}

// enqueueExtractSemanticTasksTx registers durable img_extract_text and/or
// audio_extract_text tasks for one confirmed file revision. It applies in both
// database auto-embedding and app-embedding modes: image/audio extraction does
// not depend on EMBED_TEXT. When the tenant's media LLM file quota is exceeded,
// no extraction tasks are enqueued but the file write itself succeeds normally.
// currentMediaDelta accounts for the current transaction when server quota
// usage has not converged yet.
func (b *Dat9Backend) enqueueExtractSemanticTasksTx(ctx context.Context, tx *sql.Tx, fileID string, revision int64, path, contentType string, currentMediaDelta int64) (bool, error) {
	isImage := b.hasAsyncImageTextSource(path, contentType)
	isVideo := b.shouldEnqueueVideoExtractTask(path, contentType)
	// MP4 can carry independently configured audio and video extraction. The
	// audio task is marked as a fallback; after claim the worker discards it
	// when video extraction is executable, because the video runtime handles
	// both modalities and must remain the sole content_text writer.
	isAudio := b.shouldEnqueueAudioExtractTask(path, contentType)
	if !isImage && !isAudio && !isVideo {
		return false, nil
	}
	enqueued := false
	videoAdmitted := false
	// Video has its own independent quota. Decide this first so an MP4 audio
	// task is marked as fallback only when the corresponding video task was
	// actually admitted in this transaction.
	if isVideo {
		if b.videoLLMQuotaExceededTx(ctx, tx) {
			b.recordTenantOperation("video_llm_budget", "enqueue_skip", "quota_exceeded", 0)
		} else {
			created, err := b.enqueueVideoExtractTaskTx(tx, fileID, revision, path, contentType)
			if err != nil {
				return enqueued, err
			}
			videoAdmitted = true
			enqueued = enqueued || created
		}
	}
	// Image and audio share the general media LLM quota.
	if (isImage || isAudio) && b.mediaLLMQuotaExceededCheckTx(ctx, tx, currentMediaDelta) {
		b.recordTenantOperation("media_llm_budget", "enqueue_skip", "quota_exceeded", 0)
	} else {
		if isImage {
			created, err := b.enqueueImgExtractTaskTx(tx, fileID, revision, path, contentType)
			if err != nil {
				return enqueued, err
			}
			enqueued = enqueued || created
		}
		if isAudio {
			created, err := b.enqueueAudioExtractTaskTx(tx, fileID, revision, path, contentType, videoAdmitted)
			if err != nil {
				return enqueued, err
			}
			enqueued = enqueued || created
		}
	}
	return enqueued, nil
}

func (b *Dat9Backend) shouldEnqueueAudioExtractTask(path, contentType string) bool {
	if !b.SupportsAudioExtractTasks() {
		return false
	}
	return isSupportedAudioForSemanticTask(path, contentType)
}

func (b *Dat9Backend) shouldEnqueueVideoExtractTask(path, contentType string) bool {
	if !b.SupportsVideoExtractTasks() {
		return false
	}
	// Tenant check: "*" allows all; otherwise only explicitly listed tenants.
	if !b.videoExtractCapability && !b.videoExtractAllTenants {
		if _, ok := b.videoExtractTenantAllowlist[b.tenantID]; !ok {
			return false
		}
	}
	return isSupportedVideoForSemanticTask(path, contentType)
}

func newEmbedTask(taskID, fileID string, revision int64, generation, configRevision uint64, tenantManaged bool, now time.Time) *semantic.Task {
	now = now.UTC()
	payloadJSON, _ := json.Marshal(semantic.EmbedTaskPayload{
		Generation: generation, ConfigRevision: configRevision, TenantManaged: tenantManaged,
	})
	return &semantic.Task{
		TaskID:          taskID,
		TaskType:        semantic.TaskTypeEmbed,
		ResourceID:      fileID,
		ResourceVersion: revision,
		Status:          semantic.TaskQueued,
		MaxAttempts:     5,
		AvailableAt:     now,
		PayloadJSON:     payloadJSON,
		CreatedAt:       now,
		UpdatedAt:       now,
	}
}

func newImgExtractTask(taskID, fileID string, revision int64, path, contentType string, now time.Time) (*semantic.Task, error) {
	now = now.UTC()
	payload := semantic.ImgExtractTaskPayload{Path: path, ContentType: contentType}
	var payloadJSON []byte
	if payload.Path != "" || payload.ContentType != "" {
		encoded, err := json.Marshal(payload)
		if err != nil {
			return nil, err
		}
		payloadJSON = encoded
	}
	return &semantic.Task{
		TaskID:          taskID,
		TaskType:        semantic.TaskTypeImgExtractText,
		ResourceID:      fileID,
		ResourceVersion: revision,
		Status:          semantic.TaskQueued,
		MaxAttempts:     5,
		AvailableAt:     now,
		PayloadJSON:     payloadJSON,
		CreatedAt:       now,
		UpdatedAt:       now,
	}, nil
}

func newAudioExtractTask(taskID, fileID string, revision int64, path, contentType string, videoFallback bool, now time.Time) (*semantic.Task, error) {
	now = now.UTC()
	payload := semantic.AudioExtractTaskPayload{
		Path: path, ContentType: contentType,
		VideoFallback: videoFallback,
	}
	var payloadJSON []byte
	if payload.Path != "" || payload.ContentType != "" {
		encoded, err := json.Marshal(payload)
		if err != nil {
			return nil, err
		}
		payloadJSON = encoded
	}
	return &semantic.Task{
		TaskID:          taskID,
		TaskType:        semantic.TaskTypeAudioExtractText,
		ResourceID:      fileID,
		ResourceVersion: revision,
		Status:          semantic.TaskQueued,
		MaxAttempts:     5,
		AvailableAt:     now,
		PayloadJSON:     payloadJSON,
		CreatedAt:       now,
		UpdatedAt:       now,
	}, nil
}

func newVideoExtractTask(taskID, fileID string, revision int64, path, contentType string, now time.Time) (*semantic.Task, error) {
	now = now.UTC()
	payload := semantic.VideoExtractTaskPayload{Path: path, ContentType: contentType}
	var payloadJSON []byte
	if payload.Path != "" || payload.ContentType != "" {
		encoded, err := json.Marshal(payload)
		if err != nil {
			return nil, err
		}
		payloadJSON = encoded
	}
	return &semantic.Task{
		TaskID:          taskID,
		TaskType:        semantic.TaskTypeVideoExtractVisual,
		ResourceID:      fileID,
		ResourceVersion: revision,
		Status:          semantic.TaskQueued,
		MaxAttempts:     3,
		AvailableAt:     now,
		PayloadJSON:     payloadJSON,
		CreatedAt:       now,
		UpdatedAt:       now,
	}, nil
}

func (b *Dat9Backend) shouldEnqueueEmbedForRevision(path, contentType, contentText, description string) bool {
	if !b.appSemanticTasksEnabled {
		return false
	}
	if strings.TrimSpace(contentText) != "" {
		return true
	}
	if strings.TrimSpace(description) != "" {
		return true
	}
	return false
}

// SupportsAppEmbeddingTasks reports whether this backend has one resolved
// app-managed provider and may enqueue/execute durable embed work.
func (b *Dat9Backend) SupportsAppEmbeddingTasks() bool {
	return b != nil && b.appSemanticTasksEnabled && b.queryEmbedder != nil && b.appEmbeddingGeneration > 0 && b.appEmbeddingConfigRevision > 0
}

// AppEmbeddingRuntime returns the backend's immutable app-managed runtime.
func (b *Dat9Backend) AppEmbeddingRuntime() (AppEmbeddingRuntime, bool) {
	if !b.SupportsAppEmbeddingTasks() {
		return AppEmbeddingRuntime{}, false
	}
	return AppEmbeddingRuntime{
		Client: b.queryEmbedder, Generation: b.appEmbeddingGeneration,
		ConfigRevision: b.appEmbeddingConfigRevision, TenantManaged: b.appEmbeddingTenantManaged,
	}, true
}

func (b *Dat9Backend) hasAsyncImageTextSource(path, contentType string) bool {
	if !b.SupportsImageExtractTasks() {
		return false
	}
	if isImageContentType(contentType) {
		return true
	}
	return isImageContentType(contentTypeFromPath(path))
}

// AutoSemanticTaskTypes returns durable semantic task types executed on the
// database auto-embedding (TiDB auto) path: img_extract_text and/or
// audio_extract_text when the corresponding async runtimes are configured.
//
// Image registration is capability-based (SupportsImageExtractTasks): a
// resolver-backed backend with no process-default extractor still registers
// img_extract_text because the runtime is resolved per tenant at execution
// time.
//
// It does not include app-managed embed work; embed routing uses the worker
// embedder, not the backend. A nil return means this backend contributes no
// auto semantic tasks. The returned slice must be treated as read-only.
func (b *Dat9Backend) AutoSemanticTaskTypes() []semantic.TaskType {
	if b == nil || !b.UsesDatabaseAutoEmbedding() {
		return nil
	}
	var out []semantic.TaskType
	if b.SupportsImageExtractTasks() {
		out = append(out, semantic.TaskTypeImgExtractText)
	}
	if b.SupportsAudioExtractTasks() {
		out = append(out, semantic.TaskTypeAudioExtractText)
	}
	if b.SupportsVideoExtractTasks() {
		out = append(out, semantic.TaskTypeVideoExtractVisual)
	}
	if len(out) == 0 {
		return nil
	}
	return out
}