yx-dataset-server/app/bll/impl/internal/b_user.go

package internal

import (
	"context"
	"strings"

	"github.com/gogf/gf/v2/util/guid"
	"yx-dataset-server/app/errors"
	"yx-dataset-server/app/model"
	"yx-dataset-server/app/schema"
	"yx-dataset-server/library/redis"
	"yx-dataset-server/library/utils"
)

// NewUser 创建User
func NewUser(
	mUser model.IUser,
	mRelation model.IDatasetRelation,
	mTrans model.ITrans,
	mMenu model.IMenu,
	mRoleMenu model.IRoleMenu,
	mDataset model.IDataset,
	mOrganization model.IOrganization,
	mRole model.IRole,
) *User {
	return &User{
		UserModel:     mUser,
		relationModel: mRelation,
		transModel:    mTrans,
		menuModel:     mMenu,
		roleMenuModel: mRoleMenu,
		datasetModel:  mDataset,
		orgModel:      mOrganization,
		roleModel:     mRole,
	}
}

// User 创建User对象
type User struct {
	UserModel     model.IUser
	relationModel model.IDatasetRelation
	transModel    model.ITrans
	menuModel     model.IMenu
	roleMenuModel model.IRoleMenu
	datasetModel  model.IDataset
	orgModel      model.IOrganization
	roleModel     model.IRole
}

// classifyDatasets 按 type 分桶填充
func classifyDatasets(user *schema.User, datasets schema.Datasets) {
	user.Datasets = datasets
	user.PublicDatasets = make(schema.Datasets, 0)
	user.OrgDatasets = make(schema.Datasets, 0)
	user.PersonalDatasets = make(schema.Datasets, 0)
	for _, d := range datasets {
		switch d.Type {
		case schema.DatasetTypePublic:
			user.PublicDatasets = append(user.PublicDatasets, d)
		case schema.DatasetTypeOrg:
			user.OrgDatasets = append(user.OrgDatasets, d)
		case schema.DatasetTypePersonal:
			user.PersonalDatasets = append(user.PersonalDatasets, d)
		}
	}
}

// loadUserDatasets 加载指定用户的可访问知识库并按类型分桶填充
// 特殊规则：系统管理员（root 另行处理于 getRootInfo）——
//
//	可查看全部公共/共享知识库（type=1） + 本人的个人知识库（type=3）
func (a *User) loadUserDatasets(ctx context.Context, user *schema.User) error {
	// 系统管理员：全部公共KB + 本人的个人KB
	if user.RoleId != "" {
		role, err := a.roleModel.Get(ctx, user.RoleId)
		if err != nil {
			return err
		}
		if role != nil && role.Code == RoleCodeSystemAdmin {
			return a.loadAdminDatasets(ctx, user)
		}
	}
	// 其他角色：仅限 DatasetRelation 里显式授权的
	rel, err := a.relationModel.Query(ctx, schema.DatasetRelationQueryParam{BizId: user.RecordID})
	if err != nil {
		return err
	}
	if len(rel.Data) == 0 {
		classifyDatasets(user, schema.Datasets{})
		return nil
	}
	datasets, err := a.datasetModel.Query(ctx, schema.DatasetQueryParam{RecordIds: rel.Data.ToDatasetIds()})
	if err != nil {
		return err
	}
	classifyDatasets(user, datasets.Data)
	return nil
}

// loadAdminDatasets 为 root / 系统管理员加载：全部公共KB + 本人个人KB
func (a *User) loadAdminDatasets(ctx context.Context, user *schema.User) error {
	public, err := a.datasetModel.Query(ctx, schema.DatasetQueryParam{Type: schema.DatasetTypePublic})
	if err != nil {
		return err
	}

	all := make(schema.Datasets, 0, len(public.Data))
	all = append(all, public.Data...)

	// 本人个人知识库（通过 DatasetRelation 绑定）
	rel, err := a.relationModel.Query(ctx, schema.DatasetRelationQueryParam{
		BizId: user.RecordID,
		Type:  schema.DatasetTypePersonal,
	})
	if err != nil {
		return err
	}
	if len(rel.Data) > 0 {
		personal, err := a.datasetModel.Query(ctx, schema.DatasetQueryParam{RecordIds: rel.Data.ToDatasetIds()})
		if err != nil {
			return err
		}
		all = append(all, personal.Data...)
	}
	classifyDatasets(user, all)
	return nil
}

// Query 查询数据
func (a *User) Query(ctx context.Context, params schema.UserQueryParam, opts ...schema.UserOptions) (*schema.UserQueryResult, error) {
	if isSys, err := IsSystemAdmin(ctx, a.UserModel, a.roleModel); !isSys {
		if err != nil {
			return nil, err
		}
		current, err := a.UserModel.Get(ctx, GetUserID(ctx))
		if err != nil {
			return nil, err
		}
		params.OrgId = current.OrgId
	}

	result, err := a.UserModel.Query(ctx, params, opts...)
	if err != nil {
		return nil, err
	}
	result.Data.FillCreator(result.Data)
	for _, v := range result.Data {
		v.Used, _ = redis.GetRedisClient().Get(ctx, "chart:user:"+v.RecordID).Int()
		v.Unused = v.ChartNum - v.Used
	}

	roles, err := a.roleModel.Query(ctx, schema.RoleQueryParam{})
	if err != nil {
		return nil, err
	}
	result.Data.FillRoleName(roles.Data)

	orgs, err := a.orgModel.Query(ctx, schema.OrganizationQueryParam{})
	if err != nil {
		return nil, err
	}
	result.Data.FillOrgName(orgs.Data)

	return result, nil
}

// Get 查询指定数据
func (a *User) Get(ctx context.Context, recordID string, opts ...schema.UserOptions) (*schema.User, error) {

	item, err := a.UserModel.Get(ctx, recordID, opts...)
	if err != nil {
		return nil, err
	} else if item == nil {
		return nil, errors.ErrNotFound
	}

	// 查询用户菜单
	roleMenu, err := a.roleMenuModel.Query(ctx, schema.RoleMenuQueryParam{RoleId: item.RoleId})
	if err != nil {
		return nil, err
	}
	menus, err := a.menuModel.Query(ctx, schema.MenuQueryParam{RecordIDs: roleMenu.Data.ToMenuIds()})
	if err != nil {
		return nil, err
	}
	item.Menu = menus.Data.ToTrees()

	role, err := a.roleModel.Get(ctx, item.RoleId)
	if err != nil {
		return nil, err
	}
	if role != nil {
		item.RoleCode = role.Code
		item.RoleName = role.Name
	}

	org, err := a.orgModel.Get(ctx, item.OrgId)
	if err != nil {
		return nil, err
	}
	if org != nil {
		item.OrgName = org.Name
	}

	if err := a.loadUserDatasets(ctx, item); err != nil {
		return nil, err
	}

	item.Password = "******"
	item.Used, _ = redis.GetRedisClient().Get(ctx, "chart:user:"+item.RecordID).Int()
	item.Unused = item.ChartNum - item.Used
	return item, nil
}

// GetCurrentUser 查询当前用户
func (a *User) GetCurrentUser(ctx context.Context, recordID string, opts ...schema.UserOptions) (*schema.User, error) {
	if CheckIsRootUser(ctx) {
		return a.getRootInfo(ctx)
	}
	item, err := a.UserModel.Get(ctx, recordID, opts...)
	if err != nil {
		return nil, err
	} else if item == nil {
		return nil, errors.ErrNotFound
	}

	roleMenu, err := a.roleMenuModel.Query(ctx, schema.RoleMenuQueryParam{RoleId: item.RoleId})
	if err != nil {
		return nil, err
	}
	menus, err := a.menuModel.Query(ctx, schema.MenuQueryParam{RecordIDs: roleMenu.Data.ToMenuIds()})
	if err != nil {
		return nil, err
	}
	item.Menu = menus.Data.ToTrees()

	role, err := a.roleModel.Get(ctx, item.RoleId)
	if err != nil {
		return nil, err
	}
	if role != nil {
		item.RoleCode = role.Code
		item.RoleName = role.Name
	}

	org, err := a.orgModel.Get(ctx, item.OrgId)
	if err != nil {
		return nil, err
	}
	if org != nil {
		item.OrgName = org.Name
	}

	if err := a.loadUserDatasets(ctx, item); err != nil {
		return nil, err
	}

	item.Password = "******"
	item.Used, _ = redis.GetRedisClient().Get(ctx, "chart:user:"+item.RecordID).Int()
	item.Unused = item.ChartNum - item.Used
	return item, nil
}

// GetCurrentUserDatasets 获取当前用户有权限的知识库（按组织分组）
func (a *User) GetCurrentUserDatasets(ctx context.Context, opts ...schema.UserOptions) (schema.Organizations, error) {
	rel, err := a.relationModel.Query(ctx, schema.DatasetRelationQueryParam{BizId: GetUserID(ctx)})
	if err != nil {
		return nil, err
	}
	dataset, err := a.datasetModel.Query(ctx, schema.DatasetQueryParam{RecordIds: rel.Data.ToDatasetIds()})
	if err != nil {
		return nil, err
	}
	orgs, err := a.orgModel.Query(ctx, schema.OrganizationQueryParam{RecordIds: dataset.Data.ToOrgIds()})
	if err != nil {
		return nil, err
	}
	orgs.Data.FillDataset(dataset.Data)
	return orgs.Data, nil
}

func (a *User) getUpdate(ctx context.Context, recordID string) (*schema.User, error) {
	return a.Get(ctx, recordID)
}

// Create 创建数据
// item.Datasets 为分配给该用户的知识库列表（schema.Dataset，仅 RecordID 必填）
// - 类型从知识库本身读取
// - 仅允许分配 共享(1) / 企业(2) 两类；个人(3) 不由创建/编辑用户接口管理
func (a *User) Create(ctx context.Context, item schema.User) error {
	item.RecordID = guid.S()
	item.Status = true
	item.Password = utils.SHA1HashString(strings.ToUpper(item.Password))

	return ExecTrans(ctx, a.transModel, func(ctx context.Context) error {
		if err := a.UserModel.Create(ctx, item); err != nil {
			return err
		}
		return a.syncAssignedDatasets(ctx, item.RecordID, item.Datasets)
	})
}

// syncAssignedDatasets 同步“分配给用户”的共享/企业 KB 关系（不影响个人 KB）
func (a *User) syncAssignedDatasets(ctx context.Context, userId string, datasets schema.Datasets) error {
	// 仅清理 type=1/2 的分配，保留 type=3 个人 KB
	if err := a.relationModel.DeleteByBizIdAndType(ctx, userId, schema.DatasetTypePublic); err != nil {
		return err
	}
	if err := a.relationModel.DeleteByBizIdAndType(ctx, userId, schema.DatasetTypeOrg); err != nil {
		return err
	}
	for _, d := range datasets {
		if d.RecordID == "" {
			continue
		}
		ds, err := a.datasetModel.Get(ctx, d.RecordID)
		if err != nil {
			return err
		}
		if ds == nil {
			return errors.New400Response("知识库不存在: " + d.RecordID)
		}
		if ds.Type == schema.DatasetTypePersonal {
			// 个人知识库不通过“分配”流转
			continue
		}
		if err := a.relationModel.Create(ctx, schema.DatasetRelation{
			RecordID:  guid.S(),
			DatasetId: ds.RecordID,
			BizId:     userId,
			Type:      ds.Type,
			CreatorId: GetUserID(ctx),
		}); err != nil {
			return err
		}
	}
	return nil
}

// Update 更新数据
// 知识库权限变化后是否刷新对话助手，改为前端通过独立接口触发确认：
//
//	GET  /web/v1/chat_assistants/permission/check  -> 检查是否一致
//	POST /web/v1/chat_assistants/permission/sync   -> 用户确认后同步
func (a *User) Update(ctx context.Context, recordID string, item schema.User) error {
	oldItem, err := a.UserModel.Get(ctx, recordID)
	if err != nil {
		return err
	} else if oldItem == nil {
		return errors.ErrNotFound
	}
	if item.Password != "" {
		item.Password = utils.SHA1HashString(strings.ToUpper(item.Password))
	}
	return ExecTrans(ctx, a.transModel, func(ctx context.Context) error {
		if err := a.syncAssignedDatasets(ctx, recordID, item.Datasets); err != nil {
			return err
		}
		return a.UserModel.Update(ctx, recordID, item)
	})
}

// Delete 删除数据
func (a *User) Delete(ctx context.Context, recordID string) error {
	oldItem, err := a.UserModel.Get(ctx, recordID)
	if err != nil {
		return err
	} else if oldItem == nil {
		return errors.ErrNotFound
	}
	return ExecTrans(ctx, a.transModel, func(ctx context.Context) error {
		// 清理该用户的全部授权（含个人KB的映射）
		if err := a.relationModel.DeleteByBizId(ctx, recordID); err != nil {
			return err
		}
		return a.UserModel.Delete(ctx, recordID)
	})
}

// UpdateStatus 更新状态
func (a *User) UpdateStatus(ctx context.Context, recordID string, status bool) error {
	oldItem, err := a.UserModel.Get(ctx, recordID)
	if err != nil {
		return err
	} else if oldItem == nil {
		return errors.ErrNotFound
	}
	return a.UserModel.UpdateStatus(ctx, recordID, status)
}

func (a *User) getRootInfo(ctx context.Context) (*schema.User, error) {
	user := GetRootUser()
	user.Password = "******"
	menu, err := a.menuModel.Query(ctx, schema.MenuQueryParam{})
	if err != nil {
		return nil, err
	}
	user.Menu = menu.Data.ToTrees()
	// root：全部公共/共享知识库 + root 本人的个人知识库
	if err := a.loadAdminDatasets(ctx, user); err != nil {
		return nil, err
	}
	return user, nil
}