unbot/db/db.go

package db

import (
	"database/sql"
	"time"

	_ "modernc.org/sqlite"
)

type DB struct {
	*sql.DB
}

type DictEntry struct {
	ID         int64
	Word       string
	Definition string
	AuthorID   string
	CreatedAt  time.Time
	Upvotes    int
	Downvotes  int
}

type XDCount struct {
	UserID string
	Count  int
}

func Open(path string) (*DB, error) {
	d, err := sql.Open("sqlite", path)
	if err != nil {
		return nil, err
	}
	db := &DB{d}
	if err := db.migrate(); err != nil {
		return nil, err
	}
	return db, nil
}

func (db *DB) migrate() error {
	queries := []string{
		`CREATE TABLE IF NOT EXISTS dictionary_entries (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			word TEXT NOT NULL UNIQUE,
			definition TEXT NOT NULL,
			author_id TEXT NOT NULL,
			created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
			upvotes INTEGER DEFAULT 0,
			downvotes INTEGER DEFAULT 0
		)`,
		`CREATE TABLE IF NOT EXISTS xd_counts (
			user_id TEXT PRIMARY KEY,
			count INTEGER NOT NULL DEFAULT 0
		)`,
		`CREATE TABLE IF NOT EXISTS starboard_messages (
			original_message_id TEXT PRIMARY KEY,
			starboard_message_id TEXT NOT NULL
		)`,
		`CREATE TABLE IF NOT EXISTS xd_processed_messages (
			message_id TEXT PRIMARY KEY
		)`,
	}
	for _, q := range queries {
		if _, err := db.Exec(q); err != nil {
			return err
		}
	}
	return nil
}

func (db *DB) AddDefinition(word, definition, authorID string) error {
	_, err := db.Exec(
		`INSERT INTO dictionary_entries (word, definition, author_id) VALUES (?, ?, ?)
		 ON CONFLICT(word) DO UPDATE SET definition = ?, author_id = ?`,
		word, definition, authorID, definition, authorID,
	)
	return err
}

func (db *DB) GetDefinition(word string) (*DictEntry, error) {
	var e DictEntry
	err := db.QueryRow(
		`SELECT id, word, definition, author_id, created_at, upvotes, downvotes
		 FROM dictionary_entries WHERE word = ?`, word,
	).Scan(&e.ID, &e.Word, &e.Definition, &e.AuthorID, &e.CreatedAt, &e.Upvotes, &e.Downvotes)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	return &e, nil
}

func (db *DB) RemoveDefinition(word string) error {
	_, err := db.Exec(`DELETE FROM dictionary_entries WHERE word = ?`, word)
	return err
}

func (db *DB) ListWords() ([]DictEntry, error) {
	rows, err := db.Query(
		`SELECT id, word, definition, author_id, created_at, upvotes, downvotes
		 FROM dictionary_entries ORDER BY word ASC`,
	)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var entries []DictEntry
	for rows.Next() {
		var e DictEntry
		if err := rows.Scan(&e.ID, &e.Word, &e.Definition, &e.AuthorID, &e.CreatedAt, &e.Upvotes, &e.Downvotes); err != nil {
			return nil, err
		}
		entries = append(entries, e)
	}
	return entries, nil
}

func (db *DB) IncrementXD(userID string) error {
	_, err := db.Exec(
		`INSERT INTO xd_counts (user_id, count) VALUES (?, 1)
		 ON CONFLICT(user_id) DO UPDATE SET count = count + 1`,
		userID,
	)
	return err
}

func (db *DB) GetXDCount(userID string) (int, error) {
	var count int
	err := db.QueryRow(`SELECT count FROM xd_counts WHERE user_id = ?`, userID).Scan(&count)
	if err == sql.ErrNoRows {
		return 0, nil
	}
	if err != nil {
		return 0, err
	}
	return count, nil
}

func (db *DB) GetXDLeaderboard(limit int) ([]XDCount, error) {
	rows, err := db.Query(
		`SELECT user_id, count FROM xd_counts ORDER BY count DESC LIMIT ?`, limit,
	)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var entries []XDCount
	for rows.Next() {
		var e XDCount
		if err := rows.Scan(&e.UserID, &e.Count); err != nil {
			return nil, err
		}
		entries = append(entries, e)
	}
	return entries, nil
}

func (db *DB) IsXDMessageProcessed(messageID string) (bool, error) {
	var count int
	err := db.QueryRow(`SELECT COUNT(*) FROM xd_processed_messages WHERE message_id = ?`, messageID).Scan(&count)
	return count > 0, err
}

func (db *DB) MarkXDMessageProcessed(messageID string) error {
	_, err := db.Exec(`INSERT OR IGNORE INTO xd_processed_messages (message_id) VALUES (?)`, messageID)
	return err
}

func (db *DB) BulkMarkXDMessageProcessed(messageIDs []string) error {
	tx, err := db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	stmt, err := tx.Prepare(`INSERT OR IGNORE INTO xd_processed_messages (message_id) VALUES (?)`)
	if err != nil {
		return err
	}
	defer stmt.Close()

	for _, id := range messageIDs {
		if _, err := stmt.Exec(id); err != nil {
			return err
		}
	}
	return tx.Commit()
}

func (db *DB) GetTotalXDCount() (int, error) {
	var total int
	err := db.QueryRow(`SELECT COALESCE(SUM(count), 0) FROM xd_counts`).Scan(&total)
	return total, err
}

func (db *DB) GetXDUserCount() (int, error) {
	var count int
	err := db.QueryRow(`SELECT COUNT(*) FROM xd_counts`).Scan(&count)
	return count, err
}

func (db *DB) GetUserRank(userID string) (int, error) {
	var rank int
	err := db.QueryRow(`
		SELECT COUNT(*) + 1 FROM xd_counts
		WHERE count > (SELECT COALESCE(count, 0) FROM xd_counts WHERE user_id = ?)
	`, userID).Scan(&rank)
	if err != nil {
		return 0, err
	}
	return rank, nil
}

func (db *DB) IsStarboardEntry(originalMsgID string) (bool, error) {
	var count int
	err := db.QueryRow(
		`SELECT COUNT(*) FROM starboard_messages WHERE original_message_id = ?`, originalMsgID,
	).Scan(&count)
	return count > 0, err
}

func (db *DB) AddStarboardEntry(originalMsgID, starboardMsgID string) error {
	_, err := db.Exec(
		`INSERT INTO starboard_messages (original_message_id, starboard_message_id) VALUES (?, ?)`,
		originalMsgID, starboardMsgID,
	)
	return err
}

func (db *DB) GetStarboardMessageID(originalMsgID string) (string, error) {
	var starboardMsgID string
	err := db.QueryRow(
		`SELECT starboard_message_id FROM starboard_messages WHERE original_message_id = ?`, originalMsgID,
	).Scan(&starboardMsgID)
	if err == sql.ErrNoRows {
		return "", nil
	}
	if err != nil {
		return "", err
	}
	return starboardMsgID, nil
}

func (db *DB) RemoveStarboardEntry(originalMsgID string) error {
	_, err := db.Exec(`DELETE FROM starboard_messages WHERE original_message_id = ?`, originalMsgID)
	return err
}