worktimeBot/internal/db/store.go

package db

import (
	"database/sql"
	"log/slog"
	"os"
	"time"

	_ "modernc.org/sqlite"
)

const driverName = "sqlite"

var schemaPath = "/app/db/schema.sql"

type Store struct {
	db *sql.DB
}

type User struct {
	ID             int64
	ChatID         int64
	Timezone       string
	IsActive       bool
	ReportEnabled  bool
	ReportTime     string
	LastReportDate *string
	ExportAccent   string
	Calendar       string
	CreatedAt      int64
	UpdatedAt      int64
}

type WorkType struct {
	ID        int64
	Name      string
	CreatedAt int64
}

type Day struct {
	ID                int64
	UserID            int64
	Date              string
	IsDayOff          bool
	DayOffReason      string
	CurrentWorkTypeID int64
	MinBreakThreshold int64
	CreatedAt         int64
	UpdatedAt         int64
}

type Event struct {
	ID         int64
	UserID     int64
	DayID      int64
	EventType  string
	WorkTypeID *int64
	OccurredAt int64
	Note       string
	CreatedAt  int64
}

func NewStore(path string) (*Store, error) {
	db, err := sql.Open(driverName, path)
	if err != nil {
		return nil, err
	}
	if err := db.Ping(); err != nil {
		return nil, err
	}
	db.Exec("PRAGMA journal_mode=WAL")
	if err := migrate(db); err != nil {
		return nil, err
	}
	if err := execSchema(db, schemaPath); err != nil {
		return nil, err
	}
	return &Store{db: db}, nil
}

func (s *Store) Close() error {
	return s.db.Close()
}

func (s *Store) DB() *sql.DB {
	return s.db
}

func execSchema(db *sql.DB, path string) error {
	data, err := os.ReadFile(path)
	if err != nil {
		return err
	}
	_, err = db.Exec(string(data))
	return err
}

func hasTable(db *sql.DB, name string) bool {
	var exists int
	db.QueryRow("SELECT COUNT(1) FROM sqlite_master WHERE type='table' AND name=?", name).Scan(&exists)
	return exists > 0
}

func hasColumn(db *sql.DB, table, column string) bool {
	rows, err := db.Query("PRAGMA table_info(" + table + ")")
	if err != nil {
		return false
	}
	defer rows.Close()
	for rows.Next() {
		var cid int
		var name, ctype string
		var notnull, pk int
		var dflt sql.NullString
		if err := rows.Scan(&cid, &name, &ctype, &notnull, &dflt, &pk); err != nil {
			return false
		}
		if name == column {
			return true
		}
	}
	return false
}

func migrate(db *sql.DB) error {
	if hasTable(db, "users") {
		if !hasColumn(db, "users", "calendar") {
			slog.Info("migrating: adding calendar column to users")
			db.Exec("ALTER TABLE users ADD COLUMN calendar TEXT NOT NULL DEFAULT 'gregorian' CHECK (calendar IN ('gregorian', 'jalali', 'hijri'))")
		}
	}
	if hasTable(db, "events_old") || hasTable(db, "events") && !hasColumn(db, "events", "user_id") {
		slog.Info("migrating from old schema to new schema")
		if err := migrateFromOldSchema(db); err != nil {
			return err
		}
	}
	return nil
}

func migrateFromOldSchema(db *sql.DB) error {
	// Rename old tables so CREATE TABLE IF NOT EXISTS can create new ones
	for _, t := range []string{"events", "days_off", "settings"} {
		if hasTable(db, t) {
			db.Exec("ALTER TABLE " + t + " RENAME TO " + t + "_old")
		}
	}

	if err := execSchema(db, schemaPath); err != nil {
		return err
	}

	return migrateFromOldSchemaData(db)
}

func migrateFromOldSchemaData(db *sql.DB) error {
	migrateUsers(db)
	migrateDaysOff(db)
	migrateEvents(db)
	migrateSettings(db)

	for _, t := range []string{"days_off_old", "settings_old", "events_old"} {
		if hasTable(db, t) {
			db.Exec("DROP TABLE IF EXISTS " + t)
		}
	}

	slog.Info("migration complete")
	return nil
}

func migrateUsers(db *sql.DB) {
	if !hasTable(db, "events_old") {
		return
	}
	rows, err := db.Query("SELECT DISTINCT chat_id FROM events_old")
	if err != nil {
		slog.Warn("migration: no events table to migrate users from")
		return
	}
	defer rows.Close()

	for rows.Next() {
		var chatID int64
		if err := rows.Scan(&chatID); err != nil {
			continue
		}
		db.Exec(
			"INSERT OR IGNORE INTO users (chat_id, timezone) VALUES (?, 'Asia/Tehran')",
			chatID,
		)
	}
}

func migrateDaysOff(db *sql.DB) {
	if !hasTable(db, "days_off_old") {
		return
	}
	rows, err := db.Query("SELECT chat_id, date, reason FROM days_off_old")
	if err != nil {
		return
	}
	defer rows.Close()

	for rows.Next() {
		var chatID int64
		var date, reason string
		if err := rows.Scan(&chatID, &date, &reason); err != nil {
			continue
		}
		var userID int64
		err := db.QueryRow("SELECT id FROM users WHERE chat_id=?", chatID).Scan(&userID)
		if err != nil {
			continue
		}
		db.Exec(
			"INSERT INTO days (user_id, date, is_day_off, day_off_reason) VALUES (?, ?, 1, ?) ON CONFLICT(user_id, date) DO UPDATE SET is_day_off=1, day_off_reason=excluded.day_off_reason",
			userID, date, reason,
		)
	}
}

func migrateEvents(db *sql.DB) {
	if !hasTable(db, "events_old") {
		return
	}
	rows, err := db.Query("SELECT chat_id, event_type, occurred_at, note, created_at FROM events_old ORDER BY occurred_at ASC")
	if err != nil {
		return
	}
	defer rows.Close()

	for rows.Next() {
		var chatID int64
		var eventType, note string
		var occurredAt, createdAt int64
		if err := rows.Scan(&chatID, &eventType, &occurredAt, &note, &createdAt); err != nil {
			continue
		}

		var userID int64
		err := db.QueryRow("SELECT id FROM users WHERE chat_id=?", chatID).Scan(&userID)
		if err != nil {
			continue
		}

		date := time.Unix(occurredAt, 0).UTC().Format("2006-01-02")
		db.Exec("INSERT OR IGNORE INTO days (user_id, date) VALUES (?, ?)", userID, date)

		var dayID int64
		db.QueryRow("SELECT id FROM days WHERE user_id=? AND date=?", userID, date).Scan(&dayID)

		db.Exec(
			"INSERT INTO events (user_id, day_id, event_type, occurred_at, note, created_at) VALUES (?, ?, ?, ?, ?, ?)",
			userID, dayID, eventType, occurredAt, note, createdAt,
		)
	}
}

func migrateSettings(db *sql.DB) {
	if !hasTable(db, "settings_old") {
		return
	}
	rows, err := db.Query("SELECT chat_id, key, value FROM settings_old")
	if err != nil {
		return
	}
	defer rows.Close()

	for rows.Next() {
		var chatID int64
		var key, value string
		if err := rows.Scan(&chatID, &key, &value); err != nil {
			continue
		}
		if key == "remote_flag" {
			var userID int64
			err := db.QueryRow("SELECT id FROM users WHERE chat_id=?", chatID).Scan(&userID)
			if err != nil {
				continue
			}
			wt := int64(1)
			if value == "remote" {
				wt = 2
			}
			db.Exec("UPDATE days SET current_work_type_id=? WHERE user_id=? AND date=?", wt, userID, time.Now().UTC().Format("2006-01-02"))
		}
	}
}

func (s *Store) GetOrCreateUser(chatID int64) (*User, error) {
	_, err := s.db.Exec(
		"INSERT OR IGNORE INTO users (chat_id) VALUES (?)",
		chatID,
	)
	if err != nil {
		return nil, err
	}
	return s.GetUserByChatID(chatID)
}

func (s *Store) GetUserByChatID(chatID int64) (*User, error) {
	row := s.db.QueryRow(`
		SELECT id, chat_id, timezone, is_active, report_enabled, report_time,
		       last_report_date, export_accent, calendar, created_at, updated_at
		FROM users WHERE chat_id = ?
	`, chatID)
	var u User
	var lastReport sql.NullString
	err := row.Scan(
		&u.ID, &u.ChatID, &u.Timezone, &u.IsActive, &u.ReportEnabled, &u.ReportTime,
		&lastReport, &u.ExportAccent, &u.Calendar, &u.CreatedAt, &u.UpdatedAt,
	)
	if err != nil {
		return nil, err
	}
	if lastReport.Valid {
		u.LastReportDate = &lastReport.String
	}
	return &u, nil
}

func (s *Store) UpdateUser(user *User) error {
	_, err := s.db.Exec(
		"UPDATE users SET timezone=?, is_active=?, report_enabled=?, report_time=?, last_report_date=?, export_accent=?, calendar=?, updated_at=unixepoch() WHERE id=?",
		user.Timezone, user.IsActive, user.ReportEnabled, user.ReportTime, user.LastReportDate, user.ExportAccent, user.Calendar, user.ID,
	)
	return err
}

func (s *Store) MarkReportSent(userID int64, date string) error {
	_, err := s.db.Exec(
		"UPDATE users SET last_report_date=?, updated_at=unixepoch() WHERE id=?",
		date, userID,
	)
	return err
}

func (s *Store) GetAllUsers() ([]User, error) {
	rows, err := s.db.Query(`
		SELECT id, chat_id, timezone, is_active, report_enabled, report_time,
		       last_report_date, export_accent, calendar, created_at, updated_at
		FROM users
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var users []User
	for rows.Next() {
		var u User
		var lastReport sql.NullString
		if err := rows.Scan(
			&u.ID, &u.ChatID, &u.Timezone, &u.IsActive, &u.ReportEnabled, &u.ReportTime,
			&lastReport, &u.ExportAccent, &u.Calendar, &u.CreatedAt, &u.UpdatedAt,
		); err != nil {
			return nil, err
		}
		if lastReport.Valid {
			u.LastReportDate = &lastReport.String
		}
		users = append(users, u)
	}
	return users, rows.Err()
}

func (s *Store) GetWorkTypes() ([]WorkType, error) {
	rows, err := s.db.Query("SELECT id, name, created_at FROM work_types ORDER BY id ASC")
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var wts []WorkType
	for rows.Next() {
		var wt WorkType
		if err := rows.Scan(&wt.ID, &wt.Name, &wt.CreatedAt); err != nil {
			return nil, err
		}
		wts = append(wts, wt)
	}
	return wts, rows.Err()
}

func (s *Store) GetWorkType(id int64) (*WorkType, error) {
	var wt WorkType
	err := s.db.QueryRow("SELECT id, name, created_at FROM work_types WHERE id=?", id).Scan(&wt.ID, &wt.Name, &wt.CreatedAt)
	if err != nil {
		return nil, err
	}
	return &wt, nil
}

func (s *Store) GetOrCreateDay(userID int64, date string) (*Day, error) {
	_, err := s.db.Exec(
		"INSERT OR IGNORE INTO days (user_id, date) VALUES (?, ?)",
		userID, date,
	)
	if err != nil {
		return nil, err
	}
	return s.GetDay(userID, date)
}

func (s *Store) GetDay(userID int64, date string) (*Day, error) {
	row := s.db.QueryRow(`
		SELECT id, user_id, date, is_day_off, day_off_reason,
		       current_work_type_id, min_break_threshold, created_at, updated_at
		FROM days WHERE user_id=? AND date=?
	`, userID, date)
	var d Day
	err := row.Scan(
		&d.ID, &d.UserID, &d.Date, &d.IsDayOff, &d.DayOffReason,
		&d.CurrentWorkTypeID, &d.MinBreakThreshold, &d.CreatedAt, &d.UpdatedAt,
	)
	if err != nil {
		return nil, err
	}
	return &d, nil
}

func (s *Store) UpdateDay(day *Day) error {
	_, err := s.db.Exec(
		"UPDATE days SET is_day_off=?, day_off_reason=?, current_work_type_id=?, min_break_threshold=?, updated_at=unixepoch() WHERE id=?",
		day.IsDayOff, day.DayOffReason, day.CurrentWorkTypeID, day.MinBreakThreshold, day.ID,
	)
	return err
}

func (s *Store) SetDayOff(userID int64, date string, reason string) error {
	_, err := s.GetOrCreateDay(userID, date)
	if err != nil {
		return err
	}
	_, err = s.db.Exec(
		"UPDATE days SET is_day_off=1, day_off_reason=?, updated_at=unixepoch() WHERE user_id=? AND date=?",
		reason, userID, date,
	)
	return err
}

func (s *Store) RemoveDayOff(userID int64, date string) error {
	_, err := s.db.Exec(
		"UPDATE days SET is_day_off=0, day_off_reason='', updated_at=unixepoch() WHERE user_id=? AND date=?",
		userID, date,
	)
	return err
}

func (s *Store) IsDayOff(userID int64, date string) (bool, error) {
	var isOff bool
	err := s.db.QueryRow(
		"SELECT is_day_off FROM days WHERE user_id=? AND date=?",
		userID, date,
	).Scan(&isOff)
	if err == sql.ErrNoRows {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	return isOff, nil
}

func (s *Store) GetUserDaysInRange(userID int64, startDate, endDate string) ([]Day, error) {
	rows, err := s.db.Query(`
		SELECT id, user_id, date, is_day_off, day_off_reason,
		       current_work_type_id, min_break_threshold, created_at, updated_at
		FROM days
		WHERE user_id=? AND date >= ? AND date <= ?
		ORDER BY date ASC
	`, userID, startDate, endDate)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var days []Day
	for rows.Next() {
		var d Day
		if err := rows.Scan(
			&d.ID, &d.UserID, &d.Date, &d.IsDayOff, &d.DayOffReason,
			&d.CurrentWorkTypeID, &d.MinBreakThreshold, &d.CreatedAt, &d.UpdatedAt,
		); err != nil {
			return nil, err
		}
		days = append(days, d)
	}
	return days, rows.Err()
}

func (s *Store) CreateEvent(userID int64, dayID int64, eventType string, workTypeID *int64, occurredAt int64, note string) error {
	_, err := s.db.Exec(
		"INSERT INTO events (user_id, day_id, event_type, work_type_id, occurred_at, note) VALUES (?, ?, ?, ?, ?, ?)",
		userID, dayID, eventType, workTypeID, occurredAt, note,
	)
	return err
}

func (s *Store) GetLastEvent(userID int64) (*Event, error) {
	row, err := s.db.Query(`
		SELECT id, user_id, day_id, event_type, work_type_id, occurred_at, note, created_at
		FROM events
		WHERE user_id = ?
		ORDER BY occurred_at DESC
		LIMIT 1
	`, userID)
	if err != nil {
		return nil, err
	}
	defer row.Close()
	if !row.Next() {
		return nil, sql.ErrNoRows
	}
	var e Event
	var wt sql.NullInt64
	if err := row.Scan(&e.ID, &e.UserID, &e.DayID, &e.EventType, &wt, &e.OccurredAt, &e.Note, &e.CreatedAt); err != nil {
		return nil, err
	}
	if wt.Valid {
		e.WorkTypeID = &wt.Int64
	}
	return &e, nil
}

func (s *Store) EventsForDayByDate(userID int64, date string) ([]Event, error) {
	rows, err := s.db.Query(`
		SELECT e.id, e.user_id, e.day_id, e.event_type, e.work_type_id, e.occurred_at, e.note, e.created_at
		FROM events e
		JOIN days d ON d.id = e.day_id
		WHERE d.user_id = ? AND d.date = ?
		ORDER BY e.occurred_at ASC
	`, userID, date)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var events []Event
	for rows.Next() {
		var e Event
		var wt sql.NullInt64
		if err := rows.Scan(&e.ID, &e.UserID, &e.DayID, &e.EventType, &wt, &e.OccurredAt, &e.Note, &e.CreatedAt); err != nil {
			return nil, err
		}
		if wt.Valid {
			e.WorkTypeID = &wt.Int64
		}
		events = append(events, e)
	}
	return events, rows.Err()
}

func (s *Store) EventsForDayByDayID(dayID int64) ([]Event, error) {
	rows, err := s.db.Query(`
		SELECT id, user_id, day_id, event_type, work_type_id, occurred_at, note, created_at
		FROM events
		WHERE day_id = ?
		ORDER BY occurred_at ASC
	`, dayID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var events []Event
	for rows.Next() {
		var e Event
		var wt sql.NullInt64
		if err := rows.Scan(&e.ID, &e.UserID, &e.DayID, &e.EventType, &wt, &e.OccurredAt, &e.Note, &e.CreatedAt); err != nil {
			return nil, err
		}
		if wt.Valid {
			e.WorkTypeID = &wt.Int64
		}
		events = append(events, e)
	}
	return events, rows.Err()
}