graph.go

package main

import (
	"crypto/rand"
	"database/sql"
	"encoding/gob"
	"errors"
	"math/big"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"

	bbolt "go.etcd.io/bbolt"
	_ "modernc.org/sqlite"
)

type Entity struct {
	ID           string   `json:"id"`
	Name         string   `json:"name"`
	EntityType   string   `json:"entityType"`
	Observations []string `json:"observations"`
}

func (m *GraphManager) snapshotPath() string {
	return m.sqlitePath + ".tokidx.gob"
}

type indexSnapshot struct {
	Version int
	Tok     map[string][]string
}

func (m *GraphManager) tryLoadIndexSnapshot() bool {
	ver := m.currentDataVersion()
	f, err := os.Open(m.snapshotPath())
	if err != nil {
		return false
	}
	defer f.Close()
	dec := gob.NewDecoder(f)
	var snap indexSnapshot
	if err := dec.Decode(&snap); err != nil {
		return false
	}
	if snap.Version != ver {
		return false
	}
	tok := make(map[string]map[string]struct{}, len(snap.Tok))
	for k, arr := range snap.Tok {
		s := make(map[string]struct{}, len(arr))
		for _, n := range arr {
			s[n] = struct{}{}
		}
		tok[k] = s
	}
	byName := make(map[string]Entity, len(m.cacheGraph.Entities))
	for _, e := range m.cacheGraph.Entities {
		byName[e.Name] = e
	}
	m.tokIndex = tok
	m.entsByName = byName
	m.indexReady = true
	m.indexBuilding = false
	return true
}

func (m *GraphManager) saveIndexSnapshot(version int) error {
	if version < 0 {
		return nil
	}
	tokArr := make(map[string][]string, len(m.tokIndex))
	for k, set := range m.tokIndex {
		arr := make([]string, 0, len(set))
		for n := range set {
			arr = append(arr, n)
		}
		tokArr[k] = arr
	}
	snap := indexSnapshot{Version: version, Tok: tokArr}
	tmp := m.snapshotPath() + ".tmp"
	f, err := os.Create(tmp)
	if err != nil {
		return err
	}
	enc := gob.NewEncoder(f)
	if err := enc.Encode(&snap); err != nil {
		f.Close()
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	return os.Rename(tmp, m.snapshotPath())
}

func fnv32(s string) uint32 {
	var h uint32 = 2166136261
	for i := 0; i < len(s); i++ {
		h ^= uint32(s[i])
		h *= 16777619
	}
	return h
}

func (m *GraphManager) currentDataVersion() int {
	var ver int
	_ = m.db.QueryRow("SELECT version FROM graph_meta").Scan(&ver)
	return ver
}

type execer interface {
	Exec(query string, args ...any) (sql.Result, error)
}

func (m *GraphManager) ftsUpdate(exec execer, id, name, entityType string, observations []string) error {
	var b strings.Builder
	b.WriteString(strings.ToLower(name))
	b.WriteByte(' ')
	b.WriteString(strings.ToLower(entityType))
	for _, o := range observations {
		b.WriteByte(' ')
		b.WriteString(strings.ToLower(o))
	}
	_, err := exec.Exec("DELETE FROM fts WHERE entity_id= ?", id)
	if err != nil {
		return err
	}
	_, err = exec.Exec("INSERT INTO fts(entity_id, text) VALUES(?, ?)", id, b.String())
	return err
}

type Relation struct {
	From         string `json:"from"`
	To           string `json:"to"`
	RelationType string `json:"relationType"`
}

type KnowledgeGraph struct {
	Entities  []Entity   `json:"entities"`
	Relations []Relation `json:"relations"`
}

type GraphManager struct {
	sqlitePath    string
	boltPath      string
	db            *sql.DB
	bolt          *bbolt.DB
	mu            sync.Mutex
	cacheReady    bool
	cacheGraph    KnowledgeGraph
	tokIndex      map[string]map[string]struct{}
	entsByName    map[string]Entity
	indexReady    bool
	indexBuilding bool
}

func NewGraphManager(sqlitePath, boltPath string) *GraphManager {
	return &GraphManager{sqlitePath: sqlitePath, boltPath: boltPath}
}

func (m *GraphManager) open() error {
	if m.db != nil && m.bolt != nil {
		return nil
	}
	if dir := filepath.Dir(m.sqlitePath); dir != "." && dir != "" {
		if err := os.MkdirAll(dir, 0o755); err != nil {
			return err
		}
	}
	if dir := filepath.Dir(m.boltPath); dir != "." && dir != "" {
		if err := os.MkdirAll(dir, 0o755); err != nil {
			return err
		}
	}
	db, err := sql.Open("sqlite", m.sqlitePath)
	if err != nil {
		return err
	}
	db.SetMaxOpenConns(1)
	db.SetMaxIdleConns(1)
	if _, err := db.Exec("PRAGMA foreign_keys = ON"); err != nil {
		return err
	}
	_, _ = db.Exec("PRAGMA busy_timeout = 10000")
	_, _ = db.Exec("PRAGMA journal_mode = WAL")
	_, _ = db.Exec("PRAGMA synchronous = NORMAL")
	_, _ = db.Exec("PRAGMA temp_store = MEMORY")
	_, _ = db.Exec("PRAGMA cache_size = -200000")
	_, _ = db.Exec("PRAGMA mmap_size = 1073741824")
	if err := applyMigrations(db); err != nil {
		return err
	}
	bdb, err := bbolt.Open(m.boltPath, 0o600, nil)
	if err != nil {
		return err
	}
	if err := bdb.Update(func(tx *bbolt.Tx) error {
		if _, err := tx.CreateBucketIfNotExists([]byte("term")); err != nil {
			return err
		}
		return nil
	}); err != nil {
		return err
	}
	m.db = db
	m.bolt = bdb
	return nil
}

func applyMigrations(db *sql.DB) error {
	if _, err := db.Exec(`CREATE TABLE IF NOT EXISTS schema_version (version INTEGER NOT NULL);`); err != nil {
		return err
	}
	var cnt int
	if err := db.QueryRow(`SELECT COUNT(1) FROM schema_version`).Scan(&cnt); err != nil {
		return err
	}
	if cnt == 0 {
		if _, err := db.Exec(`INSERT INTO schema_version(version) VALUES (0)`); err != nil {
			return err
		}
	}
	var v int
	if err := db.QueryRow(`SELECT version FROM schema_version`).Scan(&v); err != nil {
		return err
	}
	if v < 1 {
		if _, err := db.Exec(`
            CREATE TABLE entities (
                id TEXT PRIMARY KEY,
                name TEXT NOT NULL UNIQUE,
                entity_type TEXT NOT NULL
            );
        `); err != nil {
			return err
		}
		if _, err := db.Exec(`
            CREATE TABLE observations (
                entity_id TEXT NOT NULL,
                content TEXT NOT NULL,
                UNIQUE(entity_id, content),
                FOREIGN KEY(entity_id) REFERENCES entities(id) ON DELETE CASCADE
            );
        `); err != nil {
			return err
		}
		if _, err := db.Exec(`
            CREATE TABLE relations (
                from_id TEXT NOT NULL,
                to_id TEXT NOT NULL,
                relation_type TEXT NOT NULL,
                UNIQUE(from_id, to_id, relation_type),
                FOREIGN KEY(from_id) REFERENCES entities(id) ON DELETE CASCADE,
                FOREIGN KEY(to_id) REFERENCES entities(id) ON DELETE CASCADE
            );
        `); err != nil {
			return err
		}
		if _, err := db.Exec(`UPDATE schema_version SET version = 1`); err != nil {
			return err
		}
		v = 1
	}
	if v < 2 {
		if _, err := db.Exec(`CREATE INDEX IF NOT EXISTS idx_observations_entity ON observations(entity_id)`); err != nil {
			return err
		}
		if _, err := db.Exec(`CREATE INDEX IF NOT EXISTS idx_relations_from ON relations(from_id)`); err != nil {
			return err
		}
		if _, err := db.Exec(`CREATE INDEX IF NOT EXISTS idx_relations_to ON relations(to_id)`); err != nil {
			return err
		}
		if _, err := db.Exec(`CREATE INDEX IF NOT EXISTS idx_relations_from_to ON relations(from_id, to_id, relation_type)`); err != nil {
			return err
		}
		if _, err := db.Exec(`UPDATE schema_version SET version = 2`); err != nil {
			return err
		}
		v = 2
	}
	if v < 3 {
		if _, err := db.Exec(`CREATE VIRTUAL TABLE IF NOT EXISTS fts USING fts5(entity_id, text, tokenize='unicode61')`); err != nil {
			return err
		}
		if _, err := db.Exec(`INSERT INTO fts(entity_id, text)
            SELECT e.id,
                   lower(e.name || ' ' || e.entity_type ||
                         CASE WHEN group_concat(o.content, ' ') IS NULL THEN '' ELSE ' '||group_concat(o.content, ' ') END)
            FROM entities e
            LEFT JOIN observations o ON o.entity_id=e.id
            GROUP BY e.id`); err != nil {
			return err
		}
		if _, err := db.Exec(`UPDATE schema_version SET version = 3`); err != nil {
			return err
		}
		v = 3
	}
	if v < 4 {
		if _, err := db.Exec(`CREATE TABLE IF NOT EXISTS entity_doc (entity_id TEXT PRIMARY KEY, obs TEXT NOT NULL)`); err != nil {
			return err
		}
		if _, err := db.Exec(`INSERT INTO entity_doc(entity_id, obs)
            SELECT e.id, ifnull(group_concat(o.content, '\n'), '')
            FROM entities e LEFT JOIN observations o ON o.entity_id=e.id GROUP BY e.id`); err != nil {
			return err
		}
		if _, err := db.Exec(`CREATE TABLE IF NOT EXISTS relations_names (
            from_name TEXT NOT NULL,
            to_name   TEXT NOT NULL,
            relation_type TEXT NOT NULL,
            UNIQUE(from_name, to_name, relation_type)
        )`); err != nil {
			return err
		}
		if _, err := db.Exec(`INSERT OR IGNORE INTO relations_names(from_name, to_name, relation_type)
            SELECT e1.name, e2.name, r.relation_type FROM relations r
              JOIN entities e1 ON e1.id=r.from_id
              JOIN entities e2 ON e2.id=r.to_id`); err != nil {
			return err
		}
		if _, err := db.Exec(`UPDATE schema_version SET version = 4`); err != nil {
			return err
		}
		v = 4
	}
	if v < 5 {
		if _, err := db.Exec(`CREATE TABLE IF NOT EXISTS graph_meta (version INTEGER NOT NULL)`); err != nil {
			return err
		}
		var cnt int
		if err := db.QueryRow(`SELECT COUNT(1) FROM graph_meta`).Scan(&cnt); err != nil {
			return err
		}
		if cnt == 0 {
			if _, err := db.Exec(`INSERT INTO graph_meta(version) VALUES(0)`); err != nil {
				return err
			}
		}
		if _, err := db.Exec(`UPDATE schema_version SET version = 5`); err != nil {
			return err
		}
		v = 5
	}
	return nil
}

func (m *GraphManager) tokens(s string) []string {
	x := strings.ToLower(s)
	var out []string
	var cur []rune
	for _, r := range x {
		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') {
			cur = append(cur, r)
		} else {
			if len(cur) > 0 {
				out = append(out, string(cur))
				cur = nil
			}
		}
	}
	if len(cur) > 0 {
		out = append(out, string(cur))
	}
	return out
}

func (m *GraphManager) indexEntityTokens(id, name, entityType string, observations []string, remove bool) error {
	return m.bolt.Update(func(tx *bbolt.Tx) error {
		b := tx.Bucket([]byte("term"))
		apply := func(tok string) error {
			sb, err := b.CreateBucketIfNotExists([]byte(tok))
			if err != nil {
				return err
			}
			if remove {
				return sb.Delete([]byte(id))
			}
			return sb.Put([]byte(id), []byte{1})
		}
		for _, t := range m.tokens(name) {
			if err := apply(t); err != nil {
				return err
			}
		}
		for _, t := range m.tokens(entityType) {
			if err := apply(t); err != nil {
				return err
			}
		}
		for _, o := range observations {
			for _, t := range m.tokens(o) {
				if err := apply(t); err != nil {
					return err
				}
			}
		}
		return nil
	})
}

func (m *GraphManager) ensureOpen() error {
	if m.db == nil || m.bolt == nil {
		return m.open()
	}
	return nil
}

func generateID() (string, error) {
	const chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_/=\\|"
	var b [5]byte
	for i := 0; i < 5; i++ {
		n, err := rand.Int(rand.Reader, big.NewInt(int64(len(chars))))
		if err != nil {
			return "", err
		}
		b[i] = chars[n.Int64()]
	}
	return string(b[:]), nil
}

func (m *GraphManager) CreateEntities(entities []Entity) ([]Entity, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if err := m.ensureOpen(); err != nil {
		return nil, err
	}
	tx, err := m.db.Begin()
	if err != nil {
		return nil, err
	}
	defer func() { _ = tx.Rollback() }()
	var created []Entity
	for _, in := range entities {
		var exists int
		if err := tx.QueryRow("SELECT COUNT(1) FROM entities WHERE name = ?", in.Name).Scan(&exists); err != nil {
			return nil, err
		}
		if exists != 0 {
			continue
		}
		id, err := generateID()
		if err != nil {
			return nil, err
		}
		_, err = tx.Exec("INSERT INTO entities(id, name, entity_type) VALUES(?, ?, ?)", id, in.Name, in.EntityType)
		if err != nil {
			return nil, err
		}
		for _, o := range in.Observations {
			if _, err := tx.Exec("INSERT OR IGNORE INTO observations(entity_id, content) VALUES(?, ?)", id, o); err != nil {
				return nil, err
			}
		}
		created = append(created, Entity{ID: id, Name: in.Name, EntityType: in.EntityType, Observations: append([]string{}, in.Observations...)})
		if _, err := tx.Exec("INSERT OR REPLACE INTO entity_doc(entity_id, obs) VALUES(?, ?)", id, strings.Join(in.Observations, "\n")); err != nil {
			return nil, err
		}
		if err := m.ftsUpdate(tx, id, in.Name, in.EntityType, in.Observations); err != nil {
			return nil, err
		}
	}
	if _, err := tx.Exec("UPDATE graph_meta SET version = version + 1"); err != nil {
		_ = tx.Rollback()
		return nil, err
	}
	if err := tx.Commit(); err != nil {
		return nil, err
	}
	m.cacheReady = false
	m.indexReady = false
	return created, nil
}

func (m *GraphManager) CreateRelations(relations []Relation) ([]Relation, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if err := m.ensureOpen(); err != nil {
		return nil, err
	}
	tx, err := m.db.Begin()
	if err != nil {
		return nil, err
	}
	defer func() { _ = tx.Rollback() }()
	var created []Relation
	for _, in := range relations {
		var fromID, toID string
		if err := tx.QueryRow("SELECT id FROM entities WHERE name = ?", in.From).Scan(&fromID); err != nil {
			return nil, errors.New("entity not found: " + in.From)
		}
		if err := tx.QueryRow("SELECT id FROM entities WHERE name = ?", in.To).Scan(&toID); err != nil {
			return nil, errors.New("entity not found: " + in.To)
		}
		res, err := tx.Exec("INSERT OR IGNORE INTO relations(from_id, to_id, relation_type) VALUES(?, ?, ?)", fromID, toID, in.RelationType)
		if err != nil {
			return nil, err
		}
		if rows, _ := res.RowsAffected(); rows > 0 {
			created = append(created, in)
			if _, err := tx.Exec("INSERT OR IGNORE INTO relations_names(from_name, to_name, relation_type) VALUES(?, ?, ?)", in.From, in.To, in.RelationType); err != nil {
				return nil, err
			}
		}
	}
	if err := tx.Commit(); err != nil {
		return nil, err
	}
	return created, nil
}

func (m *GraphManager) AddObservations(obs []struct {
	EntityName string   `json:"entityName"`
	Contents   []string `json:"contents"`
}) ([]struct {
	EntityName        string   `json:"entityName"`
	AddedObservations []string `json:"addedObservations"`
}, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if err := m.ensureOpen(); err != nil {
		return nil, err
	}
	tx, err := m.db.Begin()
	if err != nil {
		return nil, err
	}
	defer func() { _ = tx.Rollback() }()
	var res []struct {
		EntityName        string   `json:"entityName"`
		AddedObservations []string `json:"addedObservations"`
	}
	for _, o := range obs {
		var id string
		if err := tx.QueryRow("SELECT id FROM entities WHERE name = ?", o.EntityName).Scan(&id); err != nil {
			return nil, errors.New("entity not found: " + o.EntityName)
		}
		var added []string
		for _, c := range o.Contents {
			res, err := tx.Exec("INSERT OR IGNORE INTO observations(entity_id, content) VALUES(?, ?)", id, c)
			if err != nil {
				return nil, err
			}
			if rows, _ := res.RowsAffected(); rows > 0 {
				added = append(added, c)
			}
		}
		res = append(res, struct {
			EntityName        string   `json:"entityName"`
			AddedObservations []string `json:"addedObservations"`
		}{EntityName: o.EntityName, AddedObservations: added})
		// refresh FTS doc
		var name, et string
		if err := tx.QueryRow("SELECT name, entity_type FROM entities WHERE id=?", id).Scan(&name, &et); err != nil {
			return nil, err
		}
		orows, err := tx.Query("SELECT content FROM observations WHERE entity_id=?", id)
		if err != nil {
			return nil, err
		}
		var all []string
		for orows.Next() {
			var c string
			if err := orows.Scan(&c); err != nil {
				orows.Close()
				return nil, err
			}
			all = append(all, c)
		}
		orows.Close()
		if _, err := tx.Exec("INSERT OR REPLACE INTO entity_doc(entity_id, obs) VALUES(?, ?)", id, strings.Join(all, "\n")); err != nil {
			return nil, err
		}
		if err := m.ftsUpdate(tx, id, name, et, all); err != nil {
			return nil, err
		}
	}
	if _, err := tx.Exec("UPDATE graph_meta SET version = version + 1"); err != nil {
		_ = tx.Rollback()
		return nil, err
	}
	if err := tx.Commit(); err != nil {
		return nil, err
	}
	m.cacheReady = false
	m.indexReady = false
	return res, nil
}

func (m *GraphManager) DeleteEntities(names []string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if err := m.ensureOpen(); err != nil {
		return err
	}
	tx, err := m.db.Begin()
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback() }()
	for _, n := range names {
		var id, et string
		if err := tx.QueryRow("SELECT id, entity_type FROM entities WHERE name = ?", n).Scan(&id, &et); err == nil {
			if _, err := tx.Exec("DELETE FROM fts WHERE entity_id=?", id); err != nil {
				return err
			}
			if _, err := tx.Exec("DELETE FROM entity_doc WHERE entity_id=?", id); err != nil {
				return err
			}
		}
		if _, err := tx.Exec("DELETE FROM relations_names WHERE from_name=? OR to_name=?", n, n); err != nil {
			return err
		}
		if _, err := tx.Exec("DELETE FROM entities WHERE name = ?", n); err != nil {
			return err
		}
	}
	if _, err := tx.Exec("UPDATE graph_meta SET version = version + 1"); err != nil {
		_ = tx.Rollback()
		return err
	}
	if err := tx.Commit(); err != nil {
		return err
	}
	m.cacheReady = false
	m.indexReady = false
	return nil
}

func (m *GraphManager) DeleteObservations(deletions []struct {
	EntityName   string   `json:"entityName"`
	Observations []string `json:"observations"`
}) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if err := m.ensureOpen(); err != nil {
		return err
	}
	tx, err := m.db.Begin()
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback() }()
	for _, d := range deletions {
		var id string
		if err := tx.QueryRow("SELECT id FROM entities WHERE name = ?", d.EntityName).Scan(&id); err != nil {
			return errors.New("entity not found: " + d.EntityName)
		}
		for _, o := range d.Observations {
			if _, err := tx.Exec("DELETE FROM observations WHERE entity_id = ? AND content = ?", id, o); err != nil {
				return err
			}
		}
		// refresh FTS doc
		var name, et string
		if err := tx.QueryRow("SELECT name, entity_type FROM entities WHERE id=?", id).Scan(&name, &et); err != nil {
			return err
		}
		orows, err := tx.Query("SELECT content FROM observations WHERE entity_id=?", id)
		if err != nil {
			return err
		}
		var all []string
		for orows.Next() {
			var c string
			if err := orows.Scan(&c); err != nil {
				orows.Close()
				return err
			}
			all = append(all, c)
		}
		orows.Close()
		if _, err := tx.Exec("INSERT OR REPLACE INTO entity_doc(entity_id, obs) VALUES(?, ?)", id, strings.Join(all, "\n")); err != nil {
			return err
		}
		if err := m.ftsUpdate(tx, id, name, et, all); err != nil {
			return err
		}
	}
	return tx.Commit()
}

func (m *GraphManager) DeleteRelations(relations []Relation) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if err := m.ensureOpen(); err != nil {
		return err
	}
	tx, err := m.db.Begin()
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback() }()
	for _, r := range relations {
		var fromID, toID string
		if err := tx.QueryRow("SELECT id FROM entities WHERE name = ?", r.From).Scan(&fromID); err != nil {
			return errors.New("entity not found: " + r.From)
		}
		if err := tx.QueryRow("SELECT id FROM entities WHERE name = ?", r.To).Scan(&toID); err != nil {
			return errors.New("entity not found: " + r.To)
		}
		if _, err := tx.Exec("DELETE FROM relations WHERE from_id=? AND to_id=? AND relation_type=?", fromID, toID, r.RelationType); err != nil {
			return err
		}
		if _, err := tx.Exec("DELETE FROM relations_names WHERE from_name=? AND to_name=? AND relation_type=?", r.From, r.To, r.RelationType); err != nil {
			return err
		}
	}
	return tx.Commit()
}

func (m *GraphManager) ReadGraph() (KnowledgeGraph, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if err := m.ensureOpen(); err != nil {
		return KnowledgeGraph{}, err
	}
	if m.cacheReady {
		return m.cacheGraph, nil
	}
	rows, err := m.db.Query("SELECT e.id, e.name, e.entity_type, ifnull(doc.obs, '') FROM entities e LEFT JOIN entity_doc doc ON doc.entity_id=e.id")
	if err != nil {
		return KnowledgeGraph{}, err
	}
	defer rows.Close()
	type raw struct{ id, name, et, cat string }
	var raws []raw
	for rows.Next() {
		var r raw
		if err := rows.Scan(&r.id, &r.name, &r.et, &r.cat); err != nil {
			return KnowledgeGraph{}, err
		}
		raws = append(raws, r)
	}
	ents := make([]Entity, len(raws))
	jobs := make(chan int, 1024)
	var wg sync.WaitGroup
	workers := runtime.GOMAXPROCS(0)
	if workers < 1 {
		workers = 1
	}
	for w := 0; w < workers; w++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := range jobs {
				r := raws[i]
				var obs []string
				if len(r.cat) > 0 {
					obs = strings.Split(r.cat, "\n")
				}
				ents[i] = Entity{ID: r.id, Name: r.name, EntityType: r.et, Observations: obs}
			}
		}()
	}
	for i := range raws {
		jobs <- i
	}
	close(jobs)
	rrows, err := m.db.Query("SELECT from_name, to_name, relation_type FROM relations_names")
	if err != nil {
		return KnowledgeGraph{}, err
	}
	defer rrows.Close()
	var rels []Relation
	for rrows.Next() {
		var fromName, toName, rt string
		if err := rrows.Scan(&fromName, &toName, &rt); err != nil {
			return KnowledgeGraph{}, err
		}
		rels = append(rels, Relation{From: fromName, To: toName, RelationType: rt})
	}
	wg.Wait()
	g := KnowledgeGraph{Entities: ents, Relations: rels}
	m.cacheGraph = g
	m.cacheReady = true
	if m.tryLoadIndexSnapshot() {
		return g, nil
	}
	if !m.indexBuilding && m.tokIndex == nil {
		m.indexBuilding = true
		entsCopy := make([]Entity, len(ents))
		copy(entsCopy, ents)
		go func() {
			jobs := make(chan int, 1024)
			var wwg sync.WaitGroup
			localTok := make([]map[string]map[string]struct{}, workers)
			localBy := make([]map[string]Entity, workers)
			for i := 0; i < workers; i++ {
				localTok[i] = make(map[string]map[string]struct{})
				localBy[i] = make(map[string]Entity)
			}
			for w := 0; w < workers; w++ {
				w := w
				wwg.Add(1)
				go func() {
					defer wwg.Done()
					lt := localTok[w]
					lb := localBy[w]
					for i := range jobs {
						e := entsCopy[i]
						lb[e.Name] = e
						add := func(s string) {
							ts := m.tokens(s)
							for _, t := range ts {
								b := lt[t]
								if b == nil {
									b = make(map[string]struct{})
									lt[t] = b
								}
								b[e.Name] = struct{}{}
							}
						}
						add(e.Name)
						add(e.EntityType)
						for _, o := range e.Observations {
							add(o)
						}
					}
				}()
			}
			for i := range entsCopy {
				jobs <- i
			}
			close(jobs)
			wwg.Wait()
			// merge
			tok := make(map[string]map[string]struct{}, len(entsCopy)*2)
			byName := make(map[string]Entity, len(entsCopy))
			for i := 0; i < workers; i++ {
				for n, e := range localBy[i] {
					byName[n] = e
				}
				for k, set := range localTok[i] {
					dst := tok[k]
					if dst == nil {
						dst = make(map[string]struct{})
						tok[k] = dst
					}
					for n := range set {
						dst[n] = struct{}{}
					}
				}
			}
			m.mu.Lock()
			m.tokIndex = tok
			m.entsByName = byName
			m.indexReady = true
			m.indexBuilding = false
			m.mu.Unlock()
			ver := m.currentDataVersion()
			_ = m.saveIndexSnapshot(ver)
		}()
	}
	return g, nil
}

func (m *GraphManager) SearchNodes(query string) (KnowledgeGraph, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if err := m.ensureOpen(); err != nil {
		return KnowledgeGraph{}, err
	}
	q := strings.ToLower(query)
	toks := m.tokens(q)
	if len(toks) == 0 {
		return KnowledgeGraph{Entities: nil, Relations: nil}, nil
	}
	if m.cacheReady && m.tokIndex != nil {
		// intersect sets
		var set map[string]struct{}
		for i, t := range toks {
			bucket := m.tokIndex[t]
			if bucket == nil {
				return KnowledgeGraph{Entities: nil, Relations: nil}, nil
			}
			if i == 0 {
				set = make(map[string]struct{}, len(bucket))
				for n := range bucket {
					set[n] = struct{}{}
				}
			} else {
				for n := range set {
					if _, ok := bucket[n]; !ok {
						delete(set, n)
					}
				}
			}
			if len(set) == 0 {
				return KnowledgeGraph{Entities: nil, Relations: nil}, nil
			}
		}
		ents := make([]Entity, 0, len(set))
		for n := range set {
			if e, ok := m.entsByName[n]; ok {
				ents = append(ents, e)
			}
		}
		// filter cached relations
		rels := make([]Relation, 0)
		nameSet := set
		for _, r := range m.cacheGraph.Relations {
			if _, ok := nameSet[r.From]; ok {
				if _, ok2 := nameSet[r.To]; ok2 {
					rels = append(rels, r)
				}
			}
		}
		return KnowledgeGraph{Entities: ents, Relations: rels}, nil
	}
	var qb strings.Builder
	for i, t := range toks {
		if i > 0 {
			qb.WriteString(" AND ")
		}
		qb.WriteString(t)
		qb.WriteByte('*')
	}
	match := qb.String()
	erows, err := m.db.Query("WITH matched AS (SELECT entity_id FROM fts WHERE fts MATCH ? LIMIT 1000) "+
		"SELECT e.id, e.name, e.entity_type, ifnull(doc.obs, '') "+
		"FROM matched m JOIN entities e ON e.id=m.entity_id LEFT JOIN entity_doc doc ON doc.entity_id=e.id", match)
	if err != nil {
		return KnowledgeGraph{}, err
	}
	defer erows.Close()
	var ents []Entity
	for erows.Next() {
		var id, name, et, obsText string
		if err := erows.Scan(&id, &name, &et, &obsText); err != nil {
			return KnowledgeGraph{}, err
		}
		var obs []string
		if len(obsText) > 0 {
			obs = strings.Split(obsText, "\n")
		}
		ents = append(ents, Entity{ID: id, Name: name, EntityType: et, Observations: obs})
	}
	if len(ents) == 0 {
		return KnowledgeGraph{Entities: nil, Relations: nil}, nil
	}
	rrows, err := m.db.Query("WITH matched AS (SELECT entity_id FROM fts WHERE fts MATCH ? LIMIT 1000) "+
		"SELECT e1.name, e2.name, r.relation_type "+
		"FROM relations r JOIN entities e1 ON e1.id=r.from_id JOIN entities e2 ON e2.id=r.to_id "+
		"JOIN matched m1 ON m1.entity_id=r.from_id JOIN matched m2 ON m2.entity_id=r.to_id", match)
	if err != nil {
		return KnowledgeGraph{}, err
	}
	defer rrows.Close()
	var rels []Relation
	setNames := make(map[string]struct{}, len(ents))
	for _, e := range ents {
		setNames[e.Name] = struct{}{}
	}
	for rrows.Next() {
		var fromName, toName, rt string
		if err := rrows.Scan(&fromName, &toName, &rt); err != nil {
			return KnowledgeGraph{}, err
		}
		if _, ok := setNames[fromName]; ok {
			if _, ok2 := setNames[toName]; ok2 {
				rels = append(rels, Relation{From: fromName, To: toName, RelationType: rt})
			}
		}
	}
	return KnowledgeGraph{Entities: ents, Relations: rels}, nil
}

func (m *GraphManager) OpenNodes(names []string) (KnowledgeGraph, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if err := m.ensureOpen(); err != nil {
		return KnowledgeGraph{}, err
	}
	rows, err := m.db.Query("SELECT e.id, e.name, e.entity_type, o.content FROM entities e LEFT JOIN observations o ON o.entity_id=e.id WHERE e.name IN ("+placeholders(len(names))+")", anyArgs(names)...)
	if err != nil {
		return KnowledgeGraph{}, err
	}
	defer rows.Close()
	emap := make(map[string]*Entity)
	for rows.Next() {
		var id, name, et string
		var content sql.NullString
		if err := rows.Scan(&id, &name, &et, &content); err != nil {
			return KnowledgeGraph{}, err
		}
		e := emap[id]
		if e == nil {
			e = &Entity{ID: id, Name: name, EntityType: et}
			emap[id] = e
		}
		if content.Valid {
			e.Observations = append(e.Observations, content.String)
		}
	}
	ents := make([]Entity, 0, len(emap))
	ids := make([]string, 0, len(emap))
	set := make(map[string]struct{}, len(emap))
	for _, e := range emap {
		ents = append(ents, *e)
		ids = append(ids, e.ID)
		set[e.Name] = struct{}{}
	}
	var rels []Relation
	rrows, err := m.db.Query("SELECT e1.name, e2.name, r.relation_type FROM relations r JOIN entities e1 ON e1.id=r.from_id JOIN entities e2 ON e2.id=r.to_id WHERE r.from_id IN ("+placeholders(len(ids))+") AND r.to_id IN ("+placeholders(len(ids))+")", append(anyArgs(ids), anyArgs(ids)...)...)
	if err != nil {
		return KnowledgeGraph{}, err
	}
	defer rrows.Close()
	for rrows.Next() {
		var fromName, toName, rt string
		if err := rrows.Scan(&fromName, &toName, &rt); err != nil {
			return KnowledgeGraph{}, err
		}
		if _, ok := set[fromName]; ok {
			if _, ok2 := set[toName]; ok2 {
				rels = append(rels, Relation{From: fromName, To: toName, RelationType: rt})
			}
		}
	}
	return KnowledgeGraph{Entities: ents, Relations: rels}, nil
}

func placeholders(n int) string {
	if n <= 0 {
		return ""
	}
	s := "?"
	for i := 1; i < n; i++ {
		s += ",?"
	}
	return s
}

func anyArgs(ss []string) []any {
	a := make([]any, len(ss))
	for i := range ss {
		a[i] = ss[i]
	}
	return a
}

func (m *GraphManager) Close() error {
	m.mu.Lock()
	defer m.mu.Unlock()
	var err error

	if m.db != nil {
		if e := m.db.Close(); e != nil {
			err = e
		}
		m.db = nil
	}

	if m.bolt != nil {
		if e := m.bolt.Close(); e != nil && err == nil {
			err = e
		}
		m.bolt = nil
	}
	return err
}