X-Agents/server/internal/repository/agent_repo.go

package repository

import (
	"x-agents/server/internal/model"

	"gorm.io/gorm"
)

type AgentRepository struct {
	db *gorm.DB
}

func NewAgentRepository(db *gorm.DB) *AgentRepository {
	return &AgentRepository{db: db}
}

func (r *AgentRepository) Create(agent *model.Agent) error {
	return r.db.Create(agent).Error
}

func (r *AgentRepository) FindByID(id string) (*model.Agent, error) {
	var agent model.Agent
	err := r.db.First(&agent, "id = ?", id).Error
	if err != nil {
		return nil, err
	}
	return &agent, nil
}

func (r *AgentRepository) FindByOwnerID(ownerID string) ([]model.Agent, error) {
	var agents []model.Agent
	err := r.db.Where("owner_id = ?", ownerID).Find(&agents).Error
	return agents, err
}

func (r *AgentRepository) FindAll() ([]model.Agent, error) {
	var agents []model.Agent
	err := r.db.Where("is_active = ?", true).Find(&agents).Error
	return agents, err
}

// FindByIDs 根据多个ID查询
func (r *AgentRepository) FindByIDs(ids []string) ([]model.Agent, error) {
	var agents []model.Agent
	err := r.db.Where("id IN ?", ids).Find(&agents).Error
	return agents, err
}

func (r *AgentRepository) Update(agent *model.Agent) error {
	return r.db.Save(agent).Error
}

func (r *AgentRepository) Delete(id string) error {
	return r.db.Delete(&model.Agent{}, "id = ?", id).Error
}

// AgentSkill 相关方法

func (r *AgentRepository) CreateAgentSkill(as *model.AgentSkill) error {
	return r.db.Create(as).Error
}

func (r *AgentRepository) FindSkillsByAgentID(agentID string) ([]model.AgentSkill, error) {
	var agentSkills []model.AgentSkill
	err := r.db.Where("agent_id = ?", agentID).Find(&agentSkills).Error
	return agentSkills, err
}

func (r *AgentRepository) DeleteAgentSkills(agentID string) error {
	return r.db.Delete(&model.AgentSkill{}, "agent_id = ?", agentID).Error
}

// AgentKnowledgeBase 相关方法

func (r *AgentRepository) CreateAgentKnowledgeBase(akb *model.AgentKnowledgeBase) error {
	return r.db.Create(akb).Error
}

func (r *AgentRepository) FindKnowledgeBasesByAgentID(agentID string) ([]model.AgentKnowledgeBase, error) {
	var agentKBs []model.AgentKnowledgeBase
	err := r.db.Where("agent_id = ?", agentID).Find(&agentKBs).Error
	return agentKBs, err
}

func (r *AgentRepository) DeleteAgentKnowledgeBases(agentID string) error {
	return r.db.Delete(&model.AgentKnowledgeBase{}, "agent_id = ?", agentID).Error
}

// AgentMemory 相关方法

func (r *AgentRepository) CreateMemory(memory *model.AgentMemory) error {
	return r.db.Create(memory).Error
}

func (r *AgentRepository) FindMemoriesByAgentID(agentID string, limit int) ([]model.AgentMemory, error) {
	var memories []model.AgentMemory
	err := r.db.Where("agent_id = ?", agentID).Order("importance DESC, created_at DESC").Limit(limit).Find(&memories).Error
	return memories, err
}

func (r *AgentRepository) FindMemoriesByUserID(agentID, userID string, limit int) ([]model.AgentMemory, error) {
	var memories []model.AgentMemory
	err := r.db.Where("agent_id = ? AND user_id = ?", agentID, userID).Order("importance DESC, created_at DESC").Limit(limit).Find(&memories).Error
	return memories, err
}

func (r *AgentRepository) DeleteMemory(id string) error {
	return r.db.Delete(&model.AgentMemory{}, "id = ?", id).Error
}

// FindMemories 通用查询记忆，支持分类和类型过滤
func (r *AgentRepository) FindMemories(agentID, userID, category, memoryType string, limit, offset int) ([]model.AgentMemory, int64, error) {
	var memories []model.AgentMemory
	query := r.db.Model(&model.AgentMemory{})

	if agentID != "" {
		query = query.Where("agent_id = ?", agentID)
	}
	if userID != "" {
		query = query.Where("user_id = ?", userID)
	}
	if category != "" {
		query = query.Where("category = ?", category)
	}
	if memoryType != "" {
		query = query.Where("memory_type = ?", memoryType)
	}

	// 统计总数
	var total int64
	countQuery := query
	if err := countQuery.Count(&total).Error; err != nil {
		return nil, 0, err
	}

	// 置顶的记忆优先，然后按重要性降序，最后按创建时间降序
	err := query.Order("is_pinned DESC, importance DESC, created_at DESC").Limit(limit).Offset(offset).Find(&memories).Error
	return memories, total, err
}

// SearchMemories 搜索记忆
func (r *AgentRepository) SearchMemories(agentID, userID, keyword, tags, category, memoryType string, minScore, limit, offset int) ([]model.AgentMemory, int64, error) {
	var memories []model.AgentMemory
	query := r.db.Model(&model.AgentMemory{})

	if agentID != "" {
		query = query.Where("agent_id = ?", agentID)
	}
	if userID != "" {
		query = query.Where("user_id = ?", userID)
	}
	if keyword != "" {
		keyword = "%" + keyword + "%"
		query = query.Where("content LIKE ? OR keywords LIKE ? OR tags LIKE ?", keyword, keyword, keyword)
	}
	if category != "" {
		query = query.Where("category = ?", category)
	}
	if memoryType != "" {
		query = query.Where("memory_type = ?", memoryType)
	}
	if minScore > 0 {
		query = query.Where("importance >= ?", minScore)
	}

	// 统计总数
	var total int64
	countQuery := query
	if err := countQuery.Count(&total).Error; err != nil {
		return nil, 0, err
	}

	err := query.Order("is_pinned DESC, importance DESC, created_at DESC").Limit(limit).Offset(offset).Find(&memories).Error
	return memories, total, err
}

// FindMemoryByID 根据ID查询记忆
func (r *AgentRepository) FindMemoryByID(id string) (*model.AgentMemory, error) {
	var memory model.AgentMemory
	err := r.db.Where("id = ?", id).First(&memory).Error
	if err != nil {
		return nil, err
	}
	return &memory, nil
}

// UpdateMemory 更新记忆
func (r *AgentRepository) UpdateMemory(memory *model.AgentMemory) error {
	return r.db.Save(memory).Error
}

// FindMemoryCategories 获取记忆分类列表
func (r *AgentRepository) FindMemoryCategories(agentID, userID string) ([]string, error) {
	var categories []string
	query := r.db.Model(&model.AgentMemory{}).Distinct("category")

	if agentID != "" {
		query = query.Where("agent_id = ?", agentID)
	}
	if userID != "" {
		query = query.Where("user_id = ?", userID)
	}

	err := query.Pluck("category", &categories).Error
	return categories, err
}

// AgentTeam 相关方法

func (r *AgentRepository) CreateAgentTeam(team *model.AgentTeam) error {
	return r.db.Create(team).Error
}

func (r *AgentRepository) FindTeamMembers(supervisorAgentID string) ([]model.AgentTeam, error) {
	var teams []model.AgentTeam
	err := r.db.Where("supervisor_agent_id = ?", supervisorAgentID).Find(&teams).Error
	return teams, err
}

func (r *AgentRepository) FindByMemberAgentID(memberAgentID string) ([]model.AgentTeam, error) {
	var teams []model.AgentTeam
	err := r.db.Where("member_agent_id = ?", memberAgentID).Find(&teams).Error
	return teams, err
}

func (r *AgentRepository) DeleteAgentTeam(id string) error {
	return r.db.Delete(&model.AgentTeam{}, "id = ?", id).Error
}

func (r *AgentRepository) DeleteTeamMembers(supervisorAgentID string) error {
	return r.db.Delete(&model.AgentTeam{}, "supervisor_agent_id = ?", supervisorAgentID).Error
}

// AgentTask 相关方法

func (r *AgentRepository) CreateTask(task *model.AgentTask) error {
	return r.db.Create(task).Error
}

func (r *AgentRepository) UpdateTask(task *model.AgentTask) error {
	return r.db.Save(task).Error
}

func (r *AgentRepository) FindTasksByAgentID(agentID string, limit int) ([]model.AgentTask, error) {
	var tasks []model.AgentTask
	err := r.db.Where("agent_id = ?", agentID).Order("created_at DESC").Limit(limit).Find(&tasks).Error
	return tasks, err
}

func (r *AgentRepository) FindTaskByID(id string) (*model.AgentTask, error) {
	var task model.AgentTask
	err := r.db.First(&task, "id = ?", id).Error
	if err != nil {
		return nil, err
	}
	return &task, nil
}