CSI-ES-2324/Projs/PD1/internal/server/datastore.go

package server

import (
	"PD1/internal/protocol"
	"database/sql"
	"log"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

type DataStore struct {
	db *sql.DB
}

func OpenDB() DataStore {
	db, err := sql.Open("sqlite3", "server.db")
	if err != nil {
		log.Fatalln("Error opening db file")
	}
	return DataStore{db: db}
}

func (ds DataStore) CreateTables() error {
	// Create users table
	_, err := ds.db.Exec(`CREATE TABLE IF NOT EXISTS users (
        UID TEXT PRIMARY KEY,
        userCert BLOB
    )`)
	if err != nil {
		return err
	}

	// Create messages table
	_, err = ds.db.Exec(`CREATE TABLE IF NOT EXISTS messages (
        fromUID TEXT,
        toUID TEXT,
        timestamp TIMESTAMP,
        content BLOB,
        PRIMARY KEY (toUID, fromUID, timestamp),
        FOREIGN KEY(fromUID) REFERENCES users(UID),
        FOREIGN KEY(toUID) REFERENCES users(UID)
    )`)
	if err != nil {
		return err
	}

	return nil
}

func (ds DataStore) GetMessage(toUID string, position int) protocol.ServerMessagePacket {

	var serverMessage protocol.ServerMessagePacket
	query := `
        SELECT fromUID, toUID, content, timestamp
        FROM messages
        WHERE toUID = ?
        AND status = 0
        ORDER BY timestamp
        LIMIT 1 OFFSET ?
    `
	// Execute the query
	row := ds.db.QueryRow(query, toUID, position)
	err := row.Scan(&serverMessage.FromUID, &serverMessage.ToUID, &serverMessage.Content, &serverMessage.Timestamp)
	if err != nil {
		log.Panicln("Could not map DB query to ServerMessage")
	}
	return serverMessage

}

func (ds DataStore) MarkMessageInQueueAsRead(toUID string, position int) error {
	query := `
        UPDATE messages
        SET status = 1
        WHERE toUID = ? AND status = 0
        ORDER BY timestamp
        LIMIT 1 OFFSET ?
    `

	// Execute the SQL statement
	_, err := ds.db.Exec(query, toUID, position)
	if err != nil {
		return err
	}

	return nil
}

func (ds DataStore) GetAllMessages(toUID string) []protocol.Packet {
	var messagePackets []protocol.Packet

	// Query to retrieve all messages from the user's queue
	query := `
        SELECT fromUID, toUID, content, timestamp
        FROM messages
        WHERE toUID = ?
        AND status = 0
        ORDER BY timestamp
    `

	// Execute the query
	rows, err := ds.db.Query(query, toUID)
	if err != nil {
		log.Panicln("Failed to execute query:", err)
	}
	defer rows.Close()

	// Iterate through the result set and scan each row into a ServerMessage struct
	for rows.Next() {
		var fromUID string
		var toUID string
		var content []byte
		var timestamp time.Time
		if err := rows.Scan(&fromUID, &toUID, &content, &timestamp); err != nil {
			log.Panicln("Failed to scan row:", err)
		}
		message := protocol.NewServerMessagePacket(fromUID, toUID, content, timestamp)
		messagePackets = append(messagePackets, message)
	}
	if err := rows.Err(); err != nil {
		log.Panicln("Error when getting user's messages")
	}

	return messagePackets
}

func (ds DataStore) AddMessageToQueue(uid string, message protocol.SubmitMessagePacket) {
	query := `
        INSERT INTO messages (fromUID, toUID, content, timestamp, status)
        VALUES (?, ?, ?, ?, 0)
    `

	// Execute the SQL statement
	currentTime := time.Now()
	_, err := ds.db.Exec(query, uid, message.ToUID, message.Content, currentTime)
	if err != nil {
		log.Panicln("Error adding message to database")
	}
}

func (ds DataStore) GetUserCertificate(uid string) protocol.Packet {
	query := `
        SELECT userCert
        FROM users
        WHERE UID = ?
    `

	// Execute the SQL query
	var userCert []byte
	err := ds.db.QueryRow(query, uid).Scan(&userCert)
	if err != nil {
		log.Panicln("Error getting user certificate from the database")
	}
	return protocol.NewSendUserCertPacket(uid, userCert)
}

func userExists(db *sql.DB, uid string) bool {
    // Prepare the SQL statement for checking if a user exists
    query := `
        SELECT COUNT(*)
        FROM users
        WHERE UID = ?
    `

    var count int
    // Execute the SQL query
    err := db.QueryRow(query, uid).Scan(&count)
    if err != nil {
        log.Panicln("Error checking if user exists")
    }

    // If count is greater than 0, the user exists
    return count > 0
}