spider/internal/telegram/client.go

package telegram

import (
	"context"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/gotd/td/session"
	"github.com/gotd/td/telegram"
	"github.com/gotd/td/tg"
	"github.com/gotd/td/tgerr"
)

var tmeRegexp = regexp.MustCompile(`https?://t\.me/[^\s"'<>)\]]+`)

// Client TG 客户端
type Client struct {
	account     Account
	sessionPath string

	mu     sync.Mutex
	tgc    *telegram.Client
	api    *tg.Client
	cancel context.CancelFunc
	ready  chan struct{} // closed when connected
	runErr error
}

// New 创建客户端（不连接，只初始化）
func New(account Account) *Client {
	return &Client{
		account:     account,
		sessionPath: account.SessionFile,
		ready:       make(chan struct{}),
	}
}

// Connect 连接并认证（从 session 文件恢复）
// session 文件不存在时返回错误（不做交互式登录，session 需要预先生成）
func (c *Client) Connect(ctx context.Context) error {
	storage := &session.FileStorage{Path: c.sessionPath}

	opts := telegram.Options{
		SessionStorage: storage,
		NoUpdates:      true,
	}

	client := telegram.NewClient(c.account.AppID, c.account.AppHash, opts)

	runCtx, cancel := context.WithCancel(ctx)

	c.mu.Lock()
	c.tgc = client
	c.cancel = cancel
	c.ready = make(chan struct{})
	c.runErr = nil
	readyCh := c.ready
	c.mu.Unlock()

	errCh := make(chan error, 1)

	go func() {
		err := client.Run(runCtx, func(ctx context.Context) error {
			c.mu.Lock()
			c.api = client.API()
			close(readyCh)
			c.mu.Unlock()

			// Block until context is cancelled (Disconnect called)
			<-ctx.Done()
			return ctx.Err()
		})
		c.mu.Lock()
		c.runErr = err
		c.mu.Unlock()
		errCh <- err
	}()

	// Wait for ready or error
	select {
	case <-readyCh:
		return nil
	case err := <-errCh:
		if err != nil && err != context.Canceled {
			return err
		}
		return nil
	case <-ctx.Done():
		cancel()
		return ctx.Err()
	}
}

// Disconnect 断开连接
func (c *Client) Disconnect() {
	c.mu.Lock()
	cancel := c.cancel
	c.mu.Unlock()

	if cancel != nil {
		cancel()
	}
}

// waitReady waits for the client to be connected and returns the api client
func (c *Client) waitReady(ctx context.Context) (*tg.Client, error) {
	c.mu.Lock()
	readyCh := c.ready
	api := c.api
	c.mu.Unlock()

	if api != nil {
		return api, nil
	}

	select {
	case <-readyCh:
		c.mu.Lock()
		api = c.api
		c.mu.Unlock()
		return api, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

// GetChannelInfo 获取频道/用户信息，通过用户名查找
func (c *Client) GetChannelInfo(ctx context.Context, username string) (*ChannelInfo, error) {
	api, err := c.waitReady(ctx)
	if err != nil {
		return nil, err
	}

	username = strings.TrimPrefix(username, "@")

	resolved, err := api.ContactsResolveUsername(ctx, &tg.ContactsResolveUsernameRequest{
		Username: username,
	})
	if err != nil {
		return nil, wrapFloodWait(err)
	}

	info := &ChannelInfo{Username: username}

	// Look for channel/chat in the resolved chats
	for _, ch := range resolved.Chats {
		switch v := ch.(type) {
		case *tg.Channel:
			title := v.Title
			info.Title = title
			info.IsChannel = v.GetBroadcast()
			info.IsGroup = v.GetMegagroup()
			if count, ok := v.GetParticipantsCount(); ok {
				info.MemberCount = count
			}
			// Get full channel info for About
			accessHash, hasHash := v.GetAccessHash()
			if hasHash {
				full, ferr := api.ChannelsGetFullChannel(ctx, &tg.InputChannel{
					ChannelID:  v.GetID(),
					AccessHash: accessHash,
				})
				if ferr == nil {
					if cf, ok := full.FullChat.(*tg.ChannelFull); ok {
						info.About = cf.GetAbout()
						if count, ok := cf.GetParticipantsCount(); ok && info.MemberCount == 0 {
							info.MemberCount = count
						}
					}
				}
			}
			return info, nil
		case *tg.Chat:
			info.Title = v.Title
			info.IsGroup = true
			info.MemberCount = v.ParticipantsCount
			return info, nil
		}
	}

	return info, nil
}

// GetMessages 获取频道历史消息
// offsetID: 从哪条消息开始（断点续传）
// limit: 最多取多少条
// 返回的消息按 ID 从小到大排序
func (c *Client) GetMessages(ctx context.Context, username string, offsetID, limit int) ([]Message, error) {
	api, err := c.waitReady(ctx)
	if err != nil {
		return nil, err
	}

	username = strings.TrimPrefix(username, "@")

	peer, err := c.resolveInputPeer(ctx, api, username)
	if err != nil {
		return nil, err
	}

	result, err := api.MessagesGetHistory(ctx, &tg.MessagesGetHistoryRequest{
		Peer:     peer,
		OffsetID: offsetID,
		Limit:    limit,
	})
	if err != nil {
		return nil, wrapFloodWait(err)
	}

	return extractMessages(result), nil
}

// GetPinnedMessages 获取置顶消息
func (c *Client) GetPinnedMessages(ctx context.Context, username string) ([]Message, error) {
	api, err := c.waitReady(ctx)
	if err != nil {
		return nil, err
	}

	username = strings.TrimPrefix(username, "@")

	peer, err := c.resolveInputPeer(ctx, api, username)
	if err != nil {
		return nil, err
	}

	result, err := api.MessagesSearch(ctx, &tg.MessagesSearchRequest{
		Peer:   peer,
		Filter: &tg.InputMessagesFilterPinned{},
		Limit:  100,
	})
	if err != nil {
		return nil, wrapFloodWait(err)
	}

	return extractMessages(result), nil
}

// VerifyUser 验证用户名是否存在，返回用户信息
func (c *Client) VerifyUser(ctx context.Context, username string) (*UserInfo, error) {
	api, err := c.waitReady(ctx)
	if err != nil {
		return nil, err
	}

	username = strings.TrimPrefix(username, "@")

	resolved, err := api.ContactsResolveUsername(ctx, &tg.ContactsResolveUsernameRequest{
		Username: username,
	})
	if err != nil {
		if tgerr.Is(err, "USERNAME_NOT_OCCUPIED", "USERNAME_INVALID") {
			return &UserInfo{Username: username, Exists: false}, nil
		}
		return nil, wrapFloodWait(err)
	}

	// Check if the peer is a user
	if _, ok := resolved.Peer.(*tg.PeerUser); ok {
		for _, u := range resolved.Users {
			if user, ok := u.(*tg.User); ok {
				info := &UserInfo{
					ID:        user.GetID(),
					Username:  username,
					IsBot:     user.GetBot(),
					IsPremium: user.GetPremium(),
					Exists:    true,
				}
				if fn, ok := user.GetFirstName(); ok {
					info.FirstName = fn
				}
				if ln, ok := user.GetLastName(); ok {
					info.LastName = ln
				}
				if status, ok := user.GetStatus(); ok {
					if offline, ok := status.(*tg.UserStatusOffline); ok {
						t := time.Unix(int64(offline.GetWasOnline()), 0)
						info.LastOnline = &t
					}
				}
				return info, nil
			}
		}
	}

	// Check if it's a channel or group
	for _, ch := range resolved.Chats {
		switch v := ch.(type) {
		case *tg.Channel:
			return &UserInfo{
				ID:        v.GetID(),
				Username:  username,
				IsChannel: v.GetBroadcast(),
				IsGroup:   v.GetMegagroup(),
				Exists:    true,
			}, nil
		case *tg.Chat:
			return &UserInfo{
				ID:      v.ID,
				IsGroup: true,
				Exists:  true,
			}, nil
		}
	}

	return &UserInfo{Username: username, Exists: false}, nil
}

// resolveInputPeer resolves a username to an InputPeer
func (c *Client) resolveInputPeer(ctx context.Context, api *tg.Client, username string) (tg.InputPeerClass, error) {
	resolved, err := api.ContactsResolveUsername(ctx, &tg.ContactsResolveUsernameRequest{
		Username: username,
	})
	if err != nil {
		return nil, wrapFloodWait(err)
	}

	switch p := resolved.Peer.(type) {
	case *tg.PeerChannel:
		for _, ch := range resolved.Chats {
			if channel, ok := ch.(*tg.Channel); ok && channel.GetID() == p.ChannelID {
				accessHash, _ := channel.GetAccessHash()
				return &tg.InputPeerChannel{
					ChannelID:  p.ChannelID,
					AccessHash: accessHash,
				}, nil
			}
		}
		return &tg.InputPeerChannel{ChannelID: p.ChannelID}, nil
	case *tg.PeerUser:
		for _, u := range resolved.Users {
			if user, ok := u.(*tg.User); ok && user.GetID() == p.UserID {
				accessHash, _ := user.GetAccessHash()
				return &tg.InputPeerUser{
					UserID:     p.UserID,
					AccessHash: accessHash,
				}, nil
			}
		}
		return &tg.InputPeerUser{UserID: p.UserID}, nil
	case *tg.PeerChat:
		return &tg.InputPeerChat{ChatID: p.ChatID}, nil
	}

	return &tg.InputPeerEmpty{}, nil
}

// extractMessages extracts messages from a MessagesMessagesClass
func extractMessages(result tg.MessagesMessagesClass) []Message {
	var rawMsgs []tg.MessageClass
	switch v := result.(type) {
	case *tg.MessagesMessages:
		rawMsgs = v.Messages
	case *tg.MessagesMessagesSlice:
		rawMsgs = v.Messages
	case *tg.MessagesChannelMessages:
		rawMsgs = v.Messages
	case *tg.MessagesMessagesNotModified:
		return nil
	}

	var msgs []Message
	for _, raw := range rawMsgs {
		switch m := raw.(type) {
		case *tg.Message:
			msg := Message{
				ID:        m.GetID(),
				Text:      m.GetMessage(),
				IsService: false,
			}

			// Extract forward source channel username
			if fwd, ok := m.GetFwdFrom(); ok {
				if fromID, ok := fwd.GetFromID(); ok {
					if peerCh, ok := fromID.(*tg.PeerChannel); ok {
						_ = peerCh // We'd need channel map to resolve username; skip for now
					}
				}
			}

			// Extract t.me links from text
			msg.Links = tmeRegexp.FindAllString(msg.Text, -1)

			msgs = append(msgs, msg)
		case *tg.MessageService:
			msgs = append(msgs, Message{
				ID:        m.GetID(),
				IsService: true,
			})
		}
	}

	// Sort by ID ascending
	sort.Slice(msgs, func(i, j int) bool {
		return msgs[i].ID < msgs[j].ID
	})

	return msgs
}

// isFloodWait 检查错误是否是 FloodWait，提取等待时间
func isFloodWait(err error) (bool, int) {
	if d, ok := tgerr.AsFloodWait(err); ok {
		return true, int(d.Seconds())
	}
	return false, 0
}

// wrapFloodWait wraps a FloodWait error into FloodWaitError
func wrapFloodWait(err error) error {
	if ok, secs := isFloodWait(err); ok {
		return &FloodWaitError{Seconds: secs}
	}
	return err
}