backend/database.go

package main

import (
	"database/sql"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"super-frpc/postLog"
	"time"

	_ "modernc.org/sqlite"
)

var db *sql.DB
var logsDB *sql.DB

func GetLogsDatabase() *sql.DB {
	return logsDB
}

type User struct {
	UserID    int
	Username  string
	Passwd    string
	Type      string
	CreatedAt string
}

type FrpcInstance struct {
	ID          int
	UserID      int
	Name        string
	BootAtStart bool
	RunUser     string
	ConfigPath  string
	CreatedAt   time.Time
	CreatedBy   string
}

func InitDatabase(dbPath_data string, dbPath_log string) error {
	InitUserDatabase(dbPath_data)
	InitFrpcDatabase(dbPath_data)
	postLog.InitLogsDatabase(dbPath_log)
	return nil
}

func CloseDatabase() error {
	if db != nil {
		if err := db.Close(); err != nil {
			return err
		}
	}
	if logsDB != nil {
		if err := logsDB.Close(); err != nil {
			return err
		}
	}
	return nil
}

func isValidInput(input string) bool {
	invalidChars := []string{"'", "\"", ";", "--", "/*", "*/", "xp_", "sp_", "EXEC", "EXECUTE", "DROP", "INSERT", "UPDATE", "DELETE", "SELECT"}
	lowerInput := strings.ToLower(input)
	for _, chars := range invalidChars {
		if strings.Contains(lowerInput, chars) {
			return false
		}
	}
	return true
}

func AddUser(username, passwd, userType string) (int, error) { // New user registration with specified type
	if !isValidInput(username) || !isValidInput(passwd) || !isValidInput(userType) {
		return 0, errors.New("invalid input: contains illegal characters")
	}

	if !isValidPassword(passwd) {
		return 0, errors.New("password does not meet complexity requirements")
	}

	hashedPasswd, err := hashPassword(passwd)
	if err != nil {
		return 0, fmt.Errorf("failed to hash password: %w", err)
	}

	result, err := db.Exec("INSERT INTO userLogin (username, passwd, type) VALUES (?, ?, ?)",
		username, hashedPasswd, userType)
	if err != nil {
		if strings.Contains(err.Error(), "UNIQUE constraint failed") {
			return 0, errors.New("username already exists")
		}
		return 0, fmt.Errorf("failed to insert user: %w", err)
	}

	lastID, err := result.LastInsertId()
	if err != nil {
		return 0, fmt.Errorf("failed to get last insert id: %w", err)
	}

	var count int
	err = db.QueryRow("SELECT COUNT(*) FROM userLogin WHERE userID = ?", lastID).Scan(&count)
	if err != nil {
		return 0, fmt.Errorf("failed to verify user insertion: %w", err)
	}

	if count == 0 {
		return 0, errors.New("user insertion verification failed")
	}

	return int(lastID), nil
}

func RemoveUser(userID int) error {
	if !isValidInput(strconv.Itoa(userID)) {
		return errors.New("invalid input: contains illegal characters")
	}

	result, err := db.Exec("DELETE FROM userLogin WHERE userID = ?", userID)
	if err != nil {
		return fmt.Errorf("failed to delete user: %w", err)
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("failed to get rows affected: %w", err)
	}

	if rowsAffected == 0 {
		return errors.New("user not found")
	}

	return nil
}

func GetUserByUsername(username string) (*User, error) {
	if !isValidInput(username) {
		return nil, errors.New("invalid input: contains illegal characters")
	}

	var user User
	err := db.QueryRow("SELECT userID, username, passwd, type, createdAt FROM userLogin WHERE username = ?", username).
		Scan(&user.UserID, &user.Username, &user.Passwd, &user.Type, &user.CreatedAt)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, errors.New("user not found")
		}
		return nil, fmt.Errorf("failed to query user: %w", err)
	}

	return &user, nil
}

func GetUserByID(userID int) (*User, error) {
	var user User
	err := db.QueryRow("SELECT userID, username, passwd, type, createdAt FROM userLogin WHERE userID = ?", userID).
		Scan(&user.UserID, &user.Username, &user.Passwd, &user.Type, &user.CreatedAt)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, errors.New("user not found")
		}
		return nil, fmt.Errorf("failed to query user: %w", err)
	}

	return &user, nil
}

func UpdateUserPassword(userID int, newPasswd string) error {
	if !isValidInput(newPasswd) {
		return errors.New("invalid input: contains illegal characters")
	}

	if !isValidPassword(newPasswd) {
		return errors.New("password does not meet complexity requirements")
	}

	hashedPasswd, err := hashPassword(newPasswd)
	if err != nil {
		return fmt.Errorf("failed to hash password: %w", err)
	}

	result, err := db.Exec("UPDATE userLogin SET passwd = ? WHERE userID = ?", hashedPasswd, userID)
	if err != nil {
		return fmt.Errorf("failed to update password: %w", err)
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("failed to get rows affected: %w", err)
	}

	if rowsAffected == 0 {
		return errors.New("user not found")
	}

	return nil
}

func UpdateUserType(userID int, newType string) error {
	validTypes := map[string]bool{
		"superuser": true,
		"admin":     true,
		"visitor":   true,
	}

	if !validTypes[newType] {
		return errors.New("invalid user type")
	}

	result, err := db.Exec("UPDATE userLogin SET type = ? WHERE userID = ?", newType, userID)
	if err != nil {
		return fmt.Errorf("failed to update user type: %w", err)
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("failed to get rows affected: %w", err)
	}

	if rowsAffected == 0 {
		return errors.New("user not found")
	}

	return nil
}

func GetNextAvailableUserID() (int, error) {
	var maxID int
	err := db.QueryRow("SELECT COALESCE(MAX(userID), 0) FROM userLogin").Scan(&maxID)
	if err != nil {
		return 0, fmt.Errorf("failed to get max userID: %w", err)
	}
	return maxID + 1, nil
}

func GetAllUsers() ([]User, error) {
	rows, err := db.Query("SELECT userID, username, type FROM userLogin")
	if err != nil {
		return nil, fmt.Errorf("failed to query users: %w", err)
	}
	defer rows.Close()

	var users []User
	for rows.Next() {
		var user User
		if err := rows.Scan(&user.UserID, &user.Username, &user.Type); err != nil {
			return nil, fmt.Errorf("failed to scan user: %w", err)
		}
		users = append(users, user)
	}

	if err = rows.Err(); err != nil {
		return nil, fmt.Errorf("rows error: %w", err)
	}

	return users, nil
}

func InitFrpcDatabase(dbPath string) error {
	var err error
	frpcDB, err = sql.Open("sqlite", dbPath)
	if err != nil {
		return fmt.Errorf("failed to open frpc database: %w", err)
	}

	if err = frpcDB.Ping(); err != nil {
		return fmt.Errorf("failed to ping frpc database: %w", err)
	}

	createTableSQL := `
	CREATE TABLE IF NOT EXISTS frpcInstances (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		userID INTEGER NOT NULL,
		name TEXT NOT NULL,
		bootAtStart INTEGER NOT NULL DEFAULT 0,
		runUser TEXT NOT NULL DEFAULT 'root',
		configPath TEXT NOT NULL,
		createdAt TEXT NOT NULL,
		UNIQUE(userID, name)
	);
	`
	_, err = frpcDB.Exec(createTableSQL)
	if err != nil {
		return fmt.Errorf("failed to create frpcInstances table: %w", err)
	}

	return nil
}

func InitUserDatabase(dbPath string) error {
	var err error
	db, err = sql.Open("sqlite", dbPath)
	if err != nil {
		return fmt.Errorf("failed to open database: %w", err)
	}

	if err = db.Ping(); err != nil {
		return fmt.Errorf("failed to ping database: %w", err)
	}

	createTableSQL := `
		CREATE TABLE IF NOT EXISTS userLogin (
			userID INTEGER PRIMARY KEY AUTOINCREMENT,
			username TEXT UNIQUE NOT NULL,
			passwd TEXT NOT NULL,
			type TEXT NOT NULL DEFAULT 'visitor',
			createdAt TEXT NOT NULL DEFAULT (datetime('now'))
		);
		`
	_, err = db.Exec(createTableSQL)
	if err != nil {
		return fmt.Errorf("failed to create table: %w", err)
	}

	return nil
}

func DBQueryUsers() ([]User, error) { // List all users
	rows, err := db.Query("SELECT userID, username, type, createdAt FROM userLogin")
	if err != nil {
		return nil, fmt.Errorf("failed to query users: %w", err)
	}
	defer rows.Close()
	var users []User
	for rows.Next() {
		var user User
		if err := rows.Scan(&user.UserID, &user.Username, &user.Type, &user.CreatedAt); err != nil {
			return nil, fmt.Errorf("failed to scan user: %w", err)
		}
		users = append(users, user)
	}

	if err = rows.Err(); err != nil {
		return nil, fmt.Errorf("rows error: %w", err)
	}

	return users, nil
}

func DBQuerySpecificUser(userID int) (User, error) { // Query user by ID
	var user User
	err := db.QueryRow("SELECT userID, username, type, createdAt FROM userLogin WHERE userID = ?", userID).Scan(&user.UserID, &user.Username, &user.Type, &user.CreatedAt)
	if err != nil {
		return user, fmt.Errorf("failed to query user: %w", err)
	}
	return user, nil
}

func DBAddFrpcInstance(instance FrpcInstance) error {
	_, err := frpcDB.Exec("INSERT INTO frpcInstances (userID, name, bootAtStart, runUser, configPath, createdAt) VALUES (?, ?, ?, ?, ?, ?)",
		instance.UserID, instance.Name, instance.BootAtStart, instance.RunUser, instance.ConfigPath, time.Now().Format(time.RFC3339))
	if err != nil {
		return fmt.Errorf("failed to insert frpc instance: %w", err)
	}
	return nil
}

func DBQueryFrpcInstanceByID(instanceID int) (FrpcInstance, error) {
	var instance FrpcInstance
	var createdAtStr string
	err := frpcDB.QueryRow("SELECT id, userID, name, bootAtStart, runUser, configPath, createdAt FROM frpcInstances WHERE id = ?", instanceID).Scan(
		&instance.ID, &instance.UserID, &instance.Name, &instance.BootAtStart, &instance.RunUser, &instance.ConfigPath, &createdAtStr)
	if err != nil {
		return instance, fmt.Errorf("failed to query frpc instance: %w", err)
	}
	instance.CreatedAt, _ = time.Parse(time.RFC3339, createdAtStr)
	return instance, nil
}

func DBQueryFrpcInstance(userID int, instanceName string) (FrpcInstance, error) {
	var instance FrpcInstance
	var createdAtStr string
	err := frpcDB.QueryRow("SELECT id, userID, name, bootAtStart, runUser, configPath, createdAt FROM frpcInstances WHERE userID = ? AND name = ?", userID, instanceName).Scan(
		&instance.ID, &instance.UserID, &instance.Name, &instance.BootAtStart, &instance.RunUser, &instance.ConfigPath, &createdAtStr)
	if err != nil {
		return instance, fmt.Errorf("failed to query frpc instance: %w", err)
	}
	instance.CreatedAt, _ = time.Parse(time.RFC3339, createdAtStr)
	return instance, nil
}

func DBRemoveFrpcInstanceByID(instanceID int) error {
	_, err := frpcDB.Exec("DELETE FROM frpcInstances WHERE id = ?", instanceID)
	if err != nil {
		return fmt.Errorf("failed to delete frpc instance: %w", err)
	}
	return nil
}

func DBUpdateFrpcInstance(instance FrpcInstance) error {
	_, err := frpcDB.Exec("UPDATE frpcInstances SET bootAtStart = ?, runUser = ?, configPath = ? WHERE id = ?",
		instance.BootAtStart, instance.RunUser, instance.ConfigPath, instance.ID)
	if err != nil {
		return fmt.Errorf("failed to update frpc instance: %w", err)
	}
	return nil
}

func DBListFrpcInstances() ([]FrpcInstance, error) {
	rows, err := frpcDB.Query(`
		SELECT fi.id, fi.userID, fi.name, fi.bootAtStart, fi.runUser, fi.configPath, fi.createdAt, u.username
		FROM frpcInstances fi
		JOIN userLogin u ON fi.userID = u.userID
	`)
	if err != nil {
		return nil, fmt.Errorf("failed to query all frpc instances: %w", err)
	}
	defer rows.Close()

	var instances []FrpcInstance
	for rows.Next() {
		var instance FrpcInstance
		var createdAtStr string
		if err := rows.Scan(&instance.ID, &instance.UserID, &instance.Name, &instance.BootAtStart, &instance.RunUser, &instance.ConfigPath, &createdAtStr, &instance.CreatedBy); err != nil {
			return nil, fmt.Errorf("failed to scan frpc instance: %w", err)
		}
		instance.CreatedAt, _ = time.Parse(time.RFC3339, createdAtStr)
		instances = append(instances, instance)
	}

	if err = rows.Err(); err != nil {
		return nil, fmt.Errorf("rows error: %w", err)
	}

	return instances, nil
}

func GetServiceNameByInstanceID(instanceID int) (string, error) {
	instance, err := DBQueryFrpcInstanceByID(instanceID)
	if err != nil {
		return "", fmt.Errorf("failed to query frpc instance: %w", err)
	}

	user, err := GetUserByID(instance.UserID)
	if err != nil {
		return "", fmt.Errorf("failed to get user info: %w", err)
	}

	serviceName := fmt.Sprintf("superfrpc_%s_%s", user.Username, instance.Name)
	postLog.Debug(fmt.Sprintf("[GetServiceNameByInstanceID] instanceID: %d, serviceName: %s", instanceID, serviceName))
	return serviceName, nil
}