fix: 优化注册逻辑

backend/src/main/java/com/yfd/platform/system/controller/SmsVerifyCodeController.java

@@ -1,5 +1,6 @@
 package com.yfd.platform.system.controller;
 
+import cn.hutool.core.util.StrUtil;
 import cn.hutool.json.JSONUtil;
 import cn.hutool.jwt.JWTUtil;
 import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
@@ -258,103 +259,129 @@ public class SmsVerifyCodeController {
         }
     }
 
+    /**
+     * 给用户分配流域/电站数据权限 + 默认角色。
+     * <p>
+     * 逻辑说明：
+     * 1. 收集前端传入的电站编码，按流域分组
+     * 2. 如果用户选了某流域的所有电站 → 保存流域权限（RVCD）
+     * 3. 如果用户只选了某流域的部分电站 → 逐个保存电站权限（STATION）
+     * 4. 数据库中查不到的孤立电站编码 → 直接保存电站权限
+     * <p>
+     * 示例：
+     * - 选了 A流域全部6个电站 + B流域全部3个电站 + C流域1个电站
+     *   → 两条 RVCD（A、B）+ 一条 STATION（C的电站）
+     */
     private boolean addDefaultRole(String userId, SmsVerifyCodeRequest smsVerifyCodeRequest) {
         if (userId == null || userId.isEmpty()) {
             return false;
         }
 
+        // 1. 收集前端传入的电站编码
         String stationCode = smsVerifyCodeRequest.getStationCode();
-        String rvcdCode = smsVerifyCodeRequest.getRvcdCode();
-
         Set<String> selectedStationCodes = new HashSet<>();
         if (StringUtils.isNotEmpty(stationCode)) {
             selectedStationCodes.addAll(Arrays.asList(stationCode.split(",")));
         }
 
-        Set<String> selectedBasinCodes = new HashSet<>();
+        // 2. 如果没有选择任何电站
-        if (StringUtils.isNotEmpty(rvcdCode)) {
+        if (selectedStationCodes.isEmpty()) {
-            selectedBasinCodes.addAll(Arrays.asList(rvcdCode.split(",")));
+            // 如果只传了流域，直接保存流域权限
-        }
+            String rvcdCode = smsVerifyCodeRequest.getRvcdCode();
-
+            if (StringUtils.isNotEmpty(rvcdCode)) {
-//        Set<String> addedStationCodes = new HashSet<>();
+                for (String basinCode : rvcdCode.split(",")) {
-
+                    if (StringUtils.isNotEmpty(basinCode)) {
-        for (String basinCode : selectedBasinCodes) {
+                        addDataScope(userId, "RVCD", basinCode);
-            if (StringUtils.isEmpty(basinCode)) {
+                    }
-                continue;
-            }
-
-            List<SdEngInfoBH> allStationsInBasin = engInfoBHService.lambdaQuery()
-                    .eq(SdEngInfoBH::getRvcd, basinCode)
-                    .list();
-
-            if (allStationsInBasin == null || allStationsInBasin.isEmpty()) {
-                SysUserDataScope scope = new SysUserDataScope();
-                scope.setUserId(userId);
-                scope.setOrgType("RVCD");
-                scope.setOrgId(basinCode);
-                scope.setStatus(1);
-                scope.setPermissionType("READ");
-                sysUserDataScopeService.addDataScope(scope);
-                continue;
-            }
-
-            Set<String> allStationCodesInBasin = allStationsInBasin.stream()
-                    .map(SdEngInfoBH::getStcd)
-                    .collect(Collectors.toSet());
-
-            boolean allStationsSelected = allStationCodesInBasin.containsAll(selectedStationCodes)
-                    && selectedStationCodes.containsAll(allStationCodesInBasin);
-
-            if (allStationsSelected) {
-                SysUserDataScope scope = new SysUserDataScope();
-                scope.setUserId(userId);
-                scope.setOrgType("RVCD");
-                scope.setOrgId(basinCode);
-                scope.setStatus(1);
-                scope.setPermissionType("READ");
-                sysUserDataScopeService.addDataScope(scope);
-//                addedStationCodes.add(basinCode);
-            } else {
-                Set<String> stationsInBasinAndSelected = allStationCodesInBasin.stream()
-                        .filter(selectedStationCodes::contains)
-                        .collect(Collectors.toSet());
-
-                for (String stationCd : stationsInBasinAndSelected) {
-                    SysUserDataScope scope = new SysUserDataScope();
-                    scope.setUserId(userId);
-                    scope.setOrgType("STATION");
-                    scope.setOrgId(stationCd);
-                    scope.setStatus(1);
-                    scope.setPermissionType("READ");
-                    sysUserDataScopeService.addDataScope(scope);
-//                    addedStationCodes.add(stationCd);
                 }
             }
+            assignRoleToUser(userId);
+            return true;
         }
 
-//        Set<String> standaloneStations = selectedStationCodes.stream()
+        // 3. 批量查询所有选中电站的工程信息，获取每个电站所属的流域
-//                .filter(code -> !addedStationCodes.contains(code))
+        List<SdEngInfoBH> selectedStations = engInfoBHService.lambdaQuery()
-//                .collect(Collectors.toSet());
+                .in(SdEngInfoBH::getStcd, selectedStationCodes)
-//
+                .list();
-//        for (String stationCd : standaloneStations) {
+
-//            if (StringUtils.isEmpty(stationCd)) {
+        // 已通过电站权限处理过的电站编码
-//                continue;
+        Set<String> processedStationCodes = new HashSet<>();
-//            }
+
-//            SysUserDataScope scope = new SysUserDataScope();
+        if (!selectedStations.isEmpty()) {
-//            scope.setUserId(userId);
+            // 4. 按流域分组：{rvcd -> Set<stcd>}
-//            scope.setOrgType("STATION");
+            Map<String, Set<String>> basinToStations = new LinkedHashMap<>();
-//            scope.setOrgId(stationCd);
+            for (SdEngInfoBH station : selectedStations) {
-//            scope.setStatus(1);
+                String reachcd = station.getReachcd();
-//            scope.setPermissionType("READ");
+                String stcd = station.getStcd();
-//            sysUserDataScopeService.addDataScope(scope);
+                if (StrUtil.isBlank(reachcd)) {
-//        }
+                    // 无流域归属的电站，直接保存电站权限
+                    addDataScope(userId, "STATION", stcd);
+                    processedStationCodes.add(stcd);
+                    continue;
+                }
+                basinToStations.computeIfAbsent(reachcd, k -> new HashSet<>()).add(stcd);
+            }
+
+            // 5. 对每个流域，判断用户是否选中了该流域的全部电站
+            for (Map.Entry<String, Set<String>> entry : basinToStations.entrySet()) {
+                String basinCode = entry.getKey();
+                Set<String> selectedInThisBasin = entry.getValue();
+
+                // 查询该流域下的所有电站
+                List<SdEngInfoBH> allStationsInBasin = engInfoBHService.lambdaQuery()
+                        .eq(SdEngInfoBH::getReachcd, basinCode)
+                        .list();
+                Set<String> allStationCodesInBasin = allStationsInBasin.stream()
+                        .map(SdEngInfoBH::getStcd)
+                        .collect(Collectors.toSet());
+
+                if (allStationCodesInBasin.equals(selectedInThisBasin)) {
+                    // 全选 → 保存流域权限
+                    addDataScope(userId, "RVCD", basinCode);
+                } else {
+                    // 部分选 → 逐个保存电站权限
+                    for (String stcd : selectedInThisBasin) {
+                        addDataScope(userId, "STATION", stcd);
+                    }
+                }
+                processedStationCodes.addAll(selectedInThisBasin);
+            }
+        }
+
+        // 6. 处理数据库中查不到的电站编码（孤立的电站），直接保存电站权限
+        for (String stcd : selectedStationCodes) {
+            if (!processedStationCodes.contains(stcd)) {
+                addDataScope(userId, "STATION", stcd);
+            }
+        }
+
+        // 7. 分配默认角色
+        assignRoleToUser(userId);
+        return true;
+    }
+
+    /**
+     * 保存一条数据权限
+     */
+    private void addDataScope(String userId, String orgType, String orgId) {
+        SysUserDataScope scope = new SysUserDataScope();
+        scope.setUserId(userId);
+        scope.setOrgType(orgType);
+        scope.setOrgId(orgId);
+        scope.setStatus(1);
+        scope.setPermissionType("READ");
+        sysUserDataScopeService.addDataScope(scope);
+    }
+
+    /**
+     * 给用户分配默认角色
+     */
+    private void assignRoleToUser(String userId) {
         SysUser user = new SysUser();
         SysRole aDefault = roleMapper.selectOne(new LambdaQueryWrapper<SysRole>().eq(SysRole::getCustom1, "default"));
         String roleids = aDefault == null ? null : aDefault.getId();
         user.setId(userId);
         userService.updateUserRoles(user, roleids);
-        // 加上角色权限
-        return true;
     }
 
     /**