package datastore

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

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

// UpdateFileEmbedding conditionally writes an embedding for the current file revision.
func (s *Store) UpdateFileEmbedding(ctx context.Context, fileID string, revision int64, vector []float32) (updated bool, err error) {
	return s.updateFileEmbedding(ctx, s.db, fileID, revision, nil, vector)
}

// UpdateFileEmbeddingForGeneration conditionally writes an embedding and its
// provider generation for the current file revision.
func (s *Store) UpdateFileEmbeddingForGeneration(ctx context.Context, fileID string, revision int64, generation uint64, vector []float32) (updated bool, err error) {
	return s.updateFileEmbedding(ctx, s.db, fileID, revision, &generation, vector)
}

// UpdateFileEmbeddingForTask writes a generation-pinned embedding only while
// the caller still owns the corresponding semantic task lease and content_text
// still matches the value sent to the embedding provider.
func (s *Store) UpdateFileEmbeddingForTask(ctx context.Context, fileID string, revision int64, generation uint64, taskID, receipt, expectedContentText string, vector []float32) (updated bool, err error) {
	return s.withEmbeddingTaskOwnership(ctx, taskID, receipt, func(tx *sql.Tx) (bool, error) {
		var currentContentText sql.NullString
		if err := tx.QueryRowContext(ctx, `SELECT content_text FROM semantic WHERE `+
			s.scope.And(`inode_id = ?`)+` FOR UPDATE`, s.scope.Args(fileID)...).Scan(&currentContentText); err != nil {
			if errors.Is(err, sql.ErrNoRows) {
				return false, nil
			}
			return false, fmt.Errorf("lock semantic content_text: %w", err)
		}
		if currentContentText.String != expectedContentText {
			return false, nil
		}
		return s.updateFileEmbedding(ctx, tx, fileID, revision, &generation, vector)
	})
}

func (s *Store) updateFileEmbedding(ctx context.Context, db execer, fileID string, revision int64, generation *uint64, vector []float32) (updated bool, err error) {
	start := time.Now()
	defer observeStoreOp(ctx, "update_file_embedding", start, &err)

	if len(vector) == 0 {
		return false, fmt.Errorf("embedding vector is required")
	}
	if s.useLegacyFiles {
		res, err := db.ExecContext(ctx, `UPDATE files SET embedding = ?, embedding_revision = ?
			WHERE file_id = ? AND revision = ? AND status = 'CONFIRMED'`,
			embedding.FormatVector(vector), revision, fileID, revision)
		if err != nil {
			return false, err
		}
		rowsAffected, _ := res.RowsAffected()
		// Always attempt semantic update regardless of legacy rowsAffected,
		// so a retry after transient semantic failure is not skipped.
		if _, err := db.ExecContext(ctx, `UPDATE semantic SET embedding = ?, embedding_revision = ?, embedding_generation = ?
			WHERE `+s.scope.And(`inode_id = ? AND EXISTS (SELECT 1 FROM inodes WHERE `+s.scope.And(`inode_id = ? AND revision = ? AND status = 'CONFIRMED'`)+`)`),
			append([]any{embedding.FormatVector(vector), revision, generation}, append(s.scope.Args(fileID), s.scope.Args(fileID, revision)...)...)...); err != nil {
			return false, fmt.Errorf("update semantic embedding: %w", err)
		}
		return rowsAffected > 0, nil
	}
	// New tenant without legacy files: write directly to semantic.
	res, err := db.ExecContext(ctx, `UPDATE semantic SET embedding = ?, embedding_revision = ?, embedding_generation = ?
		WHERE `+s.scope.And(`inode_id = ? AND EXISTS (SELECT 1 FROM inodes WHERE `+s.scope.And(`inode_id = ? AND revision = ? AND status = 'CONFIRMED'`)+`)`),
		append([]any{embedding.FormatVector(vector), revision, generation}, append(s.scope.Args(fileID), s.scope.Args(fileID, revision)...)...)...)
	if err != nil {
		return false, fmt.Errorf("update semantic embedding: %w", err)
	}
	n, _ := res.RowsAffected()
	return n > 0, nil
}

// UpdateFileDescriptionEmbedding conditionally writes a description embedding for the current file revision.
func (s *Store) UpdateFileDescriptionEmbedding(ctx context.Context, fileID string, revision int64, vector []float32) (updated bool, err error) {
	return s.updateFileDescriptionEmbedding(ctx, s.db, fileID, revision, nil, vector)
}

// UpdateFileDescriptionEmbeddingForGeneration conditionally writes a
// description embedding and its provider generation for the current file
// revision.
func (s *Store) UpdateFileDescriptionEmbeddingForGeneration(ctx context.Context, fileID string, revision int64, generation uint64, vector []float32) (updated bool, err error) {
	return s.updateFileDescriptionEmbedding(ctx, s.db, fileID, revision, &generation, vector)
}

// UpdateFileDescriptionEmbeddingForTask writes a generation-pinned description
// embedding only while the caller still owns the semantic task lease and the
// stored description still matches the value sent to the embedding provider.
// Description can change without a revision bump (setmeta), so the text
// comparison is what prevents an in-flight task from stamping a vector for a
// stale description.
func (s *Store) UpdateFileDescriptionEmbeddingForTask(ctx context.Context, fileID string, revision int64, generation uint64, taskID, receipt, expectedDescription string, vector []float32) (updated bool, err error) {
	return s.withEmbeddingTaskOwnership(ctx, taskID, receipt, func(tx *sql.Tx) (bool, error) {
		var currentDescription sql.NullString
		if err := tx.QueryRowContext(ctx, `SELECT description FROM semantic WHERE `+
			s.scope.And(`inode_id = ?`)+` FOR UPDATE`, s.scope.Args(fileID)...).Scan(&currentDescription); err != nil {
			if errors.Is(err, sql.ErrNoRows) {
				return false, nil
			}
			return false, fmt.Errorf("lock semantic description: %w", err)
		}
		if currentDescription.String != expectedDescription {
			return false, nil
		}
		return s.updateFileDescriptionEmbedding(ctx, tx, fileID, revision, &generation, vector)
	})
}

func (s *Store) updateFileDescriptionEmbedding(ctx context.Context, db execer, fileID string, revision int64, generation *uint64, vector []float32) (updated bool, err error) {
	start := time.Now()
	defer observeStoreOp(ctx, "update_file_description_embedding", start, &err)

	if len(vector) == 0 {
		return false, fmt.Errorf("embedding vector is required")
	}
	if s.useLegacyFiles {
		res, err := db.ExecContext(ctx, `UPDATE files SET description_embedding = ?, description_embedding_revision = ?
			WHERE file_id = ? AND revision = ? AND status = 'CONFIRMED'`,
			embedding.FormatVector(vector), revision, fileID, revision)
		if err != nil {
			return false, err
		}
		rowsAffected, _ := res.RowsAffected()
		// Always attempt semantic update regardless of legacy rowsAffected.
		if _, err := db.ExecContext(ctx, `UPDATE semantic SET description_embedding = ?, description_embedding_revision = ?, description_embedding_generation = ?
			WHERE `+s.scope.And(`inode_id = ? AND EXISTS (SELECT 1 FROM inodes WHERE `+s.scope.And(`inode_id = ? AND revision = ? AND status = 'CONFIRMED'`)+`)`),
			append([]any{embedding.FormatVector(vector), revision, generation}, append(s.scope.Args(fileID), s.scope.Args(fileID, revision)...)...)...); err != nil {
			return false, fmt.Errorf("update semantic description_embedding: %w", err)
		}
		return rowsAffected > 0, nil
	}
	// New tenant without legacy files: write directly to semantic.
	res, err := db.ExecContext(ctx, `UPDATE semantic SET description_embedding = ?, description_embedding_revision = ?, description_embedding_generation = ?
		WHERE `+s.scope.And(`inode_id = ? AND EXISTS (SELECT 1 FROM inodes WHERE `+s.scope.And(`inode_id = ? AND revision = ? AND status = 'CONFIRMED'`)+`)`),
		append([]any{embedding.FormatVector(vector), revision, generation}, append(s.scope.Args(fileID), s.scope.Args(fileID, revision)...)...)...)
	if err != nil {
		return false, fmt.Errorf("update semantic description_embedding: %w", err)
	}
	n, _ := res.RowsAffected()
	return n > 0, nil
}

// withEmbeddingTaskOwnership runs write inside a transaction that first locks
// the semantic_tasks row (lease fence) and then, via write, the semantic row.
// Lock-order contract: transactions that race an in-flight embed task under
// its lease must take the two tables in the same order (semantic_tasks →
// semantic); see SetFileMetadataCtx. Note the extract writebacks
// (image/audio/video) predate this contract and still lock semantic before
// semantic_tasks via EnsureSemanticTaskQueuedTx; aligning them is a
// deliberate follow-up, not covered here.
func (s *Store) withEmbeddingTaskOwnership(ctx context.Context, taskID, receipt string, write func(*sql.Tx) (bool, error)) (bool, error) {
	if taskID == "" || receipt == "" {
		return false, fmt.Errorf("embedding task ownership is required")
	}
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return false, err
	}
	defer func() { _ = tx.Rollback() }()

	var owned int
	err = tx.QueryRowContext(ctx, `SELECT 1 FROM semantic_tasks WHERE `+
		s.scope.And(`task_id = ? AND task_type = ? AND status = ? AND receipt = ? AND lease_until IS NOT NULL AND lease_until > ?`)+`
		FOR UPDATE`,
		s.scope.Args(taskID, semantic.TaskTypeEmbed, semantic.TaskProcessing, receipt, semanticTaskLeaseNow())...).Scan(&owned)
	if errors.Is(err, sql.ErrNoRows) {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	updated, err := write(tx)
	if err != nil {
		return false, err
	}
	if err := tx.Commit(); err != nil {
		return false, err
	}
	return updated, nil
}