x

internal/wechatdb/repository/chatroom.go

@@ -0,0 +1,184 @@
+package repository
+
+import (
+	"context"
+	"fmt"
+	"sort"
+	"strings"
+
+	"github.com/sjzar/chatlog/internal/model"
+)
+
+// initChatRoomCache 初始化群聊缓存
+func (r *Repository) initChatRoomCache(ctx context.Context) error {
+	// 加载所有群聊到缓存
+	chatRooms, err := r.ds.GetChatRooms(ctx, "", 0, 0)
+	if err != nil {
+		return fmt.Errorf("加载群聊失败: %w", err)
+	}
+
+	chatRoomMap := make(map[string]*model.ChatRoom)
+	remarkToChatRoom := make(map[string]*model.ChatRoom)
+	nickNameToChatRoom := make(map[string]*model.ChatRoom)
+	chatRoomList := make([]string, 0)
+	chatRoomRemark := make([]string, 0)
+	chatRoomNickName := make([]string, 0)
+
+	for _, chatRoom := range chatRooms {
+		// 补充群聊信息（从联系人中获取 Remark 和 NickName）
+		r.enrichChatRoom(chatRoom)
+		chatRoomMap[chatRoom.Name] = chatRoom
+		chatRoomList = append(chatRoomList, chatRoom.Name)
+		if chatRoom.Remark != "" {
+			remarkToChatRoom[chatRoom.Remark] = chatRoom
+			chatRoomRemark = append(chatRoomRemark, chatRoom.Remark)
+		}
+		if chatRoom.NickName != "" {
+			nickNameToChatRoom[chatRoom.NickName] = chatRoom
+			chatRoomNickName = append(chatRoomNickName, chatRoom.NickName)
+		}
+	}
+
+	for _, contact := range r.chatRoomInContact {
+		if _, ok := chatRoomMap[contact.UserName]; !ok {
+			chatRoom := &model.ChatRoom{
+				Name:     contact.UserName,
+				Remark:   contact.Remark,
+				NickName: contact.NickName,
+			}
+			chatRoomMap[contact.UserName] = chatRoom
+			chatRoomList = append(chatRoomList, contact.UserName)
+			if contact.Remark != "" {
+				remarkToChatRoom[contact.Remark] = chatRoom
+				chatRoomRemark = append(chatRoomRemark, contact.Remark)
+			}
+			if contact.NickName != "" {
+				nickNameToChatRoom[contact.NickName] = chatRoom
+				chatRoomNickName = append(chatRoomNickName, contact.NickName)
+			}
+		}
+	}
+	sort.Strings(chatRoomList)
+	sort.Strings(chatRoomRemark)
+	sort.Strings(chatRoomNickName)
+
+	r.chatRoomCache = chatRoomMap
+	r.chatRoomList = chatRoomList
+	r.remarkToChatRoom = remarkToChatRoom
+	r.nickNameToChatRoom = nickNameToChatRoom
+	return nil
+}
+
+func (r *Repository) GetChatRooms(ctx context.Context, key string, limit, offset int) ([]*model.ChatRoom, error) {
+
+	ret := make([]*model.ChatRoom, 0)
+	if key != "" {
+		ret = r.findChatRooms(key)
+		if len(ret) == 0 {
+			return nil, fmt.Errorf("未找到群聊: %s", key)
+		}
+
+		if limit > 0 {
+			end := offset + limit
+			if end > len(ret) {
+				end = len(ret)
+			}
+			if offset >= len(ret) {
+				return []*model.ChatRoom{}, nil
+			}
+			return ret[offset:end], nil
+		}
+	} else {
+		list := r.chatRoomList
+		if limit > 0 {
+			end := offset + limit
+			if end > len(list) {
+				end = len(list)
+			}
+			if offset >= len(list) {
+				return []*model.ChatRoom{}, nil
+			}
+			list = list[offset:end]
+		}
+		for _, name := range list {
+			ret = append(ret, r.chatRoomCache[name])
+		}
+	}
+
+	return ret, nil
+}
+
+func (r *Repository) GetChatRoom(ctx context.Context, key string) (*model.ChatRoom, error) {
+	chatRoom := r.findChatRoom(key)
+	if chatRoom == nil {
+		return nil, fmt.Errorf("未找到群聊: %s", key)
+	}
+	return chatRoom, nil
+}
+
+// enrichChatRoom 从联系人信息中补充群聊信息
+func (r *Repository) enrichChatRoom(chatRoom *model.ChatRoom) {
+	if contact, ok := r.contactCache[chatRoom.Name]; ok {
+		chatRoom.Remark = contact.Remark
+		chatRoom.NickName = contact.NickName
+	}
+}
+
+func (r *Repository) findChatRoom(key string) *model.ChatRoom {
+	if chatRoom, ok := r.chatRoomCache[key]; ok {
+		return chatRoom
+	}
+	if chatRoom, ok := r.remarkToChatRoom[key]; ok {
+		return chatRoom
+	}
+	if chatRoom, ok := r.nickNameToChatRoom[key]; ok {
+		return chatRoom
+	}
+
+	// Contain
+	for _, remark := range r.chatRoomRemark {
+		if strings.Contains(remark, key) {
+			return r.remarkToChatRoom[remark]
+		}
+	}
+	for _, nickName := range r.chatRoomNickName {
+		if strings.Contains(nickName, key) {
+			return r.nickNameToChatRoom[nickName]
+		}
+	}
+
+	return nil
+}
+
+func (r *Repository) findChatRooms(key string) []*model.ChatRoom {
+	ret := make([]*model.ChatRoom, 0)
+	distinct := make(map[string]bool)
+	if chatRoom, ok := r.chatRoomCache[key]; ok {
+		ret = append(ret, chatRoom)
+		distinct[chatRoom.Name] = true
+	}
+	if chatRoom, ok := r.remarkToChatRoom[key]; ok && !distinct[chatRoom.Name] {
+		ret = append(ret, chatRoom)
+		distinct[chatRoom.Name] = true
+	}
+	if chatRoom, ok := r.nickNameToChatRoom[key]; ok && !distinct[chatRoom.Name] {
+		ret = append(ret, chatRoom)
+		distinct[chatRoom.Name] = true
+	}
+
+	// Contain
+	for _, remark := range r.chatRoomRemark {
+		if strings.Contains(remark, key) && !distinct[r.remarkToChatRoom[remark].Name] {
+			ret = append(ret, r.remarkToChatRoom[remark])
+			distinct[r.remarkToChatRoom[remark].Name] = true
+		}
+	}
+	for _, nickName := range r.chatRoomNickName {
+		if strings.Contains(nickName, key) && !distinct[r.nickNameToChatRoom[nickName].Name] {
+			ret = append(ret, r.nickNameToChatRoom[nickName])
+			distinct[r.nickNameToChatRoom[nickName].Name] = true
+		}
+	}
+
+	return ret
+}

internal/wechatdb/repository/contact.go

@@ -0,0 +1,211 @@
+package repository
+
+import (
+	"context"
+	"fmt"
+	"sort"
+	"strings"
+
+	"github.com/sjzar/chatlog/internal/model"
+)
+
+// initContactCache 初始化联系人缓存
+func (r *Repository) initContactCache(ctx context.Context) error {
+	// 加载所有联系人到缓存
+	contacts, err := r.ds.GetContacts(ctx, "", 0, 0)
+	if err != nil {
+		return fmt.Errorf("加载联系人失败: %w", err)
+	}
+
+	contactMap := make(map[string]*model.Contact)
+	aliasMap := make(map[string]*model.Contact)
+	remarkMap := make(map[string]*model.Contact)
+	nickNameMap := make(map[string]*model.Contact)
+	chatRoomUserMap := make(map[string]*model.Contact)
+	chatRoomInContactMap := make(map[string]*model.Contact)
+	contactList := make([]string, 0)
+	aliasList := make([]string, 0)
+	remarkList := make([]string, 0)
+	nickNameList := make([]string, 0)
+
+	for _, contact := range contacts {
+		contactMap[contact.UserName] = contact
+		contactList = append(contactList, contact.UserName)
+
+		// 建立快速查找索引
+		if contact.Alias != "" {
+			aliasMap[contact.Alias] = contact
+			aliasList = append(aliasList, contact.Alias)
+		}
+		if contact.Remark != "" {
+			remarkMap[contact.Remark] = contact
+			remarkList = append(remarkList, contact.Remark)
+		}
+		if contact.NickName != "" {
+			nickNameMap[contact.NickName] = contact
+			nickNameList = append(nickNameList, contact.NickName)
+		}
+
+		// 如果是群聊成员（非好友），添加到群聊成员索引
+		if !contact.IsFriend {
+			chatRoomUserMap[contact.UserName] = contact
+		}
+
+		if strings.HasSuffix(contact.UserName, "@chatroom") {
+			chatRoomInContactMap[contact.UserName] = contact
+		}
+	}
+
+	sort.Strings(contactList)
+	sort.Strings(aliasList)
+	sort.Strings(remarkList)
+	sort.Strings(nickNameList)
+
+	r.contactCache = contactMap
+	r.aliasToContact = aliasMap
+	r.remarkToContact = remarkMap
+	r.nickNameToContact = nickNameMap
+	r.chatRoomUserToInfo = chatRoomUserMap
+	r.chatRoomInContact = chatRoomInContactMap
+	r.contactList = contactList
+	r.aliasList = aliasList
+	r.remarkList = remarkList
+	r.nickNameList = nickNameList
+	return nil
+}
+
+func (r *Repository) GetContact(ctx context.Context, key string) (*model.Contact, error) {
+	// 先尝试从缓存中获取
+	contact := r.findContact(key)
+	if contact == nil {
+		return nil, fmt.Errorf("未找到联系人: %s", key)
+	}
+	return contact, nil
+}
+
+func (r *Repository) GetContacts(ctx context.Context, key string, limit, offset int) ([]*model.Contact, error) {
+	ret := make([]*model.Contact, 0)
+	if key != "" {
+		ret = r.findContacts(key)
+		if len(ret) == 0 {
+			return nil, fmt.Errorf("未找到联系人: %s", key)
+		}
+		if limit > 0 {
+			end := offset + limit
+			if end > len(ret) {
+				end = len(ret)
+			}
+			if offset >= len(ret) {
+				return []*model.Contact{}, nil
+			}
+			return ret[offset:end], nil
+		}
+	} else {
+		list := r.contactList
+		if limit > 0 {
+			end := offset + limit
+			if end > len(list) {
+				end = len(list)
+			}
+			if offset >= len(list) {
+				return []*model.Contact{}, nil
+			}
+			list = list[offset:end]
+		}
+		for _, name := range list {
+			ret = append(ret, r.contactCache[name])
+		}
+	}
+	return ret, nil
+}
+
+func (r *Repository) findContact(key string) *model.Contact {
+	if contact, ok := r.contactCache[key]; ok {
+		return contact
+	}
+	if contact, ok := r.aliasToContact[key]; ok {
+		return contact
+	}
+	if contact, ok := r.remarkToContact[key]; ok {
+		return contact
+	}
+	if contact, ok := r.nickNameToContact[key]; ok {
+		return contact
+	}
+
+	// Contain
+	for _, alias := range r.aliasList {
+		if strings.Contains(alias, key) {
+			return r.aliasToContact[alias]
+		}
+	}
+	for _, remark := range r.remarkList {
+		if strings.Contains(remark, key) {
+			return r.remarkToContact[remark]
+		}
+	}
+	for _, nickName := range r.nickNameList {
+		if strings.Contains(nickName, key) {
+			return r.nickNameToContact[nickName]
+		}
+	}
+	return nil
+}
+
+func (r *Repository) findContacts(key string) []*model.Contact {
+	ret := make([]*model.Contact, 0)
+	distinct := make(map[string]bool)
+	if contact, ok := r.contactCache[key]; ok {
+		ret = append(ret, contact)
+		distinct[contact.UserName] = true
+	}
+	if contact, ok := r.aliasToContact[key]; ok && !distinct[contact.UserName] {
+		ret = append(ret, contact)
+		distinct[contact.UserName] = true
+	}
+	if contact, ok := r.remarkToContact[key]; ok && !distinct[contact.UserName] {
+		ret = append(ret, contact)
+		distinct[contact.UserName] = true
+	}
+	if contact, ok := r.nickNameToContact[key]; ok && !distinct[contact.UserName] {
+		ret = append(ret, contact)
+		distinct[contact.UserName] = true
+	}
+	// Contain
+	for _, alias := range r.aliasList {
+		if strings.Contains(alias, key) && !distinct[r.aliasToContact[alias].UserName] {
+			ret = append(ret, r.aliasToContact[alias])
+			distinct[r.aliasToContact[alias].UserName] = true
+		}
+	}
+	for _, remark := range r.remarkList {
+		if strings.Contains(remark, key) && !distinct[r.remarkToContact[remark].UserName] {
+			ret = append(ret, r.remarkToContact[remark])
+			distinct[r.remarkToContact[remark].UserName] = true
+		}
+	}
+	for _, nickName := range r.nickNameList {
+		if strings.Contains(nickName, key) && !distinct[r.nickNameToContact[nickName].UserName] {
+			ret = append(ret, r.nickNameToContact[nickName])
+			distinct[r.nickNameToContact[nickName].UserName] = true
+		}
+	}
+	return ret
+}
+
+// getFullContact 获取联系人信息，包括群聊成员
+func (r *Repository) getFullContact(userName string) *model.Contact {
+	// 先查找联系人缓存
+	if contact, ok := r.contactCache[userName]; ok {
+		return contact
+	}
+
+	// 再查找群聊成员缓存
+	contact, ok := r.chatRoomUserToInfo[userName]
+
+	if ok {
+		return contact
+	}
+
+	return nil
+}

internal/wechatdb/repository/message.go

@@ -0,0 +1,68 @@
+package repository
+
+import (
+	"context"
+	"time"
+
+	"github.com/sjzar/chatlog/internal/model"
+
+	log "github.com/sirupsen/logrus"
+)
+
+// GetMessages 实现 Repository 接口的 GetMessages 方法
+func (r *Repository) GetMessages(ctx context.Context, startTime, endTime time.Time, talker string, limit, offset int) ([]*model.Message, error) {
+
+	if contact, _ := r.GetContact(ctx, talker); contact != nil {
+		talker = contact.UserName
+	} else if chatRoom, _ := r.GetChatRoom(ctx, talker); chatRoom != nil {
+		talker = chatRoom.Name
+	}
+
+	messages, err := r.ds.GetMessages(ctx, startTime, endTime, talker, limit, offset)
+	if err != nil {
+		return nil, err
+	}
+
+	// 补充消息信息
+	if err := r.EnrichMessages(ctx, messages); err != nil {
+		log.Debugf("EnrichMessages failed: %v", err)
+	}
+
+	return messages, nil
+}
+
+// EnrichMessages 补充消息的额外信息
+func (r *Repository) EnrichMessages(ctx context.Context, messages []*model.Message) error {
+	for _, msg := range messages {
+		r.enrichMessage(msg)
+	}
+	return nil
+}
+
+// enrichMessage 补充单条消息的额外信息
+func (r *Repository) enrichMessage(msg *model.Message) {
+	talker := msg.Talker
+
+	// 处理群聊消息
+	if msg.IsChatRoom {
+		talker = msg.ChatRoomSender
+
+		// 补充群聊名称
+		if chatRoom, ok := r.chatRoomCache[msg.Talker]; ok {
+			msg.ChatRoomName = chatRoom.DisplayName()
+
+			// 补充发送者在群里的显示名称
+			if displayName, ok := chatRoom.User2DisplayName[talker]; ok {
+				msg.DisplayName = displayName
+			}
+		}
+	}
+
+	// 如果不是自己发送的消息且还没有显示名称，尝试补充发送者信息
+	if msg.DisplayName == "" && msg.IsSender != 1 {
+		contact := r.getFullContact(talker)
+		if contact != nil {
+			msg.DisplayName = contact.DisplayName()
+		}
+	}
+}

internal/wechatdb/repository/repository.go

@@ -0,0 +1,85 @@
+package repository
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/sjzar/chatlog/internal/model"
+	"github.com/sjzar/chatlog/internal/wechatdb/datasource"
+)
+
+// Repository 实现了 repository.Repository 接口
+type Repository struct {
+	ds datasource.DataSource
+
+	// Cache for contact
+	contactCache      map[string]*model.Contact
+	aliasToContact    map[string]*model.Contact
+	remarkToContact   map[string]*model.Contact
+	nickNameToContact map[string]*model.Contact
+	chatRoomInContact map[string]*model.Contact
+	contactList       []string
+	aliasList         []string
+	remarkList        []string
+	nickNameList      []string
+
+	// Cache for chat room
+	chatRoomCache      map[string]*model.ChatRoom
+	remarkToChatRoom   map[string]*model.ChatRoom
+	nickNameToChatRoom map[string]*model.ChatRoom
+	chatRoomList       []string
+	chatRoomRemark     []string
+	chatRoomNickName   []string
+
+	// 快速查找索引
+	chatRoomUserToInfo map[string]*model.Contact
+}
+
+// New 创建一个新的 Repository
+func New(ds datasource.DataSource) (*Repository, error) {
+	r := &Repository{
+		ds:                 ds,
+		contactCache:       make(map[string]*model.Contact),
+		aliasToContact:     make(map[string]*model.Contact),
+		remarkToContact:    make(map[string]*model.Contact),
+		nickNameToContact:  make(map[string]*model.Contact),
+		chatRoomUserToInfo: make(map[string]*model.Contact),
+		contactList:        make([]string, 0),
+		aliasList:          make([]string, 0),
+		remarkList:         make([]string, 0),
+		nickNameList:       make([]string, 0),
+		chatRoomCache:      make(map[string]*model.ChatRoom),
+		remarkToChatRoom:   make(map[string]*model.ChatRoom),
+		nickNameToChatRoom: make(map[string]*model.ChatRoom),
+		chatRoomList:       make([]string, 0),
+		chatRoomRemark:     make([]string, 0),
+		chatRoomNickName:   make([]string, 0),
+	}
+
+	// 初始化缓存
+	if err := r.initCache(context.Background()); err != nil {
+		return nil, fmt.Errorf("初始化缓存失败: %w", err)
+	}
+
+	return r, nil
+}
+
+// initCache 初始化缓存
+func (r *Repository) initCache(ctx context.Context) error {
+	// 初始化联系人缓存
+	if err := r.initContactCache(ctx); err != nil {
+		return err
+	}
+
+	// 初始化群聊缓存
+	if err := r.initChatRoomCache(ctx); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// Close 实现 Repository 接口的 Close 方法
+func (r *Repository) Close() error {
+	return r.ds.Close()
+}

internal/wechatdb/repository/session.go

@@ -0,0 +1,11 @@
+package repository
+
+import (
+	"context"
+
+	"github.com/sjzar/chatlog/internal/model"
+)
+
+func (r *Repository) GetSessions(ctx context.Context, key string, limit, offset int) ([]*model.Session, error) {
+	return r.ds.GetSessions(ctx, key, limit, offset)
+}