ax/service/node_service_sqlite.go

package service

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

	_ "modernc.org/sqlite"
)

type sqliteNodeService struct {
	db     *sql.DB
	userID string
}

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)`,
}

var mentionRegex = regexp.MustCompile(`@([a-z0-9_]+)`)

func mentions(t string) []string {
	seen := make(map[string]bool)
	for _, m := range mentionRegex.FindAllStringSubmatch(t, -1) {
		seen[m[1]] = true
	}
	return slices.Collect(maps.Keys(seen))
}

func InitSqliteDB(path string) error {
	if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
		return err
	}
	var err error
	db, err := sql.Open("sqlite", path)
	if err != nil {
		return err
	}
	for _, q := range append([]string{"PRAGMA journal_mode=WAL", "PRAGMA busy_timeout=5000", "PRAGMA foreign_keys=ON"}, migrations...) {
		if _, err := db.Exec(q); err != nil {
			return err
		}
	}
	return err
}

func GetSqliteDB(cfg Config) (*sql.DB, 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 {
			db, err := sql.Open("sqlite", dbpath)
			if err != nil {
				return nil, fmt.Errorf("failed to open database: %w", err)
			}
			return db, nil
		}
		if parent := filepath.Dir(dir); parent == dir {
			return nil, errors.New("no .ax.db found (run 'ax init' first)")
		} else {
			dir = parent
		}
	}
}

func (s *sqliteNodeService) GetByID(id string) (*models.Node, error) {
	n := models.NewNode()
	q := s.db.QueryRow("SELECT id, title, COALESCE(content, ''), COALESCE(due_date, ''), created_at, updated_at FROM nodes WHERE id = ?", id)
	if err := q.Scan(&n.ID, &n.Title, &n.Content, &n.DueDate, &n.CreatedAt, &n.UpdatedAt); err != nil {
		return nil, err
	}

	if rows, err := s.db.Query("SELECT tag FROM tags WHERE node_id = ?", id); err == nil {
		defer rows.Close()
		for rows.Next() {
			var tag string
			rows.Scan(&tag)
			n.AddTag(tag)
		}
	} else {
		return nil, err
	}

	if rows, err := s.db.Query("SELECT to_id, rel_type FROM rels WHERE from_id = ?", id); err == nil {
		defer rows.Close()
		for rows.Next() {
			var toID, relType string
			rows.Scan(&toID, &relType)
			n.AddRelation(models.RelType(relType), toID)
		}
	} else {
		return nil, err
	}
	return n, nil
}

func (s *sqliteNodeService) Exists(id string) (bool, error) {
	var e bool
	err := s.db.QueryRow("SELECT EXISTS(SELECT 1 FROM nodes WHERE id = ?)", id).Scan(&e)
	return e, err
}

func (s *sqliteNodeService) Delete(id string) error {
	_, err := s.db.Exec("DELETE FROM nodes WHERE id = ?", id)
	return err
}

func (s *sqliteNodeService) CanClose(id string) (bool, []string, error) {
	rows, err := s.db.Query("SELECT to_id FROM rels WHERE from_id = ? AND rel_type = ?", id, models.RelBlocks)
	if err != nil {
		return false, nil, err
	}
	defer rows.Close()
	var blocking []string
	for rows.Next() {
		var bID, tag string
		if err := rows.Scan(&bID); err != nil {
			return false, nil, err
		}
		if err := s.db.QueryRow("SELECT tag FROM tags WHERE node_id = ? AND tag LIKE '_status::%'", bID).Scan(&tag); err == sql.ErrNoRows {
			continue
		} else if err != nil {
			return false, nil, err
		}
		if strings.HasSuffix(tag, "::open") {
			blocking = append(blocking, bID)
		}
	}
	return len(blocking) == 0, blocking, nil
}

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

func (s *sqliteNodeService) generateUniqueID() string {
	for {
		id := genID()
		if exists, _ := s.Exists(id); !exists {
			return id
		}
	}
}

func (s *sqliteNodeService) Create(title, content, dueDate string, tags []string, rels map[models.RelType][]string) (*models.Node, error) {
	tx, err := s.db.Begin()
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()

	now, id := time.Now().UTC().Format(time.RFC3339), s.generateUniqueID()
	if _, err := tx.Exec("INSERT INTO nodes (id, title, content, due_date, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)",
		id, title, content, dueDate, now, now); err != nil {
		return nil, err
	}

	for _, m := range mentions(title + " " + content) {
		userID, err := s.resolveUserRef(tx, m)
		if err != nil {
			return nil, err
		}
		if _, err := tx.Exec("INSERT INTO rels (from_id, to_id, rel_type) VALUES (?, ?, ?)", id, userID, models.RelMentions); err != nil {
			return nil, err
		}
	}

	for _, t := range tags {
		if _, err := tx.Exec("INSERT INTO tags (node_id, tag) VALUES (?, ?)", id, t); err != nil {
			return nil, err
		}
	}

	hasCreated := false
	for rt, tgts := range rels {
		for _, tgt := range tgts {
			if rt == models.RelCreated {
				hasCreated = true
			}
			if rt == models.RelAssignee || rt == models.RelCreated {
				var err error
				if tgt, err = s.resolveUserRef(tx, tgt); err != nil {
					return nil, err
				}
			}
			if rt == models.RelInNamespace {
				var err error
				if tgt, err = s.resolveNamespaceRef(tx, tgt); err != nil {
					return nil, err
				}
			}
			if _, err := tx.Exec("INSERT INTO rels (from_id, to_id, rel_type) VALUES (?, ?, ?)", id, tgt, rt); err != nil {
				return nil, err
			}
		}
	}
	if !hasCreated {
		userID, err := s.resolveUserRef(tx, s.userID)
		if err != nil {
			return nil, err
		}
		if _, err := tx.Exec("INSERT INTO rels (from_id, to_id, rel_type) VALUES (?, ?, ?)", id, userID, models.RelCreated); err != nil {
			return nil, err
		}
	}

	if err := tx.Commit(); err != nil {
		return nil, err
	}
	return s.GetByID(id)
}

func (s *sqliteNodeService) Update(node *models.Node) error {
	current, err := s.GetByID(node.ID)
	if err != nil {
		return err
	}

	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	upd := func(col, val string) error {
		_, err := tx.Exec("UPDATE nodes SET "+col+" = ? WHERE id = ?", val, node.ID)
		return err
	}

	newTitle, newContent := current.Title, current.Content
	if node.Title != current.Title {
		if err := upd("title", node.Title); err != nil {
			return err
		}
		newTitle = node.Title
	}
	if node.Content != current.Content {
		if err := upd("content", node.Content); err != nil {
			return err
		}
		newContent = node.Content
	}
	if node.DueDate != current.DueDate {
		if node.DueDate == "" {
			if _, err := tx.Exec("UPDATE nodes SET due_date = NULL WHERE id = ?", node.ID); err != nil {
				return err
			}
		} else {
			if err := upd("due_date", node.DueDate); err != nil {
				return err
			}
		}
	}

	if node.Title != current.Title || node.Content != current.Content {
		newMentions := mentions(newTitle + " " + newContent)
		rows, err := tx.Query("SELECT to_id FROM rels WHERE from_id = ? AND rel_type = ?", node.ID, models.RelMentions)
		if err != nil {
			return err
		}
		existingMentionIDs := make(map[string]bool)
		for rows.Next() {
			var uid string
			if err := rows.Scan(&uid); err != nil {
				rows.Close()
				return err
			}
			existingMentionIDs[uid] = true
		}
		rows.Close()

		mentionedUserIDs := make(map[string]bool)
		for _, m := range newMentions {
			userID, err := s.resolveUserRef(tx, m)
			if err != nil {
				return err
			}
			mentionedUserIDs[userID] = true
			if !existingMentionIDs[userID] {
				if _, err := tx.Exec("INSERT INTO rels (from_id, to_id, rel_type) VALUES (?, ?, ?)", node.ID, userID, models.RelMentions); err != nil {
					return err
				}
			}
		}

		for uid := range existingMentionIDs {
			if !mentionedUserIDs[uid] {
				if _, err := tx.Exec("DELETE FROM rels WHERE from_id = ? AND to_id = ? AND rel_type = ?", node.ID, uid, models.RelMentions); err != nil {
					return err
				}
			}
		}
	}

	currentTags := current.Tags()
	nodeTags := node.Tags()
	for _, t := range currentTags {
		if !slices.Contains(nodeTags, t) {
			tx.Exec("DELETE FROM tags WHERE node_id = ? AND tag = ?", node.ID, t)
		}
	}
	for _, t := range nodeTags {
		if !slices.Contains(currentTags, t) {
			tx.Exec("INSERT OR IGNORE INTO tags (node_id, tag) VALUES (?, ?)", node.ID, t)
		}
	}

	currentRels := current.Relations()
	nodeRels := node.Relations()
	for rt, tgts := range currentRels {
		for _, tgt := range tgts {
			if nodeRels[rt] == nil || !slices.Contains(nodeRels[rt], tgt) {
				tx.Exec("DELETE FROM rels WHERE from_id = ? AND to_id = ? AND rel_type = ?", node.ID, tgt, rt)
			}
		}
	}
	for rt, tgts := range nodeRels {
		for _, tgt := range tgts {
			if currentRels[rt] == nil || !slices.Contains(currentRels[rt], tgt) {
				resolvedTgt := tgt
				if models.RelType(rt) == models.RelAssignee || models.RelType(rt) == models.RelCreated {
					var err error
					if resolvedTgt, err = s.resolveUserRef(tx, tgt); err != nil {
						return err
					}
				}
				if models.RelType(rt) == models.RelInNamespace {
					var err error
					if resolvedTgt, err = s.resolveNamespaceRef(tx, tgt); err != nil {
						return err
					}
				}
				tx.Exec("INSERT OR IGNORE INTO rels (from_id, to_id, rel_type) VALUES (?, ?, ?)", node.ID, resolvedTgt, rt)
			}
		}
	}

	if _, err := tx.Exec("UPDATE nodes SET updated_at = ? WHERE id = ?", time.Now().UTC().Format(time.RFC3339), node.ID); err != nil {
		return err
	}
	return tx.Commit()
}

func (s *sqliteNodeService) resolveIDByNameAndTypeTx(tx *sql.Tx, title, nodeType string) (string, error) {
	var id string
	query := `
		SELECT n.id FROM nodes n
		JOIN tags t ON n.id = t.node_id
		WHERE n.title = ? AND t.tag = ?
		LIMIT 1
	`
	tag := "_type::" + nodeType
	var err error
	if tx != nil {
		err = tx.QueryRow(query, title, tag).Scan(&id)
	} else {
		err = s.db.QueryRow(query, title, tag).Scan(&id)
	}
	if err == sql.ErrNoRows {
		return "", nil
	}
	return id, err
}

func (s *sqliteNodeService) resolveUserIDByNameTx(tx *sql.Tx, username string) (string, error) {
	return s.resolveIDByNameAndTypeTx(tx, username, "user")
}

func (s *sqliteNodeService) resolveUserIDByName(username string) (string, error) {
	return s.resolveUserIDByNameTx(nil, username)
}

func (s *sqliteNodeService) resolveNamespaceIDByNameTx(tx *sql.Tx, name string) (string, error) {
	return s.resolveIDByNameAndTypeTx(tx, name, "namespace")
}

func (s *sqliteNodeService) resolveNamespaceIDByName(name string) (string, error) {
	return s.resolveNamespaceIDByNameTx(nil, name)
}

func (s *sqliteNodeService) List(opts ...ListOption) ([]*models.Node, error) {
	f := &listFilter{}
	for _, opt := range opts {
		opt(f)
	}

	q, joins, whereConds, havingConds := "SELECT DISTINCT n.id FROM nodes n", []string{}, []string{}, []string{}
	var whereArgs, havingArgs []any
	if len(f.tagPrefixes) == 0 {
		f.tagPrefixes = append(f.tagPrefixes, "")
	}

	joins = append(joins, "JOIN tags t_tag ON n.id = t_tag.node_id")
	cond := ""
	for _, t := range f.tagPrefixes {
		cond += "t_tag.tag LIKE ? || '%' OR "
		havingArgs = append(havingArgs, t)
	}
	havingConds = append(havingConds, "SUM(CASE WHEN "+cond[:len(cond)-4]+" THEN 1 ELSE 0 END) >= ?")
	havingArgs = append(havingArgs, len(f.tagPrefixes))

	if len(joins) > 0 {
		q += " " + strings.Join(joins, " ") + " "
	}
	if len(whereConds) > 0 {
		q += " WHERE " + strings.Join(whereConds, " AND ")
	}
	q += " GROUP BY n.id"
	if len(havingConds) > 0 {
		q += " HAVING " + strings.Join(havingConds, " AND ")
	}

	args := append(whereArgs, havingArgs...)
	rows, err := s.db.Query(q+" ORDER BY n.created_at DESC", args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var nodes []*models.Node
	for rows.Next() {
		var id string
		if err := rows.Scan(&id); err != nil {
			return nil, err
		}
		if n, err := s.GetByID(id); err == nil {
			nodes = append(nodes, n)
		} else {
			return nil, err
		}
	}
	return nodes, nil
}

func (s *sqliteNodeService) resolveUserRef(tx *sql.Tx, ref string) (string, error) {
	if exists, _ := s.Exists(ref); exists {
		return ref, nil
	}
	return s.ensureUser(tx, ref)
}

func (s *sqliteNodeService) ensureUser(tx *sql.Tx, username string) (string, error) {
	userID, err := s.resolveUserIDByNameTx(tx, username)
	if err != nil {
		return "", err
	}
	if userID != "" {
		return userID, nil
	}

	id := s.generateUniqueID()
	now := time.Now().UTC().Format(time.RFC3339)
	if _, err := tx.Exec("INSERT INTO nodes (id, title, created_at, updated_at) VALUES (?, ?, ?, ?)",
		id, username, now, now); err != nil {
		return "", err
	}
	if _, err := tx.Exec("INSERT INTO tags (node_id, tag) VALUES (?, '_type::user')", id); err != nil {
		return "", err
	}
	return id, nil
}

func (s *sqliteNodeService) resolveNamespaceRef(tx *sql.Tx, ref string) (string, error) {
	if exists, _ := s.Exists(ref); exists {
		return ref, nil
	}
	return s.ensureNamespace(tx, ref)
}

func (s *sqliteNodeService) ensureNamespace(tx *sql.Tx, name string) (string, error) {
	nsID, err := s.resolveNamespaceIDByNameTx(tx, name)
	if err != nil {
		return "", err
	}
	if nsID != "" {
		return nsID, nil
	}

	id := s.generateUniqueID()
	now := time.Now().UTC().Format(time.RFC3339)
	if _, err := tx.Exec("INSERT INTO nodes (id, title, created_at, updated_at) VALUES (?, ?, ?, ?)",
		id, name, now, now); err != nil {
		return "", err
	}
	if _, err := tx.Exec("INSERT INTO tags (node_id, tag) VALUES (?, '_type::namespace')", id); err != nil {
		return "", err
	}
	if _, err := tx.Exec("INSERT INTO rels (from_id, to_id, rel_type) VALUES (?, ?, ?)",
		id, id, models.RelInNamespace); err != nil {
		return "", err
	}
	userID, err := s.resolveUserRef(tx, s.userID)
	if err != nil {
		return "", err
	}
	if _, err := tx.Exec("INSERT INTO rels (from_id, to_id, rel_type) VALUES (?, ?, ?)",
		id, userID, models.RelCreated); err != nil {
		return "", err
	}
	return id, nil
}