mygoTgChanBot/internal/telegram/state.go

package telegram

import (
	"sync"
	"time"

	tele "gopkg.in/telebot.v3"
)

type PostStep int

const (
	StepAwaitForward PostStep = iota
	StepAwaitCategory
	StepAwaitTitle
	StepAwaitConfirm
)

type PostState struct {
	UserID        int64
	ForwardedMsg  *tele.Message   // 单条消息
	ForwardedMsgs []*tele.Message // 相册消息
	AlbumID       string          // 相册ID
	SelectedCat   string
	Title         string
	Step          PostStep
	CreatedAt     time.Time
}

type StateManager struct {
	mu     sync.RWMutex
	states map[int64]*PostState
}

func NewStateManager() *StateManager {
	return &StateManager{
		states: make(map[int64]*PostState),
	}
}

func (sm *StateManager) Get(userID int64) *PostState {
	sm.mu.RLock()
	defer sm.mu.RUnlock()
	return sm.states[userID]
}

func (sm *StateManager) Set(state *PostState) {
	sm.mu.Lock()
	defer sm.mu.Unlock()
	sm.states[state.UserID] = state
}

func (sm *StateManager) Delete(userID int64) {
	sm.mu.Lock()
	defer sm.mu.Unlock()
	delete(sm.states, userID)
}

func (sm *StateManager) StartPost(userID int64) *PostState {
	state := &PostState{
		UserID:    userID,
		Step:      StepAwaitForward,
		CreatedAt: time.Now(),
	}
	sm.Set(state)
	return state
}

func (sm *StateManager) SetForwarded(userID int64, msg *tele.Message) *PostState {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	state := sm.states[userID]
	if state == nil {
		return nil
	}

	state.ForwardedMsg = msg
	state.ForwardedMsgs = nil
	state.AlbumID = ""
	state.Step = StepAwaitCategory
	return state
}

// AddAlbumMessage 添加相册消息，返回是否应该继续等待更多消息
func (sm *StateManager) AddAlbumMessage(userID int64, msg *tele.Message) (shouldWait bool) {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	state := sm.states[userID]
	if state == nil {
		return false
	}

	// 第一条相册消息
	if state.AlbumID == "" {
		state.AlbumID = msg.AlbumID
		state.ForwardedMsgs = []*tele.Message{msg}
		return true
	}

	// 同一相册的后续消息
	if state.AlbumID == msg.AlbumID {
		state.ForwardedMsgs = append(state.ForwardedMsgs, msg)
		return true
	}

	return false
}

// FinalizeAlbum 完成相册收集
func (sm *StateManager) FinalizeAlbum(userID int64) *PostState {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	state := sm.states[userID]
	if state == nil {
		return nil
	}

	if len(state.ForwardedMsgs) > 0 {
		state.ForwardedMsg = state.ForwardedMsgs[0] // 用第一条提取标题
	}
	state.Step = StepAwaitCategory
	return state
}

func (sm *StateManager) SetCategory(userID int64, category string) *PostState {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	state := sm.states[userID]
	if state == nil {
		return nil
	}

	state.SelectedCat = category
	state.Step = StepAwaitTitle
	return state
}

func (sm *StateManager) SetTitle(userID int64, title string) *PostState {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	state := sm.states[userID]
	if state == nil {
		return nil
	}

	state.Title = title
	state.Step = StepAwaitConfirm
	return state
}