ai-mail-bot/main.py

import logging
import logging.handlers
import queue as queue_module
import signal
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from src.config import load_config, Config, UserConfig
from src.summarizer import poll_accounts, ai_process, tg_send_and_mark, _mail_label
from src.tg_bot import poll_telegram, send_thinking
from src.database import init_db, is_uid_processed, mark_fetched, mark_processing, mark_sent, mark_failed

# Thread-safe logging
_log_queue = queue_module.Queue(-1)
_queue_handler = logging.handlers.QueueHandler(_log_queue)
_console_handler = logging.StreamHandler()
_console_handler.setFormatter(logging.Formatter(
    "%(asctime)s [%(levelname)s] [%(threadName)s] %(name)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
))
_log_listener = logging.handlers.QueueListener(_log_queue, _console_handler)

root = logging.getLogger()
root.addHandler(_queue_handler)
root.setLevel(logging.INFO)
_log_listener.start()

logger = logging.getLogger("main")

_running = True
_shutdown_event = threading.Event()

# Queue 1: raw emails → AI processor (user, acct_idx, mail, tg_msg_id)
_email_queue: queue_module.Queue = queue_module.Queue()
# Queue 2: AI summaries → TG sender (user, info_dict)
_tg_queue: queue_module.Queue = queue_module.Queue()


def _signal_handler(signum, frame):
    global _running
    logger.info("收到退出信号，正在停止...")
    _running = False
    _shutdown_event.set()


def _email_poller(cfg: Config):
    logger.info("邮件轮询线程已启动")
    while _running:
        try:
            for user in cfg.users:
                if not _running:
                    return
                for acct_idx, mail in poll_accounts(user):
                    if not _running:
                        return
                    uid_str = mail.uid.decode("utf-8", errors="replace") if isinstance(mail.uid, bytes) else str(mail.uid)
                    if not is_uid_processed(uid_str):
                        mark_fetched(uid_str, mail.account_email, mail.folder,
                                     mail.subject, mail.sender)
                        try:
                            tg_msg_id = send_thinking(cfg.bot_token, user.chat_id,
                                f"📥 *获取邮件中…*\n{_mail_label(mail)}")
                            logger.info("  [%s] 获取邮件中消息已发送: msg_id=%d", user.chat_id, tg_msg_id)
                        except Exception as e:
                            logger.error("  [%s] 获取邮件中消息发送失败: %s", user.chat_id, e)
                            tg_msg_id = 0
                        _email_queue.put((user, acct_idx, mail, tg_msg_id))
        except Exception as e:
            logger.error(f"邮件轮询线程异常: {e}", exc_info=True)
            time.sleep(5)
        if _running:
            _shutdown_event.wait(cfg.polling_interval)


def _ai_processor(cfg: Config):
    logger.info("AI 处理线程已启动 (并行 workers=10)")
    max_workers = 10
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        pending_futures = {}
        while _running:
            while len(pending_futures) < max_workers:
                try:
                    user, acct_idx, mail, tg_msg_id = _email_queue.get_nowait()
                except queue_module.Empty:
                    break
                uid_str = mail.uid.decode("utf-8", errors="replace") if isinstance(mail.uid, bytes) else str(mail.uid)
                mark_processing(uid_str)
                future = executor.submit(ai_process, cfg, user, acct_idx, mail, tg_msg_id)
                pending_futures[future] = (user, acct_idx, mail)

            completed = [f for f in pending_futures if f.done()]
            for future in completed:
                user, acct_idx, mail = pending_futures.pop(future)
                try:
                    info = future.result()
                    if info:
                        _tg_queue.put((user, info))
                except Exception as e:
                    logger.error(f"AI 处理邮件失败: {e}", exc_info=True)
                    uid_str = mail.uid.decode("utf-8", errors="replace") if isinstance(mail.uid, bytes) else str(mail.uid)
                    mark_failed(uid_str)

            time.sleep(0.1 if pending_futures else 0.5)


def _tg_poller(cfg: Config):
    logger.info("TG 轮询线程已启动（按钮回调专用）")
    last_update_id = 0
    while _running:
        try:
            last_update_id = poll_telegram(cfg.bot_token, cfg, last_update_id)
        except Exception as e:
            logger.error(f"TG 轮询错误: {e}", exc_info=True)
        if _running:
            _shutdown_event.wait(0.05)


def _tg_sender(cfg: Config):
    logger.info("TG 发送线程已启动（邮件摘要专用）")
    while _running:
        try:
            user, info = _tg_queue.get_nowait()
            tg_send_and_mark(cfg, user, info)
            uid_str = info["uid"].decode("utf-8", errors="replace") if isinstance(info["uid"], bytes) else str(info["uid"])
            mark_sent(uid_str)
        except queue_module.Empty:
            pass
        except Exception as e:
            logger.error(f"TG 发送失败: {e}", exc_info=True)
            try:
                uid_str = info["uid"].decode("utf-8", errors="replace") if isinstance(info["uid"], bytes) else str(info["uid"])
                mark_failed(uid_str)
            except Exception:
                pass
        if _running:
            _shutdown_event.wait(0.2)


def main():
    global _running
    signal.signal(signal.SIGINT, _signal_handler)
    signal.signal(signal.SIGTERM, _signal_handler)

    cfg_path = sys.argv[1] if len(sys.argv) > 1 else "config.json"
    cfg = load_config(cfg_path)
    init_db()

    user_count = len(cfg.users)
    acct_count = sum(len(u.email_accounts) for u in cfg.users)
    logger.info(f"AI邮件摘要机器人已启动: {user_count} 用户, {acct_count} 邮箱, 轮询间隔: {cfg.polling_interval}s")

    threads = [
        threading.Thread(target=_email_poller, args=(cfg,), daemon=True),
        threading.Thread(target=_ai_processor, args=(cfg,), daemon=True),
        threading.Thread(target=_tg_poller, args=(cfg,), daemon=True),
        threading.Thread(target=_tg_sender, args=(cfg,), daemon=True),
    ]

    for t in threads:
        t.start()

    logger.info("所有线程已启动，进入主循环")
    tick = 0
    try:
        while _running:
            tick += 1
            if tick % 30 == 0:
                logger.info("队列状态: email_queue=%d tg_queue=%d",
                            _email_queue.qsize(), _tg_queue.qsize())
            _shutdown_event.wait(1)
    except KeyboardInterrupt:
        _running = False
        _shutdown_event.set()

    for t in threads:
        t.join(timeout=5)
    _log_listener.stop()
    logger.info("机器人已停止")


if __name__ == "__main__":
    main()