material_manage_system/internal/repository/user.go

package repository

import (
	"github.com/wmms/wmms/internal/model"
	"github.com/wmms/wmms/pkg/database"
	"golang.org/x/crypto/bcrypt"
)

type UserRepository struct{}

func NewUserRepository() *UserRepository {
	return &UserRepository{}
}

func (r *UserRepository) GetByUsername(username string) (*model.User, error) {
	query := `SELECT id, username, password, role, description, created_at FROM users WHERE username = ?`

	var user model.User
	err := database.DB.QueryRow(query, username).Scan(
		&user.ID,
		&user.Username,
		&user.Password,
		&user.Role,
		&user.Description,
		&user.CreatedAt,
	)
	if err != nil {
		return nil, err
	}
	return &user, nil
}

func (r *UserRepository) GetByID(id int64) (*model.User, error) {
	query := `SELECT id, username, password, role, description, created_at FROM users WHERE id = ?`

	var user model.User
	err := database.DB.QueryRow(query, id).Scan(
		&user.ID,
		&user.Username,
		&user.Password,
		&user.Role,
		&user.Description,
		&user.CreatedAt,
	)
	if err != nil {
		return nil, err
	}
	return &user, nil
}

func (r *UserRepository) Create(user *model.User) error {
	query := `INSERT INTO users (username, password, role, description) VALUES (?, ?, ?, ?)`

	result, err := database.DB.Exec(query, user.Username, user.Password, user.Role, user.Description)
	if err != nil {
		return err
	}

	id, err := result.LastInsertId()
	if err != nil {
		return err
	}

	user.ID = id
	return nil
}

func (r *UserRepository) List(search string, limit, offset int) ([]model.User, int64, error) {
	countQuery := `SELECT COUNT(*) FROM users`
	var countArgs []interface{}
	if search != "" {
		countQuery += ` WHERE username LIKE ?`
		searchPattern := "%" + search + "%"
		countArgs = append(countArgs, searchPattern)
	}
	var total int64
	err := database.DB.QueryRow(countQuery, countArgs...).Scan(&total)
	if err != nil {
		return nil, 0, err
	}

	query := `SELECT id, username, password, role, description, created_at FROM users`
	var args []interface{}
	if search != "" {
		query += ` WHERE username LIKE ?`
		searchPattern := "%" + search + "%"
		args = append(args, searchPattern)
	}
	query += ` ORDER BY id LIMIT ? OFFSET ?`
	args = append(args, limit, offset)
	rows, err := database.DB.Query(query, args...)
	if err != nil {
		return nil, 0, err
	}
	defer rows.Close()

	var users []model.User
	for rows.Next() {
		var user model.User
		err := rows.Scan(
			&user.ID,
			&user.Username,
			&user.Password,
			&user.Role,
			&user.Description,
			&user.CreatedAt,
		)
		if err != nil {
			return nil, 0, err
		}
		users = append(users, user)
	}

	return users, total, nil
}

func (r *UserRepository) Delete(id int64) error {
	query := `DELETE FROM users WHERE id = ?`
	_, err := database.DB.Exec(query, id)
	return err
}

func (r *UserRepository) Update(user *model.User) error {
	query := `UPDATE users SET role = ?, description = ? WHERE id = ?`
	_, err := database.DB.Exec(query, user.Role, user.Description, user.ID)
	return err
}

func (r *UserRepository) UpdateWithPassword(user *model.User) error {
	query := `UPDATE users SET password = ?, role = ?, description = ? WHERE id = ?`
	_, err := database.DB.Exec(query, user.Password, user.Role, user.Description, user.ID)
	return err
}

// UserExists checks if a user exists by username
func (r *UserRepository) UserExists(username string) (bool, error) {
	query := `SELECT COUNT(*) FROM users WHERE username = ?`
	var count int
	err := database.DB.QueryRow(query, username).Scan(&count)
	if err != nil {
		return false, err
	}
	return count > 0, nil
}

// InitUserTable creates the users table if it doesn't exist
func InitUserTable() error {
	query := `
	CREATE TABLE IF NOT EXISTS users (
		id BIGINT AUTO_INCREMENT PRIMARY KEY,
		username VARCHAR(50) NOT NULL UNIQUE,
		password VARCHAR(255) NOT NULL,
		role ENUM('admin','user') NOT NULL DEFAULT 'user',
		description VARCHAR(500) DEFAULT '',
		created_at DATETIME DEFAULT CURRENT_TIMESTAMP
	)`
	_, err := database.DB.Exec(query)
	return err
}

// CreateDefaultUsers creates default admin and user if no users exist
func CreateDefaultUsers() error {
	repo := NewUserRepository()

	exists, err := repo.UserExists("admin")
	if err != nil {
		return err
	}
	if exists {
		return nil
	}

	// Default passwords will be hashed in CreateUser
	admin := &model.User{
		Username:    "admin",
		Password:    "", // Will be set by service
		Role:        "admin",
		Description: "系统管理员",
	}

	// Use bcrypt to hash the default passwords
	adminHash, err := bcrypt.GenerateFromPassword([]byte("admin123"), bcrypt.DefaultCost)
	if err != nil {
		return err
	}
	admin.Password = string(adminHash)

	if err := repo.Create(admin); err != nil {
		return err
	}

	userHash, err := bcrypt.GenerateFromPassword([]byte("user123"), bcrypt.DefaultCost)
	if err != nil {
		return err
	}
	user := &model.User{
		Username:    "user",
		Password:    string(userHash),
		Role:        "user",
		Description: "普通用户",
	}

	return repo.Create(user)
}