ax/store/sqlite.go

package store

import (
	"axolotl/models"
	"database/sql"
	"errors"
	"fmt"
	"math/rand"
	"os"
	"path/filepath"
	"strings"

	_ "modernc.org/sqlite"
)

var migrations = []string{
	`CREATE TABLE IF NOT EXISTS nodes (id TEXT PRIMARY KEY, title TEXT NOT NULL, content TEXT, due_date TEXT, created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP)`,
	`CREATE TABLE IF NOT EXISTS rels (from_id TEXT NOT NULL, rel_name TEXT NOT NULL, to_id TEXT NOT NULL DEFAULT '', PRIMARY KEY (from_id, rel_name, to_id), FOREIGN KEY (from_id) REFERENCES nodes(id) ON DELETE CASCADE)`,
	`CREATE INDEX IF NOT EXISTS idx_rels_from ON rels(from_id)`,
	`CREATE INDEX IF NOT EXISTS idx_rels_to ON rels(to_id)`,
}

// querier abstracts *sql.DB and *sql.Tx so SQL helpers work for both.
type querier interface {
	Exec(query string, args ...any) (sql.Result, error)
	Query(query string, args ...any) (*sql.Rows, error)
	QueryRow(query string, args ...any) *sql.Row
}

// SQLiteStore is the top-level Store backed by a SQLite database file.
type SQLiteStore struct {
	db *sql.DB
}

// txStore wraps an active transaction. Its Transaction method is a no-op
// passthrough so nested calls reuse the same transaction.
type txStore struct {
	db *sql.DB
	tx *sql.Tx
}

// InitSQLiteStore creates the database file and applies the schema.
// It is idempotent (uses CREATE TABLE IF NOT EXISTS).
func InitSQLiteStore(path string) error {
	if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
		return err
	}
	db, err := sql.Open("sqlite", path)
	if err != nil {
		return err
	}
	defer db.Close()
	pragmas := []string{"PRAGMA journal_mode=WAL", "PRAGMA busy_timeout=5000", "PRAGMA foreign_keys=ON"}
	for _, q := range append(pragmas, migrations...) {
		if _, err := db.Exec(q); err != nil {
			return err
		}
	}
	return nil
}

// FindAndOpenSQLiteStore opens the SQLite database. If the AX_DB_PATH environment
// variable is set, it uses that path directly. Otherwise, it walks up from the
// current working directory to find an .ax.db file.
func FindAndOpenSQLiteStore() (Store, error) {
	if dbpath := os.Getenv("AX_DB_PATH"); dbpath != "" {
		return NewSQLiteStore(dbpath)
	}
	dir, err := filepath.Abs(".")
	if err != nil {
		return nil, err
	}
	for {
		dbpath := filepath.Join(dir, ".ax.db")
		if _, err := os.Stat(dbpath); err == nil {
			return NewSQLiteStore(dbpath)
		}
		if parent := filepath.Dir(dir); parent == dir {
			return nil, errors.New("no .ax.db found (run 'ax init' first)")
		} else {
			dir = parent
		}
	}
}

// FindOrInitSQLiteStore is like FindAndOpenSQLiteStore but intended for server
// mode: if no .ax.db is found it creates and initialises one in the current
// working directory instead of returning an error.
func FindOrInitSQLiteStore() (Store, error) {
	if dbpath := os.Getenv("AX_DB_PATH"); dbpath != "" {
		if err := InitSQLiteStore(dbpath); err != nil {
			return nil, err
		}
		return NewSQLiteStore(dbpath)
	}
	dir, err := filepath.Abs(".")
	if err != nil {
		return nil, err
	}
	for {
		dbpath := filepath.Join(dir, ".ax.db")
		if _, err := os.Stat(dbpath); err == nil {
			return NewSQLiteStore(dbpath)
		}
		if parent := filepath.Dir(dir); parent == dir {
			break
		} else {
			dir = parent
		}
	}
	// Not found — create and initialise in CWD.
	cwd, _ := filepath.Abs(".")
	dbpath := filepath.Join(cwd, ".ax.db")
	if err := InitSQLiteStore(dbpath); err != nil {
		return nil, err
	}
	return NewSQLiteStore(dbpath)
}

// NewSQLiteStore opens a SQLite database at the given path, runs a one-time
// schema migration if needed, then applies per-connection PRAGMAs.
func NewSQLiteStore(path string) (Store, error) {
	db, err := sql.Open("sqlite", path)
	if err != nil {
		return nil, fmt.Errorf("failed to open database: %w", err)
	}
	// FK must be OFF during migration (table drops/renames).
	for _, q := range []string{"PRAGMA journal_mode=WAL", "PRAGMA busy_timeout=5000", "PRAGMA foreign_keys=OFF"} {
		if _, err := db.Exec(q); err != nil {
			db.Close()
			return nil, err
		}
	}
	if err := migrateSchema(db); err != nil {
		db.Close()
		return nil, fmt.Errorf("schema migration failed: %w", err)
	}
	if _, err := db.Exec("PRAGMA foreign_keys=ON"); err != nil {
		db.Close()
		return nil, err
	}
	return &SQLiteStore{db: db}, nil
}

// migrateSchema migrates from the legacy two-table (tags + rels) schema to the
// unified rels-only schema. It is a no-op when migration has already been applied.
func migrateSchema(db *sql.DB) error {
	var tagsExists int
	if err := db.QueryRow("SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='tags'").Scan(&tagsExists); err != nil {
		return err
	}
	if tagsExists == 0 {
		return nil // already on new schema
	}
	for _, stmt := range []string{
		`CREATE TABLE rels_new (from_id TEXT NOT NULL, rel_name TEXT NOT NULL, to_id TEXT NOT NULL DEFAULT '', PRIMARY KEY (from_id, rel_name, to_id), FOREIGN KEY (from_id) REFERENCES nodes(id) ON DELETE CASCADE)`,
		`INSERT OR IGNORE INTO rels_new (from_id, rel_name, to_id) SELECT from_id, rel_type, to_id FROM rels`,
		`INSERT OR IGNORE INTO rels_new (from_id, rel_name, to_id) SELECT node_id, tag, '' FROM tags`,
		`DROP TABLE rels`,
		`DROP TABLE tags`,
		`ALTER TABLE rels_new RENAME TO rels`,
		`CREATE INDEX IF NOT EXISTS idx_rels_from ON rels(from_id)`,
		`CREATE INDEX IF NOT EXISTS idx_rels_to ON rels(to_id)`,
	} {
		if _, err := db.Exec(stmt); err != nil {
			return err
		}
	}
	return nil
}

// --- Transaction ---

func (s *SQLiteStore) Transaction(fn func(Store) error) error {
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	if err := fn(&txStore{db: s.db, tx: tx}); err != nil {
		return err
	}
	return tx.Commit()
}

func (s *txStore) Transaction(fn func(Store) error) error {
	return fn(s) // already in a transaction — reuse it
}

// --- Node operations ---

func addNode(q querier, id, title, content, dueDate, createdAt, updatedAt string) error {
	var dd interface{}
	if dueDate != "" {
		dd = dueDate
	}
	_, err := q.Exec(
		"INSERT INTO nodes (id, title, content, due_date, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)",
		id, title, content, dd, createdAt, updatedAt,
	)
	return err
}

func getNode(q querier, id string) (*models.Node, error) {
	n := models.NewNode()
	err := q.QueryRow(
		"SELECT id, title, COALESCE(content, ''), COALESCE(due_date, ''), created_at, updated_at FROM nodes WHERE id = ?", id,
	).Scan(&n.ID, &n.Title, &n.Content, &n.DueDate, &n.CreatedAt, &n.UpdatedAt)
	if err != nil {
		return nil, err
	}
	rows, err := q.Query("SELECT rel_name, to_id FROM rels WHERE from_id = ?", id)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		var relName, toID string
		rows.Scan(&relName, &toID)
		if toID == "" {
			n.AddTag(relName)
		} else {
			n.AddRelation(models.RelType(relName), toID)
		}
	}
	return n, nil
}

func updateNode(q querier, id, title, content, dueDate, updatedAt string) error {
	var dd interface{}
	if dueDate != "" {
		dd = dueDate
	}
	_, err := q.Exec(
		"UPDATE nodes SET title = ?, content = ?, due_date = ?, updated_at = ? WHERE id = ?",
		title, content, dd, updatedAt, id,
	)
	return err
}

func deleteNode(q querier, id string) error {
	_, err := q.Exec("DELETE FROM nodes WHERE id = ?", id)
	return err
}

func nodeExists(q querier, id string) (bool, error) {
	var e bool
	err := q.QueryRow("SELECT EXISTS(SELECT 1 FROM nodes WHERE id = ?)", id).Scan(&e)
	return e, err
}

func (s *SQLiteStore) AddNode(id, title, content, dueDate, createdAt, updatedAt string) error {
	return addNode(s.db, id, title, content, dueDate, createdAt, updatedAt)
}
func (s *SQLiteStore) GetNode(id string) (*models.Node, error) { return getNode(s.db, id) }
func (s *SQLiteStore) UpdateNode(id, title, content, dueDate, updatedAt string) error {
	return updateNode(s.db, id, title, content, dueDate, updatedAt)
}
func (s *SQLiteStore) DeleteNode(id string) error         { return deleteNode(s.db, id) }
func (s *SQLiteStore) NodeExists(id string) (bool, error) { return nodeExists(s.db, id) }

func (s *txStore) AddNode(id, title, content, dueDate, createdAt, updatedAt string) error {
	return addNode(s.tx, id, title, content, dueDate, createdAt, updatedAt)
}
func (s *txStore) GetNode(id string) (*models.Node, error) { return getNode(s.tx, id) }
func (s *txStore) UpdateNode(id, title, content, dueDate, updatedAt string) error {
	return updateNode(s.tx, id, title, content, dueDate, updatedAt)
}
func (s *txStore) DeleteNode(id string) error         { return deleteNode(s.tx, id) }
func (s *txStore) NodeExists(id string) (bool, error) { return nodeExists(s.db, id) }

// --- ID generation ---

func genID() string {
	b := make([]byte, 5)
	for i := range b {
		b[i] = "abcdefghijklmnopqrstuvwxyz"[rand.Intn(26)]
	}
	return string(b)
}

func generateID(q querier) (string, error) {
	for {
		id := genID()
		exists, err := nodeExists(q, id)
		if err != nil {
			return "", err
		}
		if !exists {
			return id, nil
		}
	}
}

func (s *SQLiteStore) GenerateID() (string, error) { return generateID(s.db) }
func (s *txStore) GenerateID() (string, error)     { return generateID(s.db) }

// --- Rel operations ---

func addRel(q querier, nodeID, relName, toID string) error {
	_, err := q.Exec("INSERT OR IGNORE INTO rels (from_id, rel_name, to_id) VALUES (?, ?, ?)", nodeID, relName, toID)
	return err
}

func removeRel(q querier, nodeID, relName, toID string) error {
	_, err := q.Exec("DELETE FROM rels WHERE from_id = ? AND rel_name = ? AND to_id = ?", nodeID, relName, toID)
	return err
}

func (s *SQLiteStore) AddRel(nodeID, relName, toID string) error {
	return addRel(s.db, nodeID, relName, toID)
}
func (s *SQLiteStore) RemoveRel(nodeID, relName, toID string) error {
	return removeRel(s.db, nodeID, relName, toID)
}
func (s *txStore) AddRel(nodeID, relName, toID string) error {
	return addRel(s.tx, nodeID, relName, toID)
}
func (s *txStore) RemoveRel(nodeID, relName, toID string) error {
	return removeRel(s.tx, nodeID, relName, toID)
}

// --- FindNodes ---

func findNodes(q querier, filters []*models.Rel) ([]*models.Node, error) {
	query := "SELECT DISTINCT n.id FROM nodes n"
	var joins []string
	var args []any

	for i, f := range filters {
		alias := fmt.Sprintf("r%d", i)
		if f.Target == "" {
			// Tag/property filter: match rels with empty to_id by rel_name prefix.
			joins = append(joins, fmt.Sprintf(
				"JOIN rels %s ON n.id = %s.from_id AND %s.to_id = '' AND %s.rel_name LIKE ? || '%%'",
				alias, alias, alias, alias,
			))
			args = append(args, string(f.Type))
		} else {
			// Edge filter: match rel by exact rel_name and to_id.
			joins = append(joins, fmt.Sprintf(
				"JOIN rels %s ON n.id = %s.from_id AND %s.rel_name = ? AND %s.to_id = ?",
				alias, alias, alias, alias,
			))
			args = append(args, string(f.Type), f.Target)
		}
	}

	if len(joins) > 0 {
		query += " " + strings.Join(joins, " ")
	}
	query += " ORDER BY n.created_at DESC"

	rows, err := q.Query(query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var ids []string
	for rows.Next() {
		var id string
		if err := rows.Scan(&id); err != nil {
			return nil, err
		}
		ids = append(ids, id)
	}

	var nodes []*models.Node
	for _, id := range ids {
		n, err := getNode(q, id)
		if err != nil {
			return nil, err
		}
		nodes = append(nodes, n)
	}
	return nodes, nil
}

func (s *SQLiteStore) FindNodes(filters []*models.Rel) ([]*models.Node, error) {
	return findNodes(s.db, filters)
}

func (s *txStore) FindNodes(filters []*models.Rel) ([]*models.Node, error) {
	return findNodes(s.tx, filters)
}