package backend

import (
	"bytes"
	"context"
	"database/sql"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/pingcap/failpoint"

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

// ImageExtractRequest is the input to a pluggable image->text extractor.
type ImageExtractRequest struct {
	TenantID       string
	TiDBCloudOrgID string
	FileID         string
	Path           string
	ContentType    string
	Data           []byte
}

// ImageTextExtractor extracts searchable text from image bytes.
type ImageTextExtractor interface {
	ExtractImageText(ctx context.Context, req ImageExtractRequest) (string, ImageExtractUsage, error)
}

// ImageExtractTaskSpec carries the revision-scoped inputs needed to extract
// image text for one file version.
type ImageExtractTaskSpec struct {
	FileID      string
	Path        string
	ContentType string
	Revision    int64
}

// ImageExtractResult reports the outcome of one image extraction attempt.
type ImageExtractResult string

const (
	ImageExtractResultRuntimeNotConfigured ImageExtractResult = "runtime_not_configured"
	ImageExtractResultGetFileError         ImageExtractResult = "get_file_error"
	ImageExtractResultFileNotFound         ImageExtractResult = "file_not_found"
	ImageExtractResultNotConfirmed         ImageExtractResult = "not_confirmed"
	ImageExtractResultNotImage             ImageExtractResult = "not_image"
	ImageExtractResultStale                ImageExtractResult = "stale"
	ImageExtractResultLoadError            ImageExtractResult = "load_error"
	ImageExtractResultTooLarge             ImageExtractResult = "too_large"
	ImageExtractResultExtractError         ImageExtractResult = "extract_error"
	ImageExtractResultEmptyText            ImageExtractResult = "empty_text"
	ImageExtractResultUpdateError          ImageExtractResult = "update_error"
	ImageExtractResultWritten              ImageExtractResult = "written"
	ImageExtractResultBudgetExhausted      ImageExtractResult = "budget_exhausted"
)

func isImageContentType(contentType string) bool {
	contentType = strings.ToLower(strings.TrimSpace(contentType))
	return strings.HasPrefix(contentType, "image/")
}

// rasterImageMagicOK reports whether data begins with the magic bytes of a
// raster image format that OpenAI-compatible vision providers accept
// (jpeg/png/gif/webp). Other image/* content types (bmp/tiff/svg/ico/heic)
// intentionally match none of these. Only the leading magic bytes are
// checked: a valid prefix followed by garbage still passes.
func rasterImageMagicOK(data []byte) bool {
	switch {
	case len(data) >= 3 && data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF: // jpeg
		return true
	case bytes.HasPrefix(data, []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1A, '\n'}): // png
		return true
	case bytes.HasPrefix(data, []byte("GIF8")): // gif87a / gif89a
		return true
	case len(data) >= 12 && string(data[:4]) == "RIFF" && string(data[8:12]) == "WEBP": // webp
		return true
	}
	return false
}

var imageExtensions = map[string]string{
	".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg",
	".gif": "image/gif", ".bmp": "image/bmp", ".webp": "image/webp",
	".svg": "image/svg+xml", ".tiff": "image/tiff", ".tif": "image/tiff",
	".ico": "image/x-icon",
}

func contentTypeFromPath(path string) string {
	path = strings.ToLower(path)
	for ext, ct := range imageExtensions {
		if strings.HasSuffix(path, ext) {
			return ct
		}
	}
	return ""
}

// SupportsAsyncImageExtract reports whether this backend instance has the
// runtime dependencies needed to extract image text.
func (b *Dat9Backend) SupportsAsyncImageExtract() bool {
	return b.imageExtractEnabled && b.imageExtractor != nil
}

// SupportsImageExtractTasks reports whether img_extract_text tasks should be
// enqueued, registered, and claimed for this backend: true when the
// process-default runtime is wired (SupportsAsyncImageExtract) or when the
// resolver-backed capability is enabled (the server resolves a per-tenant
// runtime at execution time via ProcessImageExtractTaskWithRuntime).
func (b *Dat9Backend) SupportsImageExtractTasks() bool {
	return b.SupportsAsyncImageExtract() || b.imageExtractCapability
}

// ImageExtractRuntime carries the effective extraction dependencies for one
// task. A resolver-backed multi-tenant server resolves these per tenant and
// passes them explicitly; the process-default path builds them from the
// backend's configured async image extract options.
type ImageExtractRuntime struct {
	Extractor           ImageTextExtractor
	Timeout             time.Duration
	MaxSourceBytes      int64
	MaxImageBytes       int64
	MaxExtractTextBytes int
	// RequireRasterImageBytes gates the jpeg/png/gif/webp magic sniff that
	// runs before the extractor call. Set it for provider-backed extractors
	// (OpenAI-compatible endpoints accept exactly those formats; junk like
	// macOS AppleDouble sidecars would otherwise burn paid requests and
	// retries on a permanent 400). Leave it false for extractors that accept
	// arbitrary declared-image bytes, e.g. the basic metadata extractor.
	RequireRasterImageBytes bool
	// BeforeExtract, when non-nil, runs immediately before the provider
	// call. A non-nil error aborts the task as retryable without calling the
	// provider; workers use it to re-check cluster kill switches that may
	// have flipped while the task was resolved and its image loaded.
	BeforeExtract func() error
	Workspace     *MediaExtractWorkspace
}

// ProcessImageExtractTask runs the backend-owned image extraction logic for one
// revision-scoped task using the process-default runtime. Normal terminal
// business outcomes return a nil error and a non-empty result label. Runtime
// misconfiguration, misrouting, and transient failures return a retryable error.
func (b *Dat9Backend) ProcessImageExtractTask(ctx context.Context, task ImageExtractTaskSpec) (ImageExtractResult, error) {
	return b.ProcessImageExtractTaskWithGuard(ctx, task, nil)
}

// ProcessImageExtractTaskWithGuard is ProcessImageExtractTask with a guard
// invoked immediately before the provider call (nil = no check).
func (b *Dat9Backend) ProcessImageExtractTaskWithGuard(ctx context.Context, task ImageExtractTaskSpec, guard func() error) (ImageExtractResult, error) {
	if !b.SupportsAsyncImageExtract() {
		return ImageExtractResultRuntimeNotConfigured, fmt.Errorf("async image extract runtime not configured")
	}
	return b.ProcessImageExtractTaskWithRuntime(ctx, task, ImageExtractRuntime{
		Extractor:               b.imageExtractor,
		Timeout:                 b.imageExtractTimeout,
		MaxSourceBytes:          b.imageExtractMaxSourceSize,
		MaxImageBytes:           b.imageExtractMaxSize,
		MaxExtractTextBytes:     b.maxExtractTextBytes,
		RequireRasterImageBytes: extractorRequiresRasterImageBytes(b.imageExtractor),
		BeforeExtract:           guard,
		Workspace:               b.mediaExtractWorkspace,
	})
}

// extractorRequiresRasterImageBytes reports whether the extractor only accepts
// jpeg/png/gif/webp payloads. The basic metadata extractor works on any
// declared-image bytes, so the sniff must stay off for it; provider-backed
// extractors (including the OpenAI-compatible fallback wrapper) reject
// everything else with a permanent 400.
func extractorRequiresRasterImageBytes(e ImageTextExtractor) bool {
	_, basic := e.(*BasicImageTextExtractor)
	return !basic
}

// ProcessImageExtractTaskWithRuntime is ProcessImageExtractTask with an
// explicitly resolved runtime. It does not require the process-default
// extractor, so it works on backends where SupportsAsyncImageExtract is false
// (resolver-backed multi-tenant servers).
func (b *Dat9Backend) ProcessImageExtractTaskWithRuntime(ctx context.Context, task ImageExtractTaskSpec, rt ImageExtractRuntime) (ImageExtractResult, error) {
	if rt.MaxSourceBytes <= 0 {
		rt.MaxSourceBytes = b.imageExtractMaxSourceSize
	}
	if rt.MaxSourceBytes <= 0 {
		rt.MaxSourceBytes = DefaultImageExtractMaxSourceBytes
	}
	if rt.Extractor == nil {
		return ImageExtractResultRuntimeNotConfigured, fmt.Errorf("async image extract runtime not configured")
	}
	if rt.MaxImageBytes <= 0 {
		return ImageExtractResultRuntimeNotConfigured, fmt.Errorf("async image extract max image bytes not configured")
	}
	if b.monthlyLLMCostExceededCheck(ctx) {
		b.recordTenantOperation("llm_cost_budget", "process_skip", "budget_exhausted", 0)
		return ImageExtractResultBudgetExhausted, nil
	}

	f, err := b.store.GetFile(ctx, task.FileID)
	if err != nil {
		if errors.Is(err, datastore.ErrNotFound) {
			return ImageExtractResultFileNotFound, nil
		}
		return ImageExtractResultGetFileError, fmt.Errorf("get file: %w", err)
	}
	if f.Status != datastore.StatusConfirmed {
		return ImageExtractResultNotConfirmed, nil
	}
	ct := f.ContentType
	if ct == "" {
		ct = task.ContentType
	}
	if ct == "" {
		ct = contentTypeFromPath(task.Path)
	}
	if !isImageContentType(ct) {
		return ImageExtractResultNotImage, nil
	}
	if task.Revision > 0 && f.Revision != task.Revision {
		return ImageExtractResultStale, nil
	}
	if f.SizeBytes > rt.MaxSourceBytes {
		b.recordTenantOperation("image_extract", "source", "too_large", 0)
		return ImageExtractResultTooLarge, nil
	}

	normalizationSupported := supportedImageNormalizationContentType(ct)
	allowNormalization := normalizationSupported
	if f.StorageType == datastore.StorageS3 && rt.Workspace == nil {
		allowNormalization = false
	}
	oversizedByMetadata := f.SizeBytes > rt.MaxImageBytes
	var releaseNormalizeSlot func()
	releaseSlot := func() {
		if releaseNormalizeSlot != nil {
			releaseNormalizeSlot()
			releaseNormalizeSlot = nil
		}
	}
	if allowNormalization && oversizedByMetadata {
		var acquired bool
		releaseNormalizeSlot, acquired = acquireImageExtractNormalizeSlot(ctx)
		if !acquired {
			return ImageExtractResultExtractError, ctx.Err()
		}
		defer releaseSlot()
	}
	var taskWorkspace *MediaExtractTaskWorkspace
	closeTaskWorkspace := func() {
		if taskWorkspace != nil {
			_ = taskWorkspace.Close()
			taskWorkspace = nil
		}
	}
	defer closeTaskWorkspace()

	var data []byte
	if f.StorageType == datastore.StorageS3 && oversizedByMetadata && allowNormalization && rt.Workspace != nil {
		taskWorkspace, err = rt.Workspace.NewTask(ctx, task.FileID)
		if err != nil {
			return ImageExtractResultLoadError, fmt.Errorf("create image extract workspace: %w", err)
		}
		stagedSourcePath, _, stageErr := b.stageMediaFile(ctx, taskWorkspace, f, "source", rt.MaxSourceBytes)
		if stageErr != nil {
			if errors.Is(stageErr, ErrMediaExtractSourceTooLarge) {
				b.recordTenantOperation("image_extract", "source", "too_large", 0)
				return ImageExtractResultTooLarge, nil
			}
			b.recordTenantOperation("image_extract", "source", "load_error", 0)
			return ImageExtractResultLoadError, fmt.Errorf("stage image source: %w", stageErr)
		}
		if rt.RequireRasterImageBytes {
			magicOK, magicErr := rasterImageFileMagicOK(stagedSourcePath)
			if magicErr != nil {
				b.recordTenantOperation("image_extract", "source", "load_error", 0)
				return ImageExtractResultLoadError, fmt.Errorf("inspect staged image source: %w", magicErr)
			}
			if !magicOK {
				return ImageExtractResultNotImage, nil
			}
		}
		normalized, normalizedContentType, ok := normalizeImageFileForExtractLocked(ctx, stagedSourcePath, rt.MaxImageBytes)
		if !ok {
			b.recordTenantOperation("image_extract", "normalize", "too_large", 0)
			if err := ctx.Err(); err != nil {
				return ImageExtractResultExtractError, err
			}
			return ImageExtractResultTooLarge, nil
		}
		b.recordTenantOperation("image_extract", "normalize", "ok", 0)
		data = normalized
		if normalizedContentType != "" {
			ct = normalizedContentType
		}
		// The normalized provider payload is bounded and memory-backed. Remove
		// the staged source and release the heavy-media slot before network I/O.
		closeTaskWorkspace()
	} else {
		data, err = b.loadImageBytesForExtract(ctx, f, rt.MaxImageBytes, rt.MaxSourceBytes, allowNormalization && oversizedByMetadata)
		if err != nil {
			b.recordTenantOperation("image_extract", "source", "load_error", 0)
			return ImageExtractResultLoadError, fmt.Errorf("load image bytes: %w", err)
		}
		if len(data) == 0 {
			if normalizationSupported && f.StorageType == datastore.StorageS3 && rt.Workspace == nil && rt.MaxImageBytes < rt.MaxSourceBytes {
				b.recordTenantOperation("image_extract", "source", "workspace_required", 0)
				return ImageExtractResultLoadError, ErrMediaExtractWorkspaceRequired
			}
			b.recordTenantOperation("image_extract", "source", "too_large", 0)
			return ImageExtractResultTooLarge, nil
		}
		// Declared as an image but the bytes match no raster format the extractor
		// accepts (e.g. macOS AppleDouble sidecars synced as .jpg). Terminal: the
		// provider would reject them with a permanent 400 anyway, so skip without
		// calling it. A later overwrite re-enqueues a fresh task. The sniff only
		// looks at the leading magic bytes; a valid prefix followed by garbage
		// still reaches the extractor.
		if rt.RequireRasterImageBytes && !rasterImageMagicOK(data) {
			return ImageExtractResultNotImage, nil
		}
	}
	if int64(len(data)) > rt.MaxImageBytes {
		var normalized []byte
		var normalizedContentType string
		var ok bool
		if releaseNormalizeSlot != nil {
			normalized, normalizedContentType, ok = normalizeImageForExtractLocked(ctx, data, rt.MaxImageBytes)
		} else {
			normalized, normalizedContentType, ok = normalizeImageForExtract(ctx, data, rt.MaxImageBytes)
		}
		if !ok {
			b.recordTenantOperation("image_extract", "normalize", "too_large", 0)
			if err := ctx.Err(); err != nil {
				return ImageExtractResultExtractError, err
			}
			return ImageExtractResultTooLarge, nil
		}
		b.recordTenantOperation("image_extract", "normalize", "ok", 0)
		data = normalized
		if normalizedContentType != "" {
			ct = normalizedContentType
		}
	}
	releaseSlot()

	// Last guard before the provider call: kill switches that flipped while
	// the task was resolved and its image loaded still abort the call. The
	// guard error is returned unwrapped so callers can match their sentinel.
	if rt.BeforeExtract != nil {
		if err := rt.BeforeExtract(); err != nil {
			return ImageExtractResultExtractError, err
		}
	}

	taskCtx, cancel := context.WithTimeout(ctx, rt.Timeout)
	text, imageUsage, err := rt.Extractor.ExtractImageText(taskCtx, ImageExtractRequest{
		TenantID:       b.tenantID,
		TiDBCloudOrgID: b.tidbCloudOrgID,
		FileID:         task.FileID,
		Path:           task.Path,
		ContentType:    ct,
		Data:           data,
	})
	cancel()
	if err != nil {
		return ImageExtractResultExtractError, fmt.Errorf("extract image text: %w", err)
	}
	b.recordImageExtractUsage(task.FileID, imageUsage)
	writeback := prepareImageExtractWriteback(text, rt.MaxExtractTextBytes)
	if writeback.text == "" {
		return ImageExtractResultEmptyText, nil
	}

	var updated bool
	err = b.store.InTx(ctx, func(tx *sql.Tx) error {
		if err := injectedImageExtractWritebackError("imageExtractWritebackUpdateFileSearchTextError"); err != nil {
			return err
		}
		var txErr error
		updated, txErr = b.store.UpdateFileSearchTextTx(tx, task.FileID, task.Revision, writeback.text)
		if txErr != nil {
			return txErr
		}
		if !updated {
			return nil
		}
		if txErr = b.store.DeleteFileTagsBySourceTx(tx, task.FileID, datastore.TagSourceImageExtract); txErr != nil {
			return txErr
		}
		// Legacy cleanup: wipe leftover numbered drive9.image.* rows.
		if txErr = b.store.ReplaceFileTagsByPrefixTx(tx, task.FileID, imageExtractTagPrefix, nil); txErr != nil {
			return txErr
		}
		var skipped []string
		skipped, txErr = b.store.InsertFileTagsWithSourceTx(tx, task.FileID, datastore.TagSourceImageExtract, writeback.tags)
		if txErr != nil {
			return txErr
		}
		// Count each skipped key as one operation with zero duration (zero
		// durations skip the duration histogram, so only the operation
		// counter moves); never smuggle the count into a fake duration.
		// The count reflects user-tag conflicts observed at the pre-insert
		// read inside InsertFileTagsWithSourceTx; a user write committing
		// between that read and the conditional INSERT still wins but is
		// not counted, so the metric may undercount concurrent user writes.
		for range skipped {
			b.recordTenantOperation("image_extract", "attr_skip", "user_tag_conflict", 0)
		}
		if b.UsesDatabaseAutoEmbedding() || !b.appSemanticTasksEnabled {
			return nil
		}
		if err := injectedImageExtractWritebackError("imageExtractWritebackQueueEmbedTaskError"); err != nil {
			return err
		}
		_, txErr = b.store.EnsureSemanticTaskQueuedTx(tx, newEmbedTask(
			b.genID(), task.FileID, task.Revision, b.appEmbeddingGeneration,
			b.appEmbeddingConfigRevision, b.appEmbeddingTenantManaged, time.Now().UTC(),
		))
		return txErr
	})
	if err != nil {
		return ImageExtractResultUpdateError, fmt.Errorf("update file search text: %w", err)
	}
	if !updated {
		return ImageExtractResultStale, nil
	}
	return ImageExtractResultWritten, nil
}

func injectedImageExtractWritebackError(name string) error {
	var injected error
	failpoint.Inject(name, func(val failpoint.Value) {
		switch v := val.(type) {
		case string:
			injected = errors.New(v)
		case bool:
			if v {
				injected = fmt.Errorf("injected failpoint %s", name)
			}
		}
	})
	return injected
}

func (b *Dat9Backend) loadImageBytesForExtract(ctx context.Context, f *datastore.File, maxSize, maxSourceSize int64, allowNormalization bool) ([]byte, error) {
	if f.ContentLayout == datastore.ContentLayoutAppendLog {
		return nil, ErrAppendLogUnsupported
	}
	readLimit := maxSize
	if allowNormalization {
		readLimit = maxSourceSize
	} else if maxSourceSize > 0 && (readLimit <= 0 || maxSourceSize < readLimit) {
		readLimit = maxSourceSize
	}
	if readLimit > 0 && f.SizeBytes > readLimit {
		return nil, nil
	}
	if f.StorageType == datastore.StorageDB9 {
		if readLimit > 0 && int64(len(f.ContentBlob)) > readLimit {
			return nil, nil
		}
		return append([]byte(nil), f.ContentBlob...), nil
	}
	if f.StorageType != datastore.StorageS3 {
		return nil, fmt.Errorf("unsupported storage type: %s", f.StorageType)
	}
	if b.s3 == nil {
		return nil, fmt.Errorf("s3 client not configured")
	}
	rc, err := b.s3.GetObject(ctx, f.StorageRef)
	if err != nil {
		return nil, err
	}
	defer func() { _ = rc.Close() }()

	reader := io.Reader(rc)
	if readLimit > 0 {
		reader = io.LimitReader(rc, readLimit+1)
	}
	data, err := io.ReadAll(reader)
	if err != nil {
		return nil, err
	}
	if readLimit > 0 && int64(len(data)) > readLimit {
		return nil, nil
	}
	return data, nil
}

func rasterImageFileMagicOK(path string) (bool, error) {
	f, err := os.Open(path)
	if err != nil {
		return false, err
	}
	defer func() { _ = f.Close() }()
	buf := make([]byte, 12)
	n, err := io.ReadFull(f, buf)
	if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrUnexpectedEOF) {
		return false, err
	}
	return rasterImageMagicOK(buf[:n]), nil
}

func sanitizeExtractedText(in string, maxBytes int) string {
	in = strings.TrimSpace(in)
	if in == "" {
		return ""
	}
	in = strings.Join(strings.Fields(in), " ")
	if maxBytes <= 0 || len(in) <= maxBytes {
		return in
	}
	return truncateUTF8Bytes(in, maxBytes)
}

func truncateUTF8Bytes(in string, maxBytes int) string {
	var (
		total int
		out   strings.Builder
	)
	for _, r := range in {
		rb := utf8.RuneLen(r)
		if rb < 0 {
			continue
		}
		if total+rb > maxBytes {
			break
		}
		out.WriteRune(r)
		total += rb
	}
	return strings.TrimSpace(out.String())
}