package vectors

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

	vec "github.com/asg017/sqlite-vec-go-bindings/cgo"
	_ "github.com/mattn/go-sqlite3" // SQLite driver
)

// Error types for better error handling
var (
	ErrDimensionMismatch    = errors.New("dimension mismatch")
	ErrInvalidVector        = errors.New("invalid vector")
	ErrExtensionUnavailable = errors.New("sqlite-vec extension unavailable")
)

// DimensionMismatchError provides detailed dimension mismatch information
type DimensionMismatchError struct {
	Expected int
	Actual   int
}

func (e *DimensionMismatchError) Error() string {
	return fmt.Sprintf("dimension mismatch: vector has %d dimensions, table configured for %d",
		e.Actual, e.Expected)
}

func (e *DimensionMismatchError) Is(target error) bool {
	return target == ErrDimensionMismatch
}

// Logger is a minimal interface to allow structured logging without
// importing a concrete logging package here. The real application logger
// should satisfy this.
type Logger interface {
	Info(msg string, kv ...interface{})
	Warn(msg string, kv ...interface{})
	Error(msg string, kv ...interface{})
}

// noopLogger is used when no logger is provided.
type noopLogger struct{}

func (n *noopLogger) Info(string, ...interface{})  {}
func (n *noopLogger) Warn(string, ...interface{})  {}
func (n *noopLogger) Error(string, ...interface{}) {}

// SQLiteVecConfig carries construction parameters.
type SQLiteVecConfig struct {
	Path      string
	Dimension int // requested / primary dimension (API-level)
}

// SQLiteVecStorage provides sqlite-vec backed vector storage
type SQLiteVecStorage struct {
	db     *sql.DB
	cfg    SQLiteVecConfig
	logger Logger
}

// NewSQLiteVecStorage creates a new sqlite-vec storage instance
func NewSQLiteVecStorage(cfg SQLiteVecConfig, logger Logger) (*SQLiteVecStorage, error) {
	if cfg.Path == "" {
		return nil, errors.New("sqlite vec storage path required")
	}
	if cfg.Dimension == 0 {
		cfg.Dimension = DefaultEmbeddingDim
	}
	if cfg.Dimension < MinEmbeddingDim || cfg.Dimension > MaxEmbeddingDim {
		return nil, fmt.Errorf("configured dimension %d invalid: expected between %d and %d", cfg.Dimension, MinEmbeddingDim, MaxEmbeddingDim)
	}
	if logger == nil {
		logger = &noopLogger{}
	}

	db, err := sql.Open("sqlite3", cfg.Path)
	if err != nil {
		return nil, fmt.Errorf("failed to open database: %w", err)
	}

	storage := &SQLiteVecStorage{db: db, cfg: cfg, logger: logger}

	if err := storage.initSchema(); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("failed to initialize schema: %w", err)
	}

	return storage, nil
}

// initSchema creates the necessary tables and loads sqlite-vec extension
func (s *SQLiteVecStorage) initSchema() error {
	// Register the sqlite-vec extension using Go bindings
	// vec.Auto() registers the extension globally but doesn't return an error
	// (it's a wrapper around sqlite3_auto_extension which is best-effort)
	vec.Auto()

	// Configure SQLite for high performance
	if err := s.configureSQLitePerformance(); err != nil {
		s.logger.Warn("failed to configure SQLite performance settings", "error", err)
		// Continue anyway - performance settings are not critical for functionality
	}

	// Test if sqlite-vec extension loaded successfully by trying to create a virtual table
	testTable := fmt.Sprintf(`
	CREATE VIRTUAL TABLE IF NOT EXISTS vec_test USING vec0(
		embedding FLOAT[%d]
	)`, s.cfg.Dimension)

	if _, err := s.db.Exec(testTable); err != nil {
		return fmt.Errorf("sqlite-vec extension not available: %w", err)
	}

	// Clean up test table
	_, _ = s.db.Exec("DROP TABLE IF EXISTS vec_test")
	s.logger.Info("sqlite-vec extension available and working")

	// Create sqlite-vec virtual table
	createVecTable := fmt.Sprintf(`
	CREATE VIRTUAL TABLE IF NOT EXISTS embeddings USING vec0(
		id TEXT PRIMARY KEY,
		embedding FLOAT[%d],
		metadata TEXT
	)`, s.cfg.Dimension)

	if _, err := s.db.Exec(createVecTable); err != nil {
		return fmt.Errorf("failed to create sqlite-vec table: %w", err)
	}

	s.logger.Info("sqlite-vec virtual table created successfully", "dimensions", s.cfg.Dimension)
	return nil
}

// configureSQLitePerformance sets SQLite PRAGMA settings for optimal performance
func (s *SQLiteVecStorage) configureSQLitePerformance() error {
	pragmaSettings := []struct {
		setting string
		value   string
		desc    string
	}{
		// WAL mode for better concurrent read/write performance
		{"journal_mode", "WAL", "Write-Ahead Logging for concurrent access"},

		// Increase cache size (negative value = KB, positive = pages)
		{"cache_size", "-64000", "64MB cache for better read performance"},

		// Memory-mapped I/O for better performance
		{"mmap_size", "268435456", "256MB memory mapping"},

		// Optimize for fast writes
		{"synchronous", "NORMAL", "Balanced durability vs performance"},

		// Temp store in memory
		{"temp_store", "MEMORY", "Store temporary data in memory"},

		// Larger page size for better I/O efficiency
		{"page_size", "4096", "4KB page size for modern storage"},

		// Automatic index creation
		{"automatic_index", "ON", "Enable automatic indexes for queries"},
	}

	for _, pragma := range pragmaSettings {
		query := fmt.Sprintf("PRAGMA %s = %s", pragma.setting, pragma.value)
		if _, err := s.db.Exec(query); err != nil {
			s.logger.Warn("failed to set pragma",
				"setting", pragma.setting,
				"value", pragma.value,
				"desc", pragma.desc,
				"error", err)
			// Continue with other settings
		} else {
			s.logger.Info("configured SQLite pragma",
				"setting", pragma.setting,
				"value", pragma.value,
				"desc", pragma.desc)
		}
	}

	return nil
}

// Store saves a vector to the database
func (s *SQLiteVecStorage) Store(vector *Vector) error {
	// Validate the vector first
	vm := NewVectorMath()
	if err := vm.ValidateVector(vector); err != nil {
		return fmt.Errorf("invalid vector: %w", err)
	}

	// Check dimension matches the table configuration specifically
	if len(vector.Embedding) != s.cfg.Dimension {
		return &DimensionMismatchError{
			Expected: s.cfg.Dimension,
			Actual:   len(vector.Embedding),
		}
	}

	// Use sqlite-vec binary format
	embeddingData, err := serializeEmbeddingBinary(vector.Embedding)
	if err != nil {
		return fmt.Errorf("failed to serialize embedding: %w", err)
	}

	// Serialize metadata as JSON
	metadataJSON, err := json.Marshal(vector.Metadata)
	if err != nil {
		return fmt.Errorf("failed to serialize metadata: %w", err)
	}

	// Use sqlite-vec optimized storage
	query := `
	INSERT OR REPLACE INTO embeddings (id, embedding, metadata)
	VALUES (?, ?, ?)`

	_, err = s.db.Exec(query, vector.ID, embeddingData, string(metadataJSON))
	if err != nil {
		return fmt.Errorf("failed to store vector: %w", err)
	}

	s.logger.Info("vector stored successfully", "id", vector.ID, "dimensions", len(vector.Embedding))
	return nil
}

// StoreBatch saves multiple vectors in a single transaction for improved performance
func (s *SQLiteVecStorage) StoreBatch(vectors []*Vector) (int, error) {
	if len(vectors) == 0 {
		return 0, nil
	}

	tx, err := s.db.Begin()
	if err != nil {
		return 0, fmt.Errorf("failed to begin transaction: %w", err)
	}
	// Use a named error in a closure to handle rollback correctly.
	var txErr error
	defer func() {
		if txErr != nil {
			if rbErr := tx.Rollback(); rbErr != nil {
				s.logger.Warn("failed to rollback transaction", "original_error", txErr, "rollback_error", rbErr)
			}
		}
	}()

	stmt, err := tx.Prepare(`INSERT OR REPLACE INTO embeddings (id, embedding, metadata) VALUES (?, ?, ?)`)
	if err != nil {
		txErr = fmt.Errorf("failed to prepare batch insert statement: %w", err)
		return 0, txErr
	}
	defer func() {
		if closeErr := stmt.Close(); closeErr != nil && txErr == nil {
			txErr = fmt.Errorf("failed to close statement: %w", closeErr)
		}
	}()

	vm := NewVectorMath()
	successCount := 0

	for i, vector := range vectors {
		// Validate the vector
		if err := vm.ValidateVector(vector); err != nil {
			s.logger.Warn("skipping invalid vector in batch", "index", i, "id", vector.ID, "error", err)
			continue
		}

		// Check dimension matches
		if len(vector.Embedding) != s.cfg.Dimension {
			s.logger.Warn("skipping vector with dimension mismatch in batch",
				"index", i, "id", vector.ID, "expected", s.cfg.Dimension, "actual", len(vector.Embedding))
			continue
		}

		// Serialize embedding
		embeddingData, serErr := serializeEmbeddingBinary(vector.Embedding)
		if serErr != nil {
			s.logger.Warn("skipping vector with serialization error", "index", i, "id", vector.ID, "error", serErr)
			continue
		}

		// Serialize metadata
		metadataJSON, jsonErr := json.Marshal(vector.Metadata)
		if jsonErr != nil {
			s.logger.Warn("skipping vector with metadata error", "index", i, "id", vector.ID, "error", jsonErr)
			continue
		}

		// Execute prepared statement
		_, execErr := stmt.Exec(vector.ID, embeddingData, string(metadataJSON))
		if execErr != nil {
			s.logger.Warn("failed to store vector in batch", "index", i, "id", vector.ID, "error", execErr)
			continue
		}

		successCount++
	}

	// Commit transaction
	txErr = tx.Commit()
	if txErr != nil {
		return successCount, fmt.Errorf("failed to commit batch transaction: %w", txErr)
	}

	s.logger.Info("batch storage completed",
		"total_vectors", len(vectors),
		"successful_stores", successCount,
		"dimensions", s.cfg.Dimension)

	return successCount, nil
}

// Search performs vector similarity search using sqlite-vec
func (s *SQLiteVecStorage) Search(query []float64, topK int) ([]SearchResult, error) {
	// Validate query vector
	vm := NewVectorMath()
	if err := vm.ValidateDimensions(query); err != nil {
		return nil, fmt.Errorf("invalid query vector: %w", err)
	}

	if topK <= 0 {
		topK = 10 // Default
	}

	return s.vectorSearch(query, topK)
}

// vectorSearch uses sqlite-vec extension for optimized search
func (s *SQLiteVecStorage) vectorSearch(query []float64, topK int) ([]SearchResult, error) {
	// Serialize query vector for sqlite-vec using binary format
	queryData, err := serializeEmbeddingBinary(query)
	if err != nil {
		return nil, fmt.Errorf("failed to serialize query vector: %w", err)
	}

	// sqlite-vec query syntax - we don't need to retrieve the embedding data
	sqlQuery := `
	SELECT id, metadata,
		   vec_distance_cosine(embedding, ?) as distance
	FROM embeddings
	ORDER BY distance ASC
	LIMIT ?`

	rows, err := s.db.Query(sqlQuery, queryData, topK)
	if err != nil {
		return nil, fmt.Errorf("sqlite-vec query failed: %w", err)
	}
	defer func() { _ = rows.Close() }()

	var results []SearchResult

	for rows.Next() {
		var id, metadataJSON string
		var distance float64

		err := rows.Scan(&id, &metadataJSON, &distance)
		if err != nil {
			s.logger.Warn("failed to scan sqlite-vec result", "error", err)
			continue
		}

		// Parse metadata
		var metadata map[string]interface{}
		if err := json.Unmarshal([]byte(metadataJSON), &metadata); err != nil {
			s.logger.Warn("failed to parse metadata", "id", id, "error", err)
			metadata = map[string]interface{}{"raw": metadataJSON}
		}

		// Convert distance to similarity score
		// For cosine distance, values range from 0 (identical) to 2 (opposite)
		// We convert to similarity scale 0-1 where 1 is most similar
		similarity := distanceToSimilarity(distance)

		vector := &Vector{
			ID:        id,
			Embedding: nil, // Not fetched in sqlite-vec mode for efficiency
			Metadata:  metadata,
		}

		results = append(results, SearchResult{
			Vector: vector,
			Score:  similarity,
		})
	}

	s.logger.Info("sqlite-vec search completed", "query_dims", len(query), "results", len(results), "top_k", topK)
	return results, nil
}

// Delete removes all vectors associated with a specific document ID
func (s *SQLiteVecStorage) Delete(documentID string) error {
	if s.db == nil {
		return fmt.Errorf("database not initialized")
	}

	query := `DELETE FROM embeddings WHERE id = ?`
	result, err := s.db.Exec(query, documentID)
	if err != nil {
		return fmt.Errorf("failed to delete vectors for document %s: %w", documentID, err)
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		s.logger.Warn("could not get rows affected count", "error", err)
	} else {
		s.logger.Info("deleted vectors for document", "document_id", documentID, "rows_affected", rowsAffected)
	}

	return nil
}

// Close closes the database connection
func (s *SQLiteVecStorage) Close() error {
	if s.db != nil {
		return s.db.Close()
	}
	return nil
}

// Helper functions for embedding serialization

// serializeEmbeddingBinary converts float64 slice to binary format for sqlite-vec
func serializeEmbeddingBinary(embedding []float64) ([]byte, error) {
	// Convert float64 to float32 for sqlite-vec
	float32Vec := make([]float32, len(embedding))
	for i, v := range embedding {
		float32Vec[i] = float32(v)
	}

	// Use sqlite-vec-go-bindings serialization
	data, err := vec.SerializeFloat32(float32Vec)
	if err != nil {
		return nil, fmt.Errorf("failed to serialize embedding to binary: %w", err)
	}
	return data, nil
}

// distanceToSimilarity converts cosine distance to similarity score
// Cosine distance ranges from 0 (identical) to 2 (opposite)
// Similarity score ranges from 1 (most similar) to 0 (least similar)
func distanceToSimilarity(distance float64) float64 {
	if distance >= 2.0 {
		return 0.0 // Completely opposite vectors
	} else if distance <= 0.0 {
		return 1.0 // Identical vectors
	} else {
		// Linear mapping from distance [0,2] to similarity [1,0]
		return 1.0 - (distance / 2.0)
	}
}
