rikkei_simple_care/server/internal/db/db.go

package db

import (
	"fmt"
	"log"
	"strings"
	"time"

	"server/internal/models"

	"gorm.io/driver/mysql"
	"gorm.io/gorm"
)

func ConnectDB(host, port, user, password, dbName string) (*gorm.DB, error) {
	// Kết nối ban đầu không chọn database để tạo database nếu chưa tồn tại
	dsnWithoutDB := fmt.Sprintf("%s:%s@tcp(%s:%s)/?charset=utf8mb4&parseTime=True&loc=Local", user, password, host, port)
	tmpDB, err := gorm.Open(mysql.Open(dsnWithoutDB), &gorm.Config{})
	if err != nil {
		return nil, fmt.Errorf("failed to connect to MySQL server: %w", err)
	}

	createDBQuery := fmt.Sprintf("CREATE DATABASE IF NOT EXISTS `%s` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;", dbName)
	if err := tmpDB.Exec(createDBQuery).Error; err != nil {
		return nil, fmt.Errorf("failed to create database: %w", err)
	}

	// Kết nối chính thức vào database chỉ định
	dsnWithDB := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=True&loc=Local", user, password, host, port, dbName)
	gormDB, err := gorm.Open(mysql.Open(dsnWithDB), &gorm.Config{})
	if err != nil {
		return nil, fmt.Errorf("failed to connect to database %s: %w", dbName, err)
	}

	log.Printf("Successfully connected to MySQL database: %s", dbName)
	return gormDB, nil
}

func AutoMigrate(db *gorm.DB) error {
	log.Println("Running AutoMigrate...")
	if err := migrateLegacyWifiTables(db); err != nil {
		return err
	}
	if err := db.AutoMigrate(
		&models.Class{},
		&models.Student{},
		&models.ClassStudent{},
		&models.ClassSchedule{},
		&models.ClassCourse{},
		&models.ClassAllowedApp{},
		&models.AppPoolEntry{},
		&models.WifiPoolEntry{},
		&models.AcceptedWifi{},
		&models.StudentSession{},
		&models.AttendanceResult{},
	); err != nil {
		return err
	}
	if err := purgeLegacyWifiWithoutBSSID(db); err != nil {
		return err
	}
	return migrateLegacyDiscoveredApps(db)
}

// Bảng WiFi cũ dùng ssid_key làm PK; schema mới cần id + bssid_key — GORM không tự đổi PK.
func migrateLegacyWifiTables(db *gorm.DB) error {
	if db.Migrator().HasTable("accepted_wifis") {
		needsReset := !db.Migrator().HasColumn(&models.AcceptedWifi{}, "id") ||
			!db.Migrator().HasColumn(&models.AcceptedWifi{}, "bssid_key")
		if needsReset {
			log.Println("Dropping legacy accepted_wifis (SSID-only schema) — will recreate with BSSID...")
			if err := db.Migrator().DropTable("accepted_wifis"); err != nil {
				return err
			}
		}
	}
	if db.Migrator().HasTable("wifi_pool") {
		needsReset := !db.Migrator().HasColumn(&models.WifiPoolEntry{}, "bssid_key")
		if needsReset {
			log.Println("Dropping legacy wifi_pool (SSID-only schema) — will recreate with BSSID...")
			if err := db.Migrator().DropTable("wifi_pool"); err != nil {
				return err
			}
		}
	}
	return nil
}

func purgeLegacyWifiWithoutBSSID(db *gorm.DB) error {
	_ = db.Where("bssid = '' OR bssid_key = '' OR LENGTH(bssid_key) <> 12").Delete(&models.AcceptedWifi{}).Error
	_ = db.Where("bssid = '' OR bssid_key = '' OR LENGTH(bssid_key) <> 12").Delete(&models.WifiPoolEntry{}).Error
	return nil
}

// Gộp dữ liệu cũ (theo lớp) vào kho global nếu bảng discovered_apps còn tồn tại.
func migrateLegacyDiscoveredApps(db *gorm.DB) error {
	if !db.Migrator().HasTable("discovered_apps") {
		return nil
	}
	type legacyRow struct {
		ProcessName string
		WindowTitle string
		Keyword     string
		HitCount    int
		LastSeenAt  time.Time
		StudentRkID int64
		ClassRkID   int64
	}
	var rows []legacyRow
	if err := db.Table("discovered_apps").Find(&rows).Error; err != nil {
		return nil
	}
	for _, r := range rows {
		pKey := strings.ToLower(strings.TrimSpace(r.ProcessName))
		if pKey == "" {
			continue
		}
		kw := strings.TrimSuffix(pKey, ".exe")
		if r.Keyword != "" {
			kw = r.Keyword
		}
		var existing models.AppPoolEntry
		if err := db.Where("process_key = ?", pKey).First(&existing).Error; err == nil {
			existing.HitCount += r.HitCount
			if r.LastSeenAt.After(existing.LastSeenAt) {
				existing.LastSeenAt = r.LastSeenAt
				existing.WindowTitle = r.WindowTitle
				existing.ProcessName = r.ProcessName
				existing.LastStudentRkID = r.StudentRkID
				existing.LastClassRkID = r.ClassRkID
			}
			_ = db.Save(&existing).Error
		} else {
			_ = db.Create(&models.AppPoolEntry{
				ProcessName:     r.ProcessName,
				ProcessKey:      pKey,
				WindowTitle:     r.WindowTitle,
				Keyword:         kw,
				HitCount:        r.HitCount,
				LastSeenAt:      r.LastSeenAt,
				LastStudentRkID: r.StudentRkID,
				LastClassRkID:   r.ClassRkID,
			}).Error
		}
	}
	return db.Migrator().DropTable("discovered_apps")
}