package vectordb

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"code.tczkiot.com/wlw/ai-agent/internal/pkg/config"

	turso "turso.tech/database/tursogo"
)

const (
	defaultLibSQLPath = "data/agent/vectors.db"
	defaultSearchTopK = 10
	busyTimeoutMillis = 5000
)

var collectionNamePattern = regexp.MustCompile(`^[A-Za-z][A-Za-z0-9_]*$`)

type LibSQLProvider struct {
	db *sql.DB
}

func NewLibSQLProvider(cfg *config.VectorDBConfig) (*LibSQLProvider, error) {
	if cfg == nil {
		return nil, fmt.Errorf("libsql vector database config is required")
	}
	path := strings.TrimSpace(cfg.Path)
	if path == "" {
		path = defaultLibSQLPath
	}
	absPath, err := filepath.Abs(path)
	if err != nil {
		return nil, fmt.Errorf("resolve libsql vector database path: %w", err)
	}
	if err := os.MkdirAll(filepath.Dir(absPath), 0o755); err != nil {
		return nil, fmt.Errorf("create libsql vector database directory: %w", err)
	}

	connector, err := turso.NewConnector(absPath, turso.WithBusyTimeout(busyTimeoutMillis))
	if err != nil {
		return nil, fmt.Errorf("create libsql vector database connector: %w", err)
	}
	db := sql.OpenDB(connector)
	db.SetMaxOpenConns(1)
	db.SetMaxIdleConns(1)
	provider := &LibSQLProvider{db: db}
	if err := provider.initialize(context.Background()); err != nil {
		_ = db.Close()
		return nil, err
	}
	return provider, nil
}

func (p *LibSQLProvider) initialize(ctx context.Context) error {
	if p == nil || p.db == nil {
		return fmt.Errorf("libsql vector database is closed")
	}
	if err := p.db.PingContext(ctx); err != nil {
		return fmt.Errorf("connect to libsql vector database: %w", err)
	}
	_, err := p.db.ExecContext(ctx, `CREATE TABLE IF NOT EXISTS "_agent_vector_collections" (
		name TEXT PRIMARY KEY NOT NULL,
		dimension INTEGER NOT NULL,
		created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
	)`)
	if err != nil {
		return fmt.Errorf("initialize libsql collection registry: %w", err)
	}
	return nil
}

func (p *LibSQLProvider) Close() error {
	if p == nil || p.db == nil {
		return nil
	}
	err := p.db.Close()
	p.db = nil
	return err
}

func (p *LibSQLProvider) CreateCollection(ctx context.Context, name string, dimension int) error {
	tableName, err := collectionIdentifier(name)
	if err != nil {
		return err
	}
	if dimension <= 0 || dimension > 65536 {
		return fmt.Errorf("invalid libsql vector dimension: %d", dimension)
	}
	if info, getErr := p.GetCollection(ctx, name); getErr == nil {
		if info.Dimension != dimension {
			return fmt.Errorf("collection %s already uses dimension %d, requested %d", name, info.Dimension, dimension)
		}
		return nil
	} else if !errors.Is(getErr, sql.ErrNoRows) {
		return getErr
	}

	tx, err := p.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin libsql collection transaction: %w", err)
	}
	defer func() { _ = tx.Rollback() }()

	createTable := fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s (
		id TEXT PRIMARY KEY NOT NULL,
		embedding BLOB NOT NULL,
		knowledge_base_id INTEGER NOT NULL DEFAULT 0,
		document_id INTEGER NOT NULL DEFAULT 0,
		document_title TEXT NOT NULL DEFAULT '',
		faq_id INTEGER NOT NULL DEFAULT 0,
		faq_question TEXT NOT NULL DEFAULT '',
		chunk_no INTEGER NOT NULL DEFAULT 0,
		chunk_type TEXT NOT NULL DEFAULT '',
		section_path TEXT NOT NULL DEFAULT '',
		title TEXT NOT NULL DEFAULT '',
		content TEXT NOT NULL DEFAULT '',
		provider TEXT NOT NULL DEFAULT ''
	)`, tableName)
	if _, err := tx.ExecContext(ctx, createTable); err != nil {
		return fmt.Errorf("create libsql collection %s: %w", name, err)
	}
	if _, err := tx.ExecContext(ctx, fmt.Sprintf(
		`CREATE INDEX IF NOT EXISTS %s ON %s (knowledge_base_id, document_id)`,
		quoteIdentifier(name+"_payload_idx"), tableName,
	)); err != nil {
		return fmt.Errorf("create libsql payload index for %s: %w", name, err)
	}
	if _, err := tx.ExecContext(ctx,
		`INSERT INTO "_agent_vector_collections" (name, dimension) VALUES (?, ?)`, name, dimension,
	); err != nil {
		return fmt.Errorf("register libsql collection %s: %w", name, err)
	}
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("commit libsql collection %s: %w", name, err)
	}
	return nil
}

func (p *LibSQLProvider) DeleteCollection(ctx context.Context, name string) error {
	tableName, err := collectionIdentifier(name)
	if err != nil {
		return err
	}
	tx, err := p.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin libsql collection transaction: %w", err)
	}
	defer func() { _ = tx.Rollback() }()
	if _, err := tx.ExecContext(ctx, "DROP TABLE IF EXISTS "+tableName); err != nil {
		return fmt.Errorf("drop libsql collection %s: %w", name, err)
	}
	if _, err := tx.ExecContext(ctx, `DELETE FROM "_agent_vector_collections" WHERE name = ?`, name); err != nil {
		return fmt.Errorf("unregister libsql collection %s: %w", name, err)
	}
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("commit libsql collection deletion %s: %w", name, err)
	}
	return nil
}

func (p *LibSQLProvider) GetCollection(ctx context.Context, name string) (*CollectionInfo, error) {
	tableName, err := collectionIdentifier(name)
	if err != nil {
		return nil, err
	}
	var dimension int
	if err := p.db.QueryRowContext(ctx,
		`SELECT dimension FROM "_agent_vector_collections" WHERE name = ?`, name,
	).Scan(&dimension); err != nil {
		return nil, err
	}
	var count int
	if err := p.db.QueryRowContext(ctx, "SELECT COUNT(*) FROM "+tableName).Scan(&count); err != nil {
		return nil, fmt.Errorf("count libsql collection %s: %w", name, err)
	}
	return &CollectionInfo{Name: name, Dimension: dimension, PointCount: count, Status: "ready"}, nil
}

func (p *LibSQLProvider) ListCollections(ctx context.Context) ([]string, error) {
	rows, err := p.db.QueryContext(ctx, `SELECT name FROM "_agent_vector_collections" ORDER BY name`)
	if err != nil {
		return nil, fmt.Errorf("list libsql collections: %w", err)
	}
	defer rows.Close()
	collections := make([]string, 0)
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			return nil, fmt.Errorf("scan libsql collection: %w", err)
		}
		collections = append(collections, name)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate libsql collections: %w", err)
	}
	return collections, nil
}

func (p *LibSQLProvider) UpsertVectors(ctx context.Context, collectionName string, vectors []Vector) error {
	if len(vectors) == 0 {
		return nil
	}
	tableName, err := collectionIdentifier(collectionName)
	if err != nil {
		return err
	}
	info, err := p.GetCollection(ctx, collectionName)
	if err != nil {
		return fmt.Errorf("get libsql collection %s: %w", collectionName, err)
	}

	tx, err := p.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin libsql vector upsert: %w", err)
	}
	defer func() { _ = tx.Rollback() }()
	statement := fmt.Sprintf(`INSERT INTO %s (
		id, embedding, knowledge_base_id, document_id, document_title,
		faq_id, faq_question, chunk_no, chunk_type, section_path, title, content, provider
	) VALUES (?, vector32(?), ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
	ON CONFLICT(id) DO UPDATE SET
		embedding=excluded.embedding,
		knowledge_base_id=excluded.knowledge_base_id,
		document_id=excluded.document_id,
		document_title=excluded.document_title,
		faq_id=excluded.faq_id,
		faq_question=excluded.faq_question,
		chunk_no=excluded.chunk_no,
		chunk_type=excluded.chunk_type,
		section_path=excluded.section_path,
		title=excluded.title,
		content=excluded.content,
		provider=excluded.provider`, tableName)
	stmt, err := tx.PrepareContext(ctx, statement)
	if err != nil {
		return fmt.Errorf("prepare libsql vector upsert: %w", err)
	}
	defer stmt.Close()
	for _, item := range vectors {
		if strings.TrimSpace(item.ID) == "" {
			return fmt.Errorf("libsql vector id is required")
		}
		if len(item.Vector) != info.Dimension {
			return fmt.Errorf("invalid vector dimension for %s: got %d, want %d", item.ID, len(item.Vector), info.Dimension)
		}
		encoded, err := json.Marshal(item.Vector)
		if err != nil {
			return fmt.Errorf("encode vector %s: %w", item.ID, err)
		}
		payload := item.Payload
		if _, err := stmt.ExecContext(ctx,
			item.ID, string(encoded), payload.KnowledgeBaseID, payload.DocumentID, payload.DocumentTitle,
			payload.FaqID, payload.FaqQuestion, payload.ChunkNo, payload.ChunkType,
			payload.SectionPath, payload.Title, payload.Content, payload.Provider,
		); err != nil {
			return fmt.Errorf("upsert libsql vector %s: %w", item.ID, err)
		}
	}
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("commit libsql vector upsert: %w", err)
	}
	return nil
}

func (p *LibSQLProvider) DeleteVectors(ctx context.Context, collectionName string, ids []string) error {
	if len(ids) == 0 {
		return nil
	}
	tableName, err := collectionIdentifier(collectionName)
	if err != nil {
		return err
	}
	tx, err := p.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin libsql vector deletion: %w", err)
	}
	defer func() { _ = tx.Rollback() }()
	stmt, err := tx.PrepareContext(ctx, "DELETE FROM "+tableName+" WHERE id = ?")
	if err != nil {
		return fmt.Errorf("prepare libsql vector deletion: %w", err)
	}
	defer stmt.Close()
	for _, id := range ids {
		if _, err := stmt.ExecContext(ctx, id); err != nil {
			return fmt.Errorf("delete libsql vector %s: %w", id, err)
		}
	}
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("commit libsql vector deletion: %w", err)
	}
	return nil
}

func (p *LibSQLProvider) Search(ctx context.Context, req *SearchRequest) ([]SearchResult, error) {
	if req == nil {
		return nil, fmt.Errorf("libsql search request is required")
	}
	tableName, err := collectionIdentifier(req.CollectionName)
	if err != nil {
		return nil, err
	}
	info, err := p.GetCollection(ctx, req.CollectionName)
	if err != nil {
		return nil, fmt.Errorf("get libsql collection %s: %w", req.CollectionName, err)
	}
	if len(req.Vector) != info.Dimension {
		return nil, fmt.Errorf("invalid search vector dimension: got %d, want %d", len(req.Vector), info.Dimension)
	}
	topK := req.TopK
	if topK <= 0 {
		topK = defaultSearchTopK
	}
	encoded, err := json.Marshal(req.Vector)
	if err != nil {
		return nil, fmt.Errorf("encode search vector: %w", err)
	}
	vectorJSON := string(encoded)

	filterSQL, filterArgs := buildSearchFilter(req.Filter)
	innerWhere := filterSQL
	innerArgs := []any{vectorJSON}
	if filterSQL != "" {
		innerArgs = append(innerArgs, filterArgs...)
	}

	query := fmt.Sprintf(`SELECT id, score, knowledge_base_id, document_id, document_title,
		faq_id, faq_question, chunk_no, chunk_type, section_path, title, content, provider
	FROM (
		SELECT id, 1.0 - vector_distance_cos(embedding, vector32(?)) AS score,
			knowledge_base_id, document_id, document_title, faq_id, faq_question,
			chunk_no, chunk_type, section_path, title, content, provider
		FROM %s%s
	) ranked
	WHERE score >= ?
	ORDER BY score DESC
	LIMIT ?`, tableName, innerWhere)
	innerArgs = append(innerArgs, req.ScoreThreshold, topK)
	rows, err := p.db.QueryContext(ctx, query, innerArgs...)
	if err != nil {
		return nil, fmt.Errorf("search libsql collection %s: %w", req.CollectionName, err)
	}
	defer rows.Close()

	results := make([]SearchResult, 0, topK)
	for rows.Next() {
		var result SearchResult
		if err := rows.Scan(
			&result.ID, &result.Score,
			&result.Payload.KnowledgeBaseID, &result.Payload.DocumentID, &result.Payload.DocumentTitle,
			&result.Payload.FaqID, &result.Payload.FaqQuestion, &result.Payload.ChunkNo,
			&result.Payload.ChunkType, &result.Payload.SectionPath, &result.Payload.Title,
			&result.Payload.Content, &result.Payload.Provider,
		); err != nil {
			return nil, fmt.Errorf("scan libsql search result: %w", err)
		}
		results = append(results, result)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate libsql search results: %w", err)
	}
	return results, nil
}

func collectionIdentifier(name string) (string, error) {
	name = strings.TrimSpace(name)
	if !collectionNamePattern.MatchString(name) {
		return "", fmt.Errorf("invalid libsql collection name %q", name)
	}
	return quoteIdentifier(name), nil
}

func quoteIdentifier(value string) string {
	return `"` + value + `"`
}

func buildSearchFilter(filter *SearchFilter) (string, []any) {
	if filter == nil {
		return "", nil
	}
	clauses := make([]string, 0, 2)
	args := make([]any, 0, len(filter.KnowledgeBaseIDs)+len(filter.DocumentIDs))
	if len(filter.KnowledgeBaseIDs) > 0 {
		clauses = append(clauses, "knowledge_base_id IN ("+placeholders(len(filter.KnowledgeBaseIDs))+")")
		for _, id := range filter.KnowledgeBaseIDs {
			args = append(args, id)
		}
	}
	if len(filter.DocumentIDs) > 0 {
		clauses = append(clauses, "document_id IN ("+placeholders(len(filter.DocumentIDs))+")")
		for _, id := range filter.DocumentIDs {
			args = append(args, id)
		}
	}
	if len(clauses) == 0 {
		return "", nil
	}
	return " WHERE " + strings.Join(clauses, " AND "), args
}

func placeholders(count int) string {
	values := make([]string, count)
	for i := range values {
		values[i] = "?"
	}
	return strings.Join(values, ",")
}