// Package vec 在 mattn/go-sqlite3 上集成 sqlite-vec 向量检索扩展。
//
// 使用该驱动后，可直接创建 vec0 虚拟表并执行向量相似度检索：
//
//	db, _ := sql.Open(vec.DriverName, "data.db")
//	db.Exec(`CREATE VIRTUAL TABLE IF NOT EXISTS embeddings USING vec0(vector float[1536])`)
//	db.Exec(`INSERT INTO embeddings(rowid, vector) VALUES (?, ?)`, id, vec.SerializeFloat32(embedding))
//	db.QueryRow(`SELECT rowid, distance FROM embeddings WHERE vector MATCH ? ORDER BY distance LIMIT 10`, vec.SerializeFloat32(query))
package vec

import (
	"database/sql"
	"encoding/binary"
	"encoding/json"
	"math"
	"regexp"

	"git.fsdpf.net/go/db"
	"git.fsdpf.net/go/db/dialect/sqlite3"
	sqlitevec "github.com/asg017/sqlite-vec-go-bindings/cgo"
	gosqlite3 "github.com/mattn/go-sqlite3"
)

// DriverName 是加载了 sqlite-vec 扩展的 SQLite3 驱动名称。
const DriverName = "sqlite3_vec"

func init() {
	sqlitevec.Auto() // 全局加载 sqlite-vec 扩展到所有后续连接

	sql.Register(DriverName, &gosqlite3.SQLiteDriver{
		ConnectHook: func(conn *gosqlite3.SQLiteConn) error {
			// sqlitevec.LoadIntoConn 在当前版本库中不存在，改用 init() 中的 Auto() 全局注册
			// if err := sqlitevec.LoadIntoConn(conn); err != nil {
			// 	return err
			// }

			if err := conn.RegisterFunc("IF", func(cond int64, trueVal, falseVal interface{}) interface{} {
				if cond != 0 {
					return trueVal
				}
				return falseVal
			}, true); err != nil {
				return err
			}

			if err := conn.RegisterFunc("REGEXP", func(expr, item string) (bool, error) {
				return regexp.MatchString(expr, item)
			}, true); err != nil {
				return err
			}

			if err := conn.RegisterFunc("vec_serialize", func(jsonText string) ([]byte, error) {
				var floats []float64
				if err := json.Unmarshal([]byte(jsonText), &floats); err != nil {
					return nil, err
				}
				f32 := make([]float32, len(floats))
				for i, f := range floats {
					f32[i] = float32(f)
				}
				return SerializeFloat32(f32), nil
			}, true); err != nil {
				return err
			}

			if err := conn.RegisterFunc("vec_deserialize", func(b []byte) (string, error) {
				floats := DeserializeFloat32(b)
				f64 := make([]float64, len(floats))
				for i, f := range floats {
					f64[i] = float64(f)
				}
				out, err := json.Marshal(f64)
				return string(out), err
			}, true); err != nil {
				return err
			}

			return nil
		},
	})

	db.RegisterDialect(DriverName, sqlite3.DialectOptions())
}

// SerializeFloat32 将 float32 切片序列化为 sqlite-vec 接受的小端 IEEE 754 字节序列。
func SerializeFloat32(v []float32) []byte {
	buf := make([]byte, len(v)*4)
	for i, f := range v {
		binary.LittleEndian.PutUint32(buf[i*4:], math.Float32bits(f))
	}
	return buf
}

// DeserializeFloat32 将 sqlite-vec 返回的字节序列反序列化为 float32 切片。
func DeserializeFloat32(b []byte) []float32 {
	v := make([]float32, len(b)/4)
	for i := range v {
		v[i] = math.Float32frombits(binary.LittleEndian.Uint32(b[i*4:]))
	}
	return v
}