package memory

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

	sq "github.com/Masterminds/squirrel"
	"github.com/pgvector/pgvector-go"
)

// tableName is the PostgreSQL table backing memory records.
// The DML patterns (INSERT INTO memory_records, FROM memory_records,
// UPDATE memory_records, DELETE FROM memory_records) are generated
// by the squirrel query builder at runtime.
const tableName = "memory_records"

// errNotFoundFmt is the format string for not-found errors.
const errNotFoundFmt = "memory record not found: %s"

// SQL column names. Defined as constants because squirrel queries
// reference them in column lists, predicates, and updates — repeating
// the literals would mean drift if a column is ever renamed.
const (
	colCreatedBy      = "created_by"
	colCategory       = "category"
	colEntityURNs     = "entity_urns"
	colCreatedAt      = "created_at"
	colDimension      = "dimension"
	colConfidence     = "confidence"
	colMetadata       = "metadata"
	colPersona        = "persona"
	colContent        = "content"
	colRelatedColumns = "related_columns"
	colSource         = "source"
	colStatus         = "status"
)

// psq is the PostgreSQL statement builder with dollar placeholders.
var psq = sq.StatementBuilder.PlaceholderFormat(sq.Dollar)

// postgresStore implements Store using PostgreSQL with pgvector.
type postgresStore struct {
	db *sql.DB
}

// NewPostgresStore creates a new PostgreSQL memory store.
func NewPostgresStore(db *sql.DB) Store {
	return &postgresStore{db: db}
}

// Insert creates a new memory record.
func (s *postgresStore) Insert(ctx context.Context, record Record) error {
	entityURNs, err := json.Marshal(record.EntityURNs)
	if err != nil {
		return fmt.Errorf("marshaling entity_urns: %w", err)
	}

	relatedCols, err := json.Marshal(record.RelatedColumns)
	if err != nil {
		return fmt.Errorf("marshaling related_columns: %w", err)
	}

	metadata, err := json.Marshal(record.Metadata)
	if err != nil {
		return fmt.Errorf("marshaling metadata: %w", err)
	}

	columns := []string{
		"id", colCreatedBy, colPersona, colDimension,
		colContent, colCategory, colConfidence, colSource,
		colEntityURNs, colRelatedColumns, colMetadata, colStatus,
	}
	values := []any{
		record.ID, record.CreatedBy, record.Persona, record.Dimension,
		record.Content, record.Category, record.Confidence, record.Source,
		entityURNs, relatedCols, metadata, record.Status,
	}

	if len(record.Embedding) > 0 {
		columns = append(columns, "embedding")
		values = append(values, pgvector.NewVector(record.Embedding))
	}

	qb := psq.Insert(tableName).Columns(columns...).Values(values...)

	query, args, err := qb.ToSql()
	if err != nil {
		return fmt.Errorf("building insert query: %w", err)
	}

	_, err = s.db.ExecContext(ctx, query, args...)
	if err != nil {
		return fmt.Errorf("inserting memory record: %w", err)
	}

	return nil
}

// Get retrieves a single memory record by ID.
func (s *postgresStore) Get(ctx context.Context, id string) (*Record, error) {
	query, args, err := psq.Select(recordColumns()...).
		From(tableName).
		Where(sq.Eq{"id": id}).
		ToSql()
	if err != nil {
		return nil, fmt.Errorf("building get query: %w", err)
	}

	record, err := scanRecord(s.db.QueryRowContext(ctx, query, args...))
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, fmt.Errorf(errNotFoundFmt, id)
		}
		return nil, fmt.Errorf("querying memory record: %w", err)
	}

	return record, nil
}

// Update modifies fields on an existing memory record.
func (s *postgresStore) Update(ctx context.Context, id string, updates RecordUpdate) error {
	qb, hasUpdates, err := buildUpdateColumns(updates)
	if err != nil {
		return err
	}

	if !hasUpdates {
		return fmt.Errorf("no fields to update")
	}

	qb = qb.Set("updated_at", time.Now()).Where(sq.Eq{"id": id})

	query, args, err := qb.ToSql()
	if err != nil {
		return fmt.Errorf("building update query: %w", err)
	}

	result, err := s.db.ExecContext(ctx, query, args...)
	if err != nil {
		return fmt.Errorf("updating memory record: %w", err)
	}

	rows, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("checking rows affected: %w", err)
	}
	if rows == 0 {
		return fmt.Errorf(errNotFoundFmt, id)
	}

	return nil
}

// buildUpdateColumns applies non-empty fields from RecordUpdate to a squirrel UpdateBuilder.
func buildUpdateColumns(updates RecordUpdate) (sq.UpdateBuilder, bool, error) {
	qb := psq.Update(tableName)
	hasUpdates := false

	if updates.Content != "" {
		qb = qb.Set(colContent, updates.Content)
		hasUpdates = true
	}
	if updates.Category != "" {
		qb = qb.Set(colCategory, updates.Category)
		hasUpdates = true
	}
	if updates.Confidence != "" {
		qb = qb.Set(colConfidence, updates.Confidence)
		hasUpdates = true
	}
	if updates.Dimension != "" {
		qb = qb.Set(colDimension, updates.Dimension)
		hasUpdates = true
	}
	if updates.Metadata != nil {
		meta, err := json.Marshal(updates.Metadata)
		if err != nil {
			return qb, false, fmt.Errorf("marshaling metadata: %w", err)
		}
		qb = qb.Set(colMetadata, sq.Expr("metadata || ?::jsonb", meta))
		hasUpdates = true
	}
	if len(updates.Embedding) > 0 {
		qb = qb.Set("embedding", pgvector.NewVector(updates.Embedding))
		hasUpdates = true
	}

	return qb, hasUpdates, nil
}

// Delete soft-deletes a memory record by setting status to archived.
func (s *postgresStore) Delete(ctx context.Context, id string) error {
	query, args, err := psq.Update(tableName).
		Set(colStatus, StatusArchived).
		Set("updated_at", time.Now()).
		Where(sq.Eq{"id": id}).
		ToSql()
	if err != nil {
		return fmt.Errorf("building delete query: %w", err)
	}

	result, err := s.db.ExecContext(ctx, query, args...)
	if err != nil {
		return fmt.Errorf("archiving memory record: %w", err)
	}

	rows, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("checking rows affected: %w", err)
	}
	if rows == 0 {
		return fmt.Errorf(errNotFoundFmt, id)
	}

	return nil
}

// List returns memory records matching the filter with pagination.
func (s *postgresStore) List(ctx context.Context, filter Filter) ([]Record, int, error) {
	// Count total.
	countQB := applyFilter(psq.Select("COUNT(*)").From(tableName), filter)
	countQuery, countArgs, err := countQB.ToSql()
	if err != nil {
		return nil, 0, fmt.Errorf("building count query: %w", err)
	}

	var total int
	if err := s.db.QueryRowContext(ctx, countQuery, countArgs...).Scan(&total); err != nil {
		return nil, 0, fmt.Errorf("counting memory records: %w", err)
	}

	// Fetch paginated results.
	selectQB := applyPagination(
		applyFilter(psq.Select(recordColumns()...).From(tableName), filter),
		filter,
	)

	selectQuery, selectArgs, err := selectQB.ToSql()
	if err != nil {
		return nil, 0, fmt.Errorf("building select query: %w", err)
	}

	rows, err := s.db.QueryContext(ctx, selectQuery, selectArgs...)
	if err != nil {
		return nil, 0, fmt.Errorf("querying memory records: %w", err)
	}
	defer rows.Close() //nolint:errcheck // best-effort cleanup

	records, err := collectRecordRows(rows)
	if err != nil {
		return nil, 0, err
	}

	return records, total, nil
}

// applyPagination adds ORDER BY, LIMIT, and OFFSET clauses to a query builder.
func applyPagination(qb sq.SelectBuilder, filter Filter) sq.SelectBuilder {
	orderClause := "created_at DESC"
	if filter.OrderBy != "" {
		orderClause = filter.OrderBy
	}
	qb = qb.OrderBy(orderClause)

	if limit := filter.EffectiveLimit(); limit > 0 {
		qb = qb.Limit(uint64(limit))
	}
	if filter.Offset > 0 {
		qb = qb.Offset(uint64(filter.Offset))
	}
	return qb
}

// collectRecordRows scans all rows into a slice of Record.
func collectRecordRows(rows *sql.Rows) ([]Record, error) {
	var records []Record
	for rows.Next() {
		record, err := scanRecordRow(rows)
		if err != nil {
			return nil, err
		}
		records = append(records, *record)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterating memory rows: %w", err)
	}
	return records, nil
}

// VectorSearch performs cosine similarity search over embeddings.
func (s *postgresStore) VectorSearch(ctx context.Context, query VectorQuery) ([]ScoredRecord, error) {
	limit := query.Limit
	if limit <= 0 {
		limit = DefaultLimit
	}
	if limit > MaxLimit {
		limit = MaxLimit
	}

	// Build SQL manually to avoid squirrel placeholder collision with the
	// vector parameter ($1) used in the ORDER BY and SELECT expressions.
	args := []any{pgvector.NewVector(query.Embedding)}
	paramIdx := 2

	where := "WHERE embedding IS NOT NULL AND status <> 'archived'"
	if query.Persona != "" {
		where += fmt.Sprintf(" AND persona = $%d", paramIdx)
		args = append(args, query.Persona)
		paramIdx++
	}
	if query.Status != "" {
		where += fmt.Sprintf(" AND status = $%d", paramIdx)
		args = append(args, query.Status)
	}

	cols := "id, created_at, updated_at, created_by, persona, dimension, " +
		"content, category, confidence, source, " +
		"entity_urns, related_columns, metadata, " +
		"status, stale_reason, stale_at, last_verified, " +
		"1 - (embedding <=> $1) AS score"

	sqlStr := fmt.Sprintf( // #nosec G201 -- tableName is a constant, cols are hardcoded, where uses parameterized placeholders, limit is a sanitized int
		"SELECT %s FROM %s %s ORDER BY embedding <=> $1 LIMIT %d",
		cols, tableName, where, limit,
	)

	rows, err := s.db.QueryContext(ctx, sqlStr, args...)
	if err != nil {
		return nil, fmt.Errorf("executing vector search: %w", err)
	}
	defer rows.Close() //nolint:errcheck // best-effort cleanup

	return collectScoredRows(rows, query.MinScore)
}

// collectScoredRows scans all scored rows and filters by minimum score.
func collectScoredRows(rows *sql.Rows, minScore float64) ([]ScoredRecord, error) {
	var results []ScoredRecord
	for rows.Next() {
		record, score, err := scanScoredRow(rows)
		if err != nil {
			return nil, err
		}
		if minScore > 0 && score < minScore {
			continue
		}
		results = append(results, ScoredRecord{Record: *record, Score: score})
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterating vector search rows: %w", err)
	}

	return results, nil
}

// EntityLookup returns active memories linked to a DataHub URN.
func (s *postgresStore) EntityLookup(ctx context.Context, urn, persona string) ([]Record, error) {
	urnJSON, err := json.Marshal([]string{urn})
	if err != nil {
		return nil, fmt.Errorf("marshaling entity URN filter: %w", err)
	}

	qb := psq.Select(recordColumns()...).
		From(tableName).
		Where(sq.Expr("entity_urns @> ?::jsonb", urnJSON)).
		Where(sq.Eq{colStatus: StatusActive}).
		OrderBy("created_at DESC").
		Limit(uint64(DefaultLimit))

	if persona != "" {
		qb = qb.Where(sq.Eq{colPersona: persona})
	}

	query, args, err := qb.ToSql()
	if err != nil {
		return nil, fmt.Errorf("building entity lookup query: %w", err)
	}

	rows, err := s.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("querying entity memories: %w", err)
	}
	defer rows.Close() //nolint:errcheck // best-effort cleanup

	var records []Record
	for rows.Next() {
		record, err := scanRecordRow(rows)
		if err != nil {
			return nil, err
		}
		records = append(records, *record)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterating entity lookup rows: %w", err)
	}

	return records, nil
}

// MarkStale flags memory records as stale with a reason.
func (s *postgresStore) MarkStale(ctx context.Context, ids []string, reason string) error {
	if len(ids) == 0 {
		return nil
	}

	now := time.Now()
	query, args, err := psq.Update(tableName).
		Set(colStatus, StatusStale).
		Set("stale_reason", reason).
		Set("stale_at", now).
		Set("updated_at", now).
		Where(sq.Eq{"id": ids}).
		ToSql()
	if err != nil {
		return fmt.Errorf("building mark stale query: %w", err)
	}

	_, err = s.db.ExecContext(ctx, query, args...)
	if err != nil {
		return fmt.Errorf("marking records as stale: %w", err)
	}

	return nil
}

// MarkVerified updates the last_verified timestamp for records.
func (s *postgresStore) MarkVerified(ctx context.Context, ids []string) error {
	if len(ids) == 0 {
		return nil
	}

	now := time.Now()
	query, args, err := psq.Update(tableName).
		Set("last_verified", now).
		Set("updated_at", now).
		Where(sq.Eq{"id": ids}).
		ToSql()
	if err != nil {
		return fmt.Errorf("building mark verified query: %w", err)
	}

	_, err = s.db.ExecContext(ctx, query, args...)
	if err != nil {
		return fmt.Errorf("marking records as verified: %w", err)
	}

	return nil
}

// Supersede marks an old record as superseded by a new one.
// Uses an atomic UPDATE with jsonb || merge to avoid read-modify-write races.
func (s *postgresStore) Supersede(ctx context.Context, oldID, newID string) error {
	patch, err := json.Marshal(map[string]any{"superseded_by": newID})
	if err != nil {
		return fmt.Errorf("marshaling supersede metadata: %w", err)
	}

	now := time.Now()
	query, args, err := psq.Update(tableName).
		Set(colStatus, StatusSuperseded).
		Set(colMetadata, sq.Expr("metadata || ?::jsonb", patch)).
		Set("updated_at", now).
		Where(sq.Eq{"id": oldID}).
		ToSql()
	if err != nil {
		return fmt.Errorf("building supersede query: %w", err)
	}

	result, err := s.db.ExecContext(ctx, query, args...)
	if err != nil {
		return fmt.Errorf("superseding record: %w", err)
	}

	rows, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("checking rows affected: %w", err)
	}
	if rows == 0 {
		return fmt.Errorf(errNotFoundFmt, oldID)
	}

	return nil
}

// applyFilter adds filter conditions to a SELECT builder.
func applyFilter(qb sq.SelectBuilder, filter Filter) sq.SelectBuilder {
	if filter.CreatedBy != "" {
		qb = qb.Where(sq.Eq{colCreatedBy: filter.CreatedBy})
	}
	if filter.Persona != "" {
		qb = qb.Where(sq.Eq{colPersona: filter.Persona})
	}
	if filter.Dimension != "" {
		qb = qb.Where(sq.Eq{colDimension: filter.Dimension})
	}
	if filter.Category != "" {
		qb = qb.Where(sq.Eq{colCategory: filter.Category})
	}
	if filter.Status != "" {
		qb = qb.Where(sq.Eq{colStatus: filter.Status})
	}
	if filter.Source != "" {
		qb = qb.Where(sq.Eq{colSource: filter.Source})
	}
	if filter.EntityURN != "" {
		urnJSON, _ := json.Marshal([]string{filter.EntityURN})
		qb = qb.Where(sq.Expr("entity_urns @> ?::jsonb", urnJSON))
	}
	if filter.Since != nil {
		qb = qb.Where(sq.GtOrEq{colCreatedAt: *filter.Since})
	}
	if filter.Until != nil {
		qb = qb.Where(sq.LtOrEq{colCreatedAt: *filter.Until})
	}
	return qb
}

// recordColumns returns the column list for memory record queries.
func recordColumns() []string {
	return []string{
		"id", colCreatedAt, "updated_at", colCreatedBy, colPersona, colDimension,
		colContent, colCategory, colConfidence, colSource,
		colEntityURNs, colRelatedColumns, colMetadata,
		colStatus, "stale_reason", "stale_at", "last_verified",
	}
}

// scanRecord scans a single row from QueryRow into a Record.
func scanRecord(row *sql.Row) (*Record, error) {
	var r Record
	var entityURNs, relatedCols, metadata []byte
	var staleReason sql.NullString
	var staleAt, lastVerified sql.NullTime

	err := row.Scan(
		&r.ID, &r.CreatedAt, &r.UpdatedAt, &r.CreatedBy, &r.Persona, &r.Dimension,
		&r.Content, &r.Category, &r.Confidence, &r.Source,
		&entityURNs, &relatedCols, &metadata,
		&r.Status, &staleReason, &staleAt, &lastVerified,
	)
	if err != nil {
		return nil, fmt.Errorf("scanning memory record: %w", err)
	}

	if err := unmarshalRecordJSON(&r, entityURNs, relatedCols, metadata); err != nil {
		return nil, err
	}

	applyNullables(&r, staleReason, staleAt, lastVerified)
	return &r, nil
}

// scanRecordRow scans a single row from Rows into a Record.
func scanRecordRow(rows *sql.Rows) (*Record, error) {
	var r Record
	var entityURNs, relatedCols, metadata []byte
	var staleReason sql.NullString
	var staleAt, lastVerified sql.NullTime

	err := rows.Scan(
		&r.ID, &r.CreatedAt, &r.UpdatedAt, &r.CreatedBy, &r.Persona, &r.Dimension,
		&r.Content, &r.Category, &r.Confidence, &r.Source,
		&entityURNs, &relatedCols, &metadata,
		&r.Status, &staleReason, &staleAt, &lastVerified,
	)
	if err != nil {
		return nil, fmt.Errorf("scanning memory row: %w", err)
	}

	if err := unmarshalRecordJSON(&r, entityURNs, relatedCols, metadata); err != nil {
		return nil, err
	}

	applyNullables(&r, staleReason, staleAt, lastVerified)
	return &r, nil
}

// scanScoredRow scans a row with an appended score column.
func scanScoredRow(rows *sql.Rows) (*Record, float64, error) {
	var r Record
	var entityURNs, relatedCols, metadata []byte
	var staleReason sql.NullString
	var staleAt, lastVerified sql.NullTime
	var score float64

	err := rows.Scan(
		&r.ID, &r.CreatedAt, &r.UpdatedAt, &r.CreatedBy, &r.Persona, &r.Dimension,
		&r.Content, &r.Category, &r.Confidence, &r.Source,
		&entityURNs, &relatedCols, &metadata,
		&r.Status, &staleReason, &staleAt, &lastVerified,
		&score,
	)
	if err != nil {
		return nil, 0, fmt.Errorf("scanning scored row: %w", err)
	}

	if err := unmarshalRecordJSON(&r, entityURNs, relatedCols, metadata); err != nil {
		return nil, 0, err
	}

	applyNullables(&r, staleReason, staleAt, lastVerified)
	return &r, score, nil
}

// unmarshalRecordJSON unmarshals JSON columns into Record fields.
func unmarshalRecordJSON(r *Record, entityURNs, relatedCols, metadata []byte) error {
	if err := json.Unmarshal(entityURNs, &r.EntityURNs); err != nil {
		return fmt.Errorf("unmarshaling entity_urns: %w", err)
	}
	if err := json.Unmarshal(relatedCols, &r.RelatedColumns); err != nil {
		return fmt.Errorf("unmarshaling related_columns: %w", err)
	}
	if err := json.Unmarshal(metadata, &r.Metadata); err != nil {
		return fmt.Errorf("unmarshaling metadata: %w", err)
	}
	return nil
}

// applyNullables applies nullable SQL values to a Record.
func applyNullables(r *Record, staleReason sql.NullString, staleAt, lastVerified sql.NullTime) {
	if staleReason.Valid {
		r.StaleReason = staleReason.String
	}
	if staleAt.Valid {
		r.StaleAt = &staleAt.Time
	}
	if lastVerified.Valid {
		r.LastVerified = &lastVerified.Time
	}
}

// Verify interface compliance.
var _ Store = (*postgresStore)(nil)
