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 tags (node_id TEXT NOT NULL, tag TEXT NOT NULL, PRIMARY KEY (node_id, tag), FOREIGN KEY (node_id) REFERENCES nodes(id) ON DELETE CASCADE)`,
	`CREATE TABLE IF NOT EXISTS rels (from_id TEXT NOT NULL, to_id TEXT NOT NULL, rel_type TEXT NOT NULL, PRIMARY KEY (from_id, to_id, rel_type), FOREIGN KEY (from_id) REFERENCES nodes(id) ON DELETE CASCADE, FOREIGN KEY (to_id) REFERENCES nodes(id) ON DELETE CASCADE)`,
	`CREATE INDEX IF NOT EXISTS idx_tags_tag ON tags(tag)`,
	`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 walks up from the current working directory to find
// an .ax.db file and opens it.
func FindAndOpenSQLiteStore() (Store, error) {
	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
		}
	}
}

// NewSQLiteStore opens a SQLite database at the given path and applies
// per-connection PRAGMAs. foreign_keys must be set on every connection.
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)
	}
	for _, q := range []string{"PRAGMA journal_mode=WAL", "PRAGMA busy_timeout=5000", "PRAGMA foreign_keys=ON"} {
		if _, err := db.Exec(q); err != nil {
			db.Close()
			return nil, err
		}
	}
	return &SQLiteStore{db: db}, 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 tag FROM tags WHERE node_id = ?", id)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		var tag string
		rows.Scan(&tag)
		n.AddTag(tag)
	}
	rows2, err := q.Query("SELECT to_id, rel_type FROM rels WHERE from_id = ?", id)
	if err != nil {
		return nil, err
	}
	defer rows2.Close()
	for rows2.Next() {
		var toID, relType string
		rows2.Scan(&toID, &relType)
		n.AddRelation(models.RelType(relType), 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) }

// --- Tag operations ---

func addTag(q querier, nodeID, tag string) error {
	_, err := q.Exec("INSERT OR IGNORE INTO tags (node_id, tag) VALUES (?, ?)", nodeID, tag)
	return err
}

func removeTag(q querier, nodeID, tag string) error {
	_, err := q.Exec("DELETE FROM tags WHERE node_id = ? AND tag = ?", nodeID, tag)
	return err
}

func (s *SQLiteStore) AddTag(nodeID, tag string) error    { return addTag(s.db, nodeID, tag) }
func (s *SQLiteStore) RemoveTag(nodeID, tag string) error { return removeTag(s.db, nodeID, tag) }
func (s *txStore) AddTag(nodeID, tag string) error        { return addTag(s.tx, nodeID, tag) }
func (s *txStore) RemoveTag(nodeID, tag string) error     { return removeTag(s.tx, nodeID, tag) }

// --- Edge operations ---

func addEdge(q querier, fromID, toID string, relType models.RelType) error {
	_, err := q.Exec("INSERT OR IGNORE INTO rels (from_id, to_id, rel_type) VALUES (?, ?, ?)", fromID, toID, relType)
	return err
}

func removeEdge(q querier, fromID, toID string, relType models.RelType) error {
	_, err := q.Exec("DELETE FROM rels WHERE from_id = ? AND to_id = ? AND rel_type = ?", fromID, toID, relType)
	return err
}

func (s *SQLiteStore) AddEdge(fromID, toID string, relType models.RelType) error {
	return addEdge(s.db, fromID, toID, relType)
}
func (s *SQLiteStore) RemoveEdge(fromID, toID string, relType models.RelType) error {
	return removeEdge(s.db, fromID, toID, relType)
}
func (s *txStore) AddEdge(fromID, toID string, relType models.RelType) error {
	return addEdge(s.tx, fromID, toID, relType)
}
func (s *txStore) RemoveEdge(fromID, toID string, relType models.RelType) error {
	return removeEdge(s.tx, fromID, toID, relType)
}

// --- FindNodes ---

func findNodes(q querier, tagPrefixes []string, edgeFilters []*models.Rel) ([]*models.Node, error) {
	// Edge filter args come first (in JOIN ON clauses), tag args come last (in HAVING).
	query := "SELECT DISTINCT n.id FROM nodes n"
	var joins []string
	var args []any

	// Edge filtering: each filter becomes an inner join with conditions in the ON clause.
	for i, ef := range edgeFilters {
		alias := fmt.Sprintf("r%d", i)
		joins = append(joins, fmt.Sprintf(
			"JOIN rels %s ON n.id = %s.from_id AND %s.rel_type = ? AND %s.to_id = ?",
			alias, alias, alias, alias,
		))
		args = append(args, string(ef.Type), ef.Target)
	}

	// Tag filtering: join tags and use HAVING with aggregate to require all prefixes.
	joins = append(joins, "JOIN tags t_tag ON n.id = t_tag.node_id")
	if len(tagPrefixes) == 0 {
		tagPrefixes = []string{""}
	}
	cond := ""
	for _, t := range tagPrefixes {
		cond += "t_tag.tag LIKE ? || '%' OR "
		args = append(args, t)
	}
	having := "SUM(CASE WHEN " + cond[:len(cond)-4] + " THEN 1 ELSE 0 END) >= ?"
	args = append(args, len(tagPrefixes))

	query += " " + strings.Join(joins, " ")
	query += " GROUP BY n.id HAVING " + having
	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(tagPrefixes []string, edgeFilters []*models.Rel) ([]*models.Node, error) {
	return findNodes(s.db, tagPrefixes, edgeFilters)
}

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