path: root/broker.py

"""
Claude Bridge - Broker Server
消息中转 + Telegram Bot + 任务队列
部署在私人服务器上，24h 运行
"""

import asyncio
import hashlib
import hmac
import os
import shlex
import sqlite3
import subprocess
import time
import uuid
from contextlib import asynccontextmanager

import httpx
from fastapi import Depends, FastAPI, HTTPException, Request, UploadFile, File, Form
from fastapi.responses import JSONResponse, FileResponse
from pydantic import BaseModel

# === Config ===

TELEGRAM_TOKEN = os.environ["TELEGRAM_TOKEN"]
TELEGRAM_CHAT_ID = os.environ["TELEGRAM_CHAT_ID"]  # 你自己的 chat id
API_SECRET = os.environ["API_SECRET"]  # broker 认证密钥
SLACK_BOT_TOKEN = os.environ.get("SLACK_BOT_TOKEN", "")
SLACK_APP_TOKEN = os.environ.get("SLACK_APP_TOKEN", "")
SLACK_ALLOWED_USERS = os.environ.get("SLACK_ALLOWED_USERS", "").split(",")  # 逗号分隔的 user ID 白名单，空=全部放行
DEFAULT_DISPATCHER = "dispatcher"  # 默认 dispatcher session（owner 的）
DB_PATH = os.environ.get("DB_PATH", "bridge.db")
POLL_INTERVAL = int(os.environ.get("POLL_INTERVAL", "2"))  # Telegram 轮询间隔(秒)
DISPATCHER_TMUX = os.environ.get("DISPATCHER_TMUX", "dispatcher")  # dispatcher 的 tmux session 名
HEARTBEAT_TIMEOUT = int(os.environ.get("HEARTBEAT_TIMEOUT", "600"))  # 心跳超时秒数，默认 10 分钟
FILES_DIR = os.environ.get("FILES_DIR", os.path.join(os.path.dirname(os.path.abspath(__file__)), "files"))
os.makedirs(FILES_DIR, exist_ok=True)


# === Database ===

def get_db() -> sqlite3.Connection:
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    return conn


def init_db():
    db = get_db()
    db.executescript("""
        CREATE TABLE IF NOT EXISTS messages (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            chat_id TEXT NOT NULL,
            text TEXT NOT NULL,
            timestamp REAL NOT NULL,
            processed INTEGER DEFAULT 0
        );
        CREATE TABLE IF NOT EXISTS tasks (
            id TEXT PRIMARY KEY,
            target TEXT NOT NULL DEFAULT '',
            type TEXT NOT NULL DEFAULT 'task',
            content TEXT NOT NULL,
            status TEXT DEFAULT 'pending',
            result TEXT DEFAULT '',
            created_at REAL NOT NULL,
            updated_at REAL NOT NULL
        );
    """)
    db.execute("""
        CREATE TABLE IF NOT EXISTS workers (
            id TEXT PRIMARY KEY,
            session_name TEXT NOT NULL,
            host TEXT NOT NULL,
            path TEXT NOT NULL,
            status TEXT DEFAULT 'active',
            created_at REAL NOT NULL
        );
    """)
    db.execute("""
        CREATE TABLE IF NOT EXISTS channel_bindings (
            channel_id TEXT PRIMARY KEY,
            worker_id TEXT NOT NULL,
            channel_name TEXT DEFAULT '',
            created_by TEXT DEFAULT '',
            created_at REAL NOT NULL
        );
    """)
    db.execute("""
        CREATE TABLE IF NOT EXISTS schedules (
            id TEXT PRIMARY KEY,
            action TEXT NOT NULL,
            trigger_at REAL NOT NULL,
            repeat_seconds REAL DEFAULT 0,
            status TEXT DEFAULT 'active',
            created_at REAL NOT NULL
        );
    """)
    db.execute("""
        CREATE TABLE IF NOT EXISTS commands (
            id TEXT PRIMARY KEY,
            target TEXT NOT NULL,
            action TEXT NOT NULL,
            params TEXT NOT NULL DEFAULT '{}',
            status TEXT DEFAULT 'pending',
            result TEXT DEFAULT '',
            created_at REAL NOT NULL
        );
    """)
    db.execute("""
        CREATE TABLE IF NOT EXISTS pending_context (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            text TEXT NOT NULL,
            timestamp REAL NOT NULL
        );
    """)
    db.execute("""
        CREATE TABLE IF NOT EXISTS host_bindings (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            slack_user_id TEXT NOT NULL,
            host TEXT NOT NULL,
            username TEXT NOT NULL,
            home_path TEXT NOT NULL,
            UNIQUE(slack_user_id, host)
        );
    """)
    db.execute("""
        CREATE TABLE IF NOT EXISTS users (
            id TEXT PRIMARY KEY,
            slack_user_id TEXT DEFAULT '',
            telegram_chat_id TEXT DEFAULT '',
            dispatcher_session TEXT NOT NULL,
            display_name TEXT DEFAULT '',
            status TEXT DEFAULT 'active',
            created_at REAL NOT NULL
        );
    """)
    # 兼容旧表：加新列
    migrations = [
        ("tasks", "target", "''"),
        ("tasks", "type", "'task'"),
        ("tasks", "dispatcher_id", "''"),
        ("schedules", "dispatcher_id", "''"),
        ("messages", "dispatcher_id", "''"),
        ("pending_context", "dispatcher_id", "''"),
        ("commands", "dispatcher_id", "''"),
        ("channel_bindings", "created_by", "''"),
    ]
    for table, col, default in migrations:
        try:
            db.execute(f"ALTER TABLE {table} ADD COLUMN {col} TEXT NOT NULL DEFAULT {default}")
        except Exception:
            pass
    # 创建默认用户（迁移：把现有 dispatcher 绑定给 owner）
    existing = db.execute("SELECT id FROM users WHERE id = 'owner'").fetchone()
    if not existing:
        db.execute(
            "INSERT INTO users (id, slack_user_id, telegram_chat_id, dispatcher_session, display_name, status, created_at) VALUES (?, ?, ?, ?, ?, 'active', ?)",
            ("owner", os.environ.get("SLACK_OWNER_ID", ""), TELEGRAM_CHAT_ID, "dispatcher", "Owner", time.time()),
        )
    db.commit()
    db.close()


# === User Helpers ===

def _get_user(user_id: str) -> dict | None:
    db = get_db()
    row = db.execute("SELECT * FROM users WHERE id = ?", (user_id,)).fetchone()
    db.close()
    return dict(row) if row else None


def _get_user_by_slack(slack_user_id: str) -> dict | None:
    db = get_db()
    row = db.execute("SELECT * FROM users WHERE slack_user_id = ?", (slack_user_id,)).fetchone()
    db.close()
    return dict(row) if row else None


def _get_user_by_telegram(telegram_chat_id: str) -> dict | None:
    db = get_db()
    row = db.execute("SELECT * FROM users WHERE telegram_chat_id = ?", (telegram_chat_id,)).fetchone()
    db.close()
    return dict(row) if row else None


def _get_dispatcher_session(dispatcher_id: str) -> str:
    """Look up tmux session name for a dispatcher_id"""
    if not dispatcher_id:
        return DEFAULT_DISPATCHER
    user = _get_user(dispatcher_id)
    return user["dispatcher_session"] if user else DEFAULT_DISPATCHER


# worker session → last dispatcher_id mapping (for reply_to_dispatcher routing)
_worker_dispatcher_map: dict = {}


# === Auth ===

async def verify_token(request: Request):
    auth = request.headers.get("Authorization", "")
    if not auth.startswith("Bearer "):
        raise HTTPException(401, "Missing token")
    token = auth[7:]
    if not hmac.compare_digest(token, API_SECRET):
        raise HTTPException(403, "Invalid token")


# === Notify Dispatcher ===

_last_long_notify = 0  # 防抖：上次发长消息通知的时间


def _notify_dispatcher(text: str, dispatcher_id: str = ""):
    """推送消息到指定 dispatcher。短消息直接 send-keys，长消息写文件。"""
    global _last_long_notify
    session = _get_dispatcher_session(dispatcher_id)
    try:
        safe_text = text.replace("\n", " ").replace("\r", "")
        if len(safe_text) <= 400:
            subprocess.run(
                ["tmux", "send-keys", "-t", session, safe_text, "Enter"],
                timeout=5, capture_output=True,
            )
        else:
            msg_id = uuid.uuid4().hex[:8]
            msg_path = os.path.join(FILES_DIR, f"msg_{msg_id}.txt")
            with open(msg_path, "w") as f:
                f.write(text)
            preview = safe_text[:80]
            subprocess.run(
                ["tmux", "send-keys", "-t", session,
                 f"[Long message saved to {msg_path} — Read it first then reply] Preview: {preview}...", "Enter"],
                timeout=5, capture_output=True,
            )
        print(f"[Broker] Pushed to dispatcher: {text[:60]}...")
    except Exception as e:
        print(f"[Broker] Failed to notify dispatcher: {e}")


# === Telegram Polling ===

class TelegramPoller:
    def __init__(self):
        self.offset = 0
        self.running = False

    async def start(self):
        self.running = True
        async with httpx.AsyncClient(timeout=60) as client:
            while self.running:
                try:
                    resp = await client.get(
                        f"https://api.telegram.org/bot{TELEGRAM_TOKEN}/getUpdates",
                        params={"offset": self.offset, "timeout": 30},
                    )
                    data = resp.json()
                    if data.get("ok"):
                        for update in data.get("result", []):
                            self.offset = update["update_id"] + 1
                            await self._handle_update(update)
                except Exception as e:
                    print(f"[Telegram Poll Error] {e}")
                    await asyncio.sleep(POLL_INTERVAL)

    async def _handle_update(self, update: dict):
        msg = update.get("message", {})
        chat_id = str(msg.get("chat", {}).get("id", ""))

        if not chat_id:
            return

        # 只接收指定 chat_id 的消息（安全）
        if chat_id != TELEGRAM_CHAT_ID:
            print(f"[Telegram] Ignored message from unknown chat: {chat_id}")
            return

        # 处理语音消息
        voice = msg.get("voice")
        if voice:
            text = await self._transcribe_voice(voice["file_id"])
            if not text:
                return
            text = f"[voice] {text}"
        # 处理文件/文档
        elif msg.get("document"):
            file_id = msg["document"]["file_id"]
            file_name = msg["document"].get("file_name", f"file_{file_id}")
            saved_name = await self._download_telegram_file(file_id, file_name)
            if saved_name:
                caption = msg.get("caption", "")
                text = f"[file] {saved_name}" + (f" caption: {caption}" if caption else "")
            else:
                return
        # 处理图片
        elif msg.get("photo"):
            # Telegram 发多个尺寸，取最大的
            photo = msg["photo"][-1]
            file_id = photo["file_id"]
            saved_name = await self._download_telegram_file(file_id, f"photo_{int(time.time())}.jpg")
            if saved_name:
                caption = msg.get("caption", "")
                text = f"[image] {saved_name}" + (f" caption: {caption}" if caption else "")
            else:
                return
        else:
            text = msg.get("text", "")

        if not text:
            return

        db = get_db()
        db.execute(
            "INSERT INTO messages (chat_id, text, timestamp) VALUES (?, ?, ?)",
            (chat_id, text, time.time()),
        )
        db.commit()
        db.close()
        print(f"[Telegram] Stored message: {text[:50]}...")

        # 推送到 dispatcher
        # Telegram 只给 owner 的 dispatcher
        user = _get_user_by_telegram(chat_id)
        did = user["id"] if user else "owner"
        _notify_dispatcher(f"[from_telegram reply with send_telegram_message] {text}", dispatcher_id=did)

    async def _download_telegram_file(self, file_id: str, file_name: str) -> str:
        """从 Telegram 下载文件，存到 FILES_DIR，返回保存的文件名"""
        try:
            async with httpx.AsyncClient(timeout=60) as client:
                resp = await client.get(
                    f"https://api.telegram.org/bot{TELEGRAM_TOKEN}/getFile",
                    params={"file_id": file_id},
                )
                file_path = resp.json()["result"]["file_path"]
                resp = await client.get(
                    f"https://api.telegram.org/file/bot{TELEGRAM_TOKEN}/{file_path}"
                )
                # 防重名
                safe_name = f"{uuid.uuid4().hex[:8]}_{file_name}"
                save_path = os.path.join(FILES_DIR, safe_name)
                with open(save_path, "wb") as f:
                    f.write(resp.content)
                print(f"[Telegram] File saved: {safe_name} ({len(resp.content)} bytes)")
                return safe_name
        except Exception as e:
            print(f"[Telegram] File download failed: {e}")
            return ""

    async def _transcribe_voice(self, file_id: str) -> str:
        """下载 Telegram 语音文件并用 whisper 转文字"""
        try:
            async with httpx.AsyncClient(timeout=30) as client:
                # 获取文件路径
                resp = await client.get(
                    f"https://api.telegram.org/bot{TELEGRAM_TOKEN}/getFile",
                    params={"file_id": file_id},
                )
                file_path = resp.json()["result"]["file_path"]

                # 下载文件
                resp = await client.get(
                    f"https://api.telegram.org/file/bot{TELEGRAM_TOKEN}/{file_path}"
                )
                ogg_path = f"/tmp/voice_{file_id}.ogg"
                with open(ogg_path, "wb") as f:
                    f.write(resp.content)

            print(f"[Telegram] Voice downloaded: {ogg_path}")

            # 用 faster-whisper 转文字
            from faster_whisper import WhisperModel
            if not hasattr(self, '_whisper_model'):
                print("[Whisper] Loading model (first time)...")
                self._whisper_model = WhisperModel("base", device="cpu", compute_type="int8")
                print("[Whisper] Model loaded.")

            segments, info = self._whisper_model.transcribe(ogg_path, language=None)
            text = "".join(seg.text for seg in segments).strip()

            # 清理临时文件
            os.remove(ogg_path)

            print(f"[Whisper] Transcribed ({info.language}): {text[:80]}...")
            return text

        except Exception as e:
            print(f"[Whisper] Transcription failed: {e}")
            return ""

    def stop(self):
        self.running = False


poller = TelegramPoller()


# === Slack Socket Mode ===

_slack_handler = None
_slack_task_channels = {}  # task_id → slack channel_id，worker 完成后回发

def _start_slack_listener():
    """在后台线程启动 Slack Socket Mode listener"""
    if not SLACK_BOT_TOKEN or not SLACK_APP_TOKEN:
        print("[Slack] No tokens configured, skipping")
        return

    import threading
    from slack_sdk.web import WebClient
    from slack_sdk.socket_mode import SocketModeClient
    from slack_sdk.socket_mode.request import SocketModeRequest
    from slack_sdk.socket_mode.response import SocketModeResponse

    slack_client = WebClient(token=SLACK_BOT_TOKEN)

    # 缓存 bot 自己的 user id
    try:
        auth = slack_client.auth_test()
        bot_user_id = auth["user_id"]
        print(f"[Slack] Bot user ID: {bot_user_id}")
    except Exception as e:
        print(f"[Slack] Auth failed: {e}")
        return

    # 缓存 user id → display name
    _user_cache = {}

    def _get_user_name(user_id):
        if user_id not in _user_cache:
            try:
                info = slack_client.users_info(user=user_id)
                _user_cache[user_id] = info["user"]["profile"].get("display_name") or info["user"]["real_name"] or user_id
            except Exception:
                _user_cache[user_id] = user_id
        return _user_cache[user_id]

    def _get_channel_binding(channel_id):
        """查询 channel 是否绑定了 worker"""
        db = get_db()
        row = db.execute(
            "SELECT w.session_name, w.host, w.path, w.id FROM channel_bindings cb JOIN workers w ON cb.worker_id = w.id WHERE cb.channel_id = ?",
            (channel_id,),
        ).fetchone()
        db.close()
        return dict(row) if row else None

    def _find_worker_by_path(host, path):
        """根据 host:path 查找已有 worker"""
        db = get_db()
        row = db.execute(
            "SELECT * FROM workers WHERE host = ? AND path = ? AND status = 'active'",
            (host, path),
        ).fetchone()
        db.close()
        return dict(row) if row else None

    def _create_worker(session_name, host, path):
        """注册新 worker 到 DB"""
        worker_id = uuid.uuid4().hex[:8]
        db = get_db()
        db.execute(
            "INSERT INTO workers (id, session_name, host, path, status, created_at) VALUES (?, ?, ?, ?, 'active', ?)",
            (worker_id, session_name, host, path, time.time()),
        )
        db.commit()
        db.close()
        return worker_id

    def _bind_channel(channel_id, worker_id, channel_name="", created_by=""):
        """绑定 channel 到 worker"""
        db = get_db()
        db.execute(
            "INSERT OR REPLACE INTO channel_bindings (channel_id, worker_id, channel_name, created_by, created_at) VALUES (?, ?, ?, ?, ?)",
            (channel_id, worker_id, channel_name, created_by, time.time()),
        )
        db.commit()
        db.close()

    def _handle_event(client: SocketModeClient, req: SocketModeRequest):
        # 立刻 ack
        client.send_socket_mode_response(SocketModeResponse(envelope_id=req.envelope_id))

        # Debug: log all request types
        print(f"[Slack] req.type={req.type}")

        # --- Slash Command ---
        if req.type in ("slash_commands", "interactive"):
            payload = req.payload
            command = payload.get("command", "")
            cmd_text = payload.get("text", "").strip()
            channel = payload.get("channel_id", "")
            user_id = payload.get("user_id", "")
            channel_name = payload.get("channel_name", "")

            ch_type = "directmessage" if channel_name == "directmessage" else "channel"
            print(f"[Slack] Slash command: {command} text={cmd_text} channel={channel} type={ch_type}")

            if command == "/help":
                slack_client.chat_postMessage(channel=channel, text=(
                    "*Claude Bridge Commands*\n\n"
                    "`/register` — Register for a personal dispatcher\n"
                    "`/init-worker host:path` — Create or bind a worker to this channel\n"
                    "`/meta` — Show worker info (status, heartbeat, tasks)\n"
                    "`/restart-worker` — Restart the worker bound to this channel\n"
                    "`/stop-worker` — Stop worker and unbind from channel\n"
                    "`/unbind` — Unbind worker without stopping it\n"
                    "`/bind-host host:user:home` — Bind a host so you can use `~/` in paths\n"
                    "`/help` — Show this message\n\n"
                    "*In channels with a bound worker:*\n"
                    "• `@bot message` — Chat with the worker (instant, stays in context)\n"
                    "• `@bot /task description` — Submit a formal task (queued, tracked)\n\n"
                    "*In DM:*\n"
                    "• Message the bot directly to chat with your personal dispatcher"
                ))
            elif command == "/register":
                _handle_register(channel, user_id)
            elif command == "/init-worker":
                _handle_init_worker(channel, channel_name, user_id, cmd_text)
            elif command == "/meta":
                _handle_meta(channel, user_id, ch_type)
            elif command == "/restart-worker":
                _handle_worker_command(channel, "restart", user_id, cmd_text, ch_type)
            elif command == "/stop-worker":
                _handle_worker_command(channel, "stop", user_id, cmd_text, ch_type)
            elif command == "/unbind":
                _handle_unbind(channel, user_id)
            elif command == "/bind-host":
                _handle_bind_host(channel, user_id, cmd_text)
            return

        if req.type != "events_api":
            return

        event = req.payload.get("event", {})
        event_type = event.get("type", "")
        user_id = event.get("user", "")
        text = event.get("text", "")
        channel = event.get("channel", "")
        channel_type = event.get("channel_type", "")

        # --- Bot 被加到 channel ---
        if event_type == "member_joined_channel" and user_id == bot_user_id:
            try:
                ch_info = slack_client.conversations_info(channel=channel)
                ch_name = ch_info["channel"].get("name", channel)
            except Exception:
                ch_name = channel
            binding = _get_channel_binding(channel)
            if binding:
                slack_client.chat_postMessage(
                    channel=channel,
                    text=f"Bound to worker `{binding['session_name']}` ({binding['host']}:{binding['path']}). @mention me to chat.",
                )
            else:
                slack_client.chat_postMessage(
                    channel=channel,
                    text=f"👋 Hi! No worker bound to this channel yet.\nUse `/init-worker host:path` to initialize, e.g.:\n`/init-worker timan1:/home/yurenh2/graph-grape`",
                )
            return

        # 忽略 bot 自己的消息
        if user_id == bot_user_id:
            return
        if not text:
            return

        # 权限过滤：已注册用户 或 白名单里的用户
        is_registered = _get_user_by_slack(user_id) is not None
        in_whitelist = SLACK_ALLOWED_USERS and SLACK_ALLOWED_USERS != [''] and user_id in SLACK_ALLOWED_USERS
        if not is_registered and not in_whitelist:
            return

        # 去掉 @mention 标记
        text = text.replace(f"<@{bot_user_id}>", "").strip()

        # 处理文件附件
        files = event.get("files", [])
        file_tags = []
        for f in files:
            try:
                file_url = f.get("url_private_download") or f.get("url_private", "")
                file_name = f.get("name", "unknown")
                file_type = f.get("mimetype", "")
                if file_url:
                    import urllib.request as _ur
                    dl_req = _ur.Request(file_url, headers={"Authorization": f"Bearer {SLACK_BOT_TOKEN}"})
                    with _ur.urlopen(dl_req, timeout=60) as resp:
                        content = resp.read()
                    safe_name = f"{uuid.uuid4().hex[:8]}_{file_name}"
                    save_path = os.path.join(FILES_DIR, safe_name)
                    with open(save_path, "wb") as fh:
                        fh.write(content)
                    if file_type and file_type.startswith("image/"):
                        file_tags.append(f"[image] {safe_name}")
                    else:
                        file_tags.append(f"[file] {safe_name}")
                    print(f"[Slack] File saved: {safe_name} ({len(content)} bytes)")
            except Exception as e:
                print(f"[Slack] File download failed: {e}")
        if file_tags:
            text = " ".join(file_tags) + (" " + text if text else "")

        user_name = _get_user_name(user_id)

        # 获取 channel 名
        channel_name = channel
        try:
            ch_info = slack_client.conversations_info(channel=channel)
            channel_name = ch_info["channel"].get("name", channel)
        except Exception:
            pass

        if event_type == "message" and channel_type == "im":
            # DM → 查用户 → 路由到该用户的 dispatcher
            user_rec = _get_user_by_slack(user_id)
            if not user_rec:
                slack_client.chat_postMessage(channel=channel, text="Please register first with /register")
                return
            did = user_rec["id"]
            tag = f"[from_slack DM @{user_name} reply with send_slack_message channel={channel}]"
            full_text = f"{tag} {text}"
            db = get_db()
            db.execute("INSERT INTO messages (chat_id, text, timestamp, dispatcher_id) VALUES (?, ?, ?, ?)",
                       (f"slack:{channel}:{user_id}", full_text, time.time(), did))
            db.commit()
            db.close()
            print(f"[Slack] DM from {user_name} → dispatcher {did}: {text[:50]}...")
            _notify_dispatcher(full_text, dispatcher_id=did)

        elif event_type == "app_mention":
            # Channel @bot → 检查是否绑定了 worker
            binding = _get_channel_binding(channel)
            if binding:
                session = binding["session_name"]
                # 默认走 message（问答模式），worker 用 reply_to_slack 回复
                # 只有 /task 命令才走任务模式
                if text.startswith("/task "):
                    # 任务模式
                    task_content = text[6:].strip()
                    task_id = uuid.uuid4().hex[:8]
                    now = time.time()
                    db = get_db()
                    db.execute(
                        "INSERT INTO tasks (id, target, type, content, status, created_at, updated_at) VALUES (?, ?, 'task', ?, 'pending', ?, ?)",
                        (task_id, session, f"[from Slack #{channel_name} @{user_name}] {task_content}\n\n[Use report_result when done — result will be posted to Slack channel]", now, now),
                    )
                    db.commit()
                    db.close()
                    _slack_task_channels[task_id] = channel
                    slack_client.chat_postMessage(channel=channel, text=f"📋 Task dispatched to `{session}`, ID: `{task_id}`")
                    print(f"[Slack] #{channel_name} TASK → {session}: {task_content[:50]}...")
                else:
                    # 问答模式：message 直接注入 worker context
                    # worker 用 reply_to_slack 回复，broker 转发到 channel
                    task_id = uuid.uuid4().hex[:8]
                    now = time.time()
                    db = get_db()
                    db.execute(
                        "INSERT INTO tasks (id, target, type, content, status, created_at, updated_at) VALUES (?, ?, 'message', ?, 'pending', ?, ?)",
                        (task_id, session, f"[from Slack #{channel_name} @{user_name} — reply with reply_to_slack] {text}", now, now),
                    )
                    db.commit()
                    db.close()
                    # 记录 message→channel 映射，reply_to_slack 时回发
                    _slack_task_channels[task_id] = channel
                    print(f"[Slack] #{channel_name} MSG → {session}: {text[:50]}...")
            else:
                # 没绑定 → 路由到发送者的 dispatcher
                user_rec = _get_user_by_slack(user_id)
                if not user_rec:
                    slack_client.chat_postMessage(channel=channel, text="Please register first with /register")
                    return
                did = user_rec["id"]
                tag = f"[from_slack #{channel_name} @{user_name} reply with send_slack_message channel={channel}]"
                full_text = f"{tag} {text}"
                db = get_db()
                db.execute("INSERT INTO messages (chat_id, text, timestamp, dispatcher_id) VALUES (?, ?, ?, ?)",
                           (f"slack:{channel}:{user_id}", full_text, time.time(), did))
                db.commit()
                db.close()
                print(f"[Slack] #{channel_name} (unbound) → dispatcher {did}: {text[:50]}...")
                _notify_dispatcher(full_text, dispatcher_id=did)

    # task→channel 映射用全局变量
    global _slack_task_channels

    def _handle_register(channel, slack_user_id):
        """处理 /register — 注册新用户 + 创建 dispatcher"""
        existing = _get_user_by_slack(slack_user_id)
        if existing:
            slack_client.chat_postMessage(
                channel=channel,
                text=f"You are already registered. Dispatcher: `{existing['dispatcher_session']}`\nDM me to chat with your dispatcher.",
            )
            return

        user_name = _get_user_name(slack_user_id)
        user_id = uuid.uuid4().hex[:8]
        session_name = f"dispatcher-{user_id}"

        # 创建 user 记录
        db = get_db()
        db.execute(
            "INSERT INTO users (id, slack_user_id, dispatcher_session, display_name, status, created_at) VALUES (?, ?, ?, ?, 'active', ?)",
            (user_id, slack_user_id, session_name, user_name, time.time()),
        )
        db.commit()
        db.close()

        # 创建 dispatcher tmux session
        workspace = os.path.expanduser(f"~/dispatchers/{user_id}")
        os.makedirs(workspace, exist_ok=True)
        # 复制英文 CLAUDE.md
        en_md = os.path.join(os.path.dirname(os.path.abspath(__file__)), "CLAUDE.md.en")
        if os.path.exists(en_md):
            import shutil
            shutil.copy2(en_md, os.path.join(workspace, "CLAUDE.md"))

        subprocess.run(["tmux", "new-session", "-d", "-s", session_name], timeout=5)
        subprocess.run(["tmux", "send-keys", "-t", session_name, f"cd {workspace}", "Enter"], timeout=5)
        time.sleep(1)
        subprocess.run(["tmux", "send-keys", "-t", session_name, f"export DISPATCHER_USER_ID={user_id}", "Enter"], timeout=5)
        time.sleep(0.5)
        subprocess.run(["tmux", "send-keys", "-t", session_name, "claude", "Enter"], timeout=5)

        slack_client.chat_postMessage(
            channel=channel,
            text=f"✅ Registered! Your dispatcher: `{session_name}`\nDM me to start chatting with your personal assistant.",
        )
        print(f"[Register] Created user {user_id} ({user_name}) with dispatcher {session_name}")

    def _handle_meta(channel, slack_user_id="", channel_type=""):
        """处理 /meta — channel 里显示 worker 信息，DM 里显示 dispatcher + 所有 worker"""
        # DM 模式：显示 dispatcher 信息 + 所有 worker
        if channel_type == "directmessage":
            user_rec = _get_user_by_slack(slack_user_id)
            if not user_rec:
                slack_client.chat_postMessage(channel=channel, text="Not registered. Use `/register` first.")
                return
            lines = [f"*Your Dispatcher*", f"• Session: `{user_rec['dispatcher_session']}`", f"• Status: {'active' if user_rec['status'] == 'active' else 'inactive'}", ""]
            db = get_db()
            workers = db.execute("""
                SELECT w.session_name, w.host, w.path, cb.channel_name
                FROM workers w LEFT JOIN channel_bindings cb ON cb.worker_id = w.id
                WHERE w.status = 'active' ORDER BY w.created_at
            """).fetchall()
            db.close()
            if workers:
                lines.append("*All Workers*")
                for w in workers:
                    session = w["session_name"]
                    hb = heartbeats.get(session, {})
                    status = "🟢" if hb.get("last_seen", 0) > time.time() - HEARTBEAT_TIMEOUT else ("🔴" if hb else "⚪")
                    ch = f" → #{w['channel_name']}" if w["channel_name"] else " (unbound)"
                    lines.append(f"• `{session}` {status} {w['host']}:{w['path']}{ch}")
            else:
                lines.append("No active workers.")
            slack_client.chat_postMessage(channel=channel, text="\n".join(lines))
            return

        # Channel 模式：显示绑定的 worker 信息
        binding = _get_channel_binding(channel)
        if not binding:
            slack_client.chat_postMessage(channel=channel, text="No worker bound to this channel. Use `/init-worker host:path` first.")
            return

        session = binding["session_name"]
        host = binding["host"]
        path = binding["path"]
        worker_id = binding["id"]

        # 心跳
        hb = heartbeats.get(session, {})
        if hb:
            age = int(time.time() - hb["last_seen"])
            alive = "🟢 online" if age < HEARTBEAT_TIMEOUT else f"🔴 offline ({age}s ago)"
            last_hb = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(hb["last_seen"]))
        else:
            alive = "⚪ no heartbeat"
            last_hb = "N/A"

        # 任务统计
        db = get_db()
        running = db.execute("SELECT COUNT(*) FROM tasks WHERE target = ? AND status = 'running'", (session,)).fetchone()[0]
        pending = db.execute("SELECT COUNT(*) FROM tasks WHERE target = ? AND status = 'pending'", (session,)).fetchone()[0]
        done = db.execute("SELECT COUNT(*) FROM tasks WHERE target = ? AND status = 'done'", (session,)).fetchone()[0]

        # 获取 worker 创建时间
        w_row = db.execute("SELECT created_at FROM workers WHERE id = ?", (worker_id,)).fetchone()
        db.close()
        created = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(w_row[0])) if w_row else "N/A"

        text = (
            f"*Worker Meta*\n"
            f"• Session: `{session}`\n"
            f"• Host: `{host}`\n"
            f"• Path: `{path}`\n"
            f"• Status: {alive}\n"
            f"• Last heartbeat: {last_hb}\n"
            f"• Created: {created}\n"
            f"• Tasks — running: {running}, pending: {pending}, done: {done}"
        )
        slack_client.chat_postMessage(channel=channel, text=text)

    def _check_pm(channel, slack_user_id):
        """检查用户是否是该 channel 的 PM（绑定者）"""
        db = get_db()
        row = db.execute("SELECT created_by FROM channel_bindings WHERE channel_id = ?", (channel,)).fetchone()
        db.close()
        if not row or not row["created_by"]:
            return True  # 没记录创建者，放行
        return row["created_by"] == slack_user_id

    def _handle_worker_command(channel, action, slack_user_id="", cmd_text="", channel_type=""):
        """处理 /restart-worker 和 /stop-worker。Channel 里操作绑定的 worker，DM 里接受 session name 参数。"""
        if channel_type == "directmessage" and cmd_text:
            # DM 模式：参数是 session name
            session_name = cmd_text.strip()
            db = get_db()
            w = db.execute("SELECT * FROM workers WHERE session_name = ? AND status = 'active'", (session_name,)).fetchone()
            db.close()
            if not w:
                slack_client.chat_postMessage(channel=channel, text=f"Worker `{session_name}` not found.")
                return
            cmd_id = uuid.uuid4().hex[:8]
            db = get_db()
            db.execute(
                "INSERT INTO commands (id, target, action, params, status, created_at) VALUES (?, ?, ?, '{}', 'pending', ?)",
                (cmd_id, session_name, action, time.time()),
            )
            if action == "stop":
                db.execute("UPDATE workers SET status = 'inactive' WHERE session_name = ?", (session_name,))
                db.execute("DELETE FROM channel_bindings WHERE worker_id = ?", (w["id"],))
            db.commit()
            db.close()
            emoji = "🔄" if action == "restart" else "🛑"
            slack_client.chat_postMessage(channel=channel, text=f"{emoji} {action.title()}ing worker `{session_name}` on {w['host']}...")
            return

        # Channel 模式
        binding = _get_channel_binding(channel)
        if not binding:
            if channel_type == "directmessage":
                slack_client.chat_postMessage(channel=channel, text=f"Usage in DM: `/{action}-worker session_name`")
            else:
                slack_client.chat_postMessage(channel=channel, text="No worker bound to this channel.")
            return
        if not _check_pm(channel, slack_user_id):
            slack_client.chat_postMessage(channel=channel, text="Permission denied. Only the PM (who ran `/init-worker`) can do this.")
            return

        session = binding["session_name"]
        host = binding["host"]
        cmd_id = uuid.uuid4().hex[:8]
        import json as _json
        db = get_db()

        if action == "restart":
            db.execute(
                "INSERT INTO commands (id, target, action, params, status, created_at) VALUES (?, ?, 'restart', '{}', 'pending', ?)",
                (cmd_id, session, time.time()),
            )
            db.commit()
            db.close()
            slack_client.chat_postMessage(channel=channel, text=f"🔄 Restarting worker `{session}` on {host}...")
        elif action == "stop":
            db.execute(
                "INSERT INTO commands (id, target, action, params, status, created_at) VALUES (?, ?, 'stop', '{}', 'pending', ?)",
                (cmd_id, session, time.time()),
            )
            # 解绑 channel + 标记 worker inactive
            db.execute("DELETE FROM channel_bindings WHERE channel_id = ?", (channel,))
            db.execute("UPDATE workers SET status = 'inactive' WHERE session_name = ?", (session,))
            db.commit()
            db.close()
            slack_client.chat_postMessage(channel=channel, text=f"🛑 Stopping worker `{session}` on {host}... Channel unbound.")

    def _handle_unbind(channel, slack_user_id=""):
        """解绑 channel 和 worker"""
        binding = _get_channel_binding(channel)
        if not binding:
            slack_client.chat_postMessage(channel=channel, text="No worker bound to this channel.")
            return
        if not _check_pm(channel, slack_user_id):
            slack_client.chat_postMessage(channel=channel, text="Permission denied. Only the PM (who ran `/init-worker`) can do this.")
            return
        session = binding["session_name"]
        db = get_db()
        db.execute("DELETE FROM channel_bindings WHERE channel_id = ?", (channel,))
        db.commit()
        db.close()
        slack_client.chat_postMessage(channel=channel, text=f"Unbound from worker `{session}`. Use `/init-worker` to bind a new one.")

    def _handle_bind_host(channel, slack_user_id, cmd_text):
        """/bind-host host:username:home_path — e.g. /bind-host timan1:yurenh2:/home/yurenh2"""
        parts = cmd_text.strip().split(":")
        if len(parts) < 3:
            slack_client.chat_postMessage(channel=channel, text=(
                "Usage: `/bind-host host:username:home_path`\n"
                "Example: `/bind-host timan1:yurenh2:/home/yurenh2`\n"
                "After binding, you can use `~/` in paths: `/init-worker timan1:~/project`"
            ))
            return
        host = parts[0].strip()
        username = parts[1].strip()
        home_path = ":".join(parts[2:]).strip()  # 防止路径里有冒号
        db = get_db()
        db.execute(
            "INSERT OR REPLACE INTO host_bindings (slack_user_id, host, username, home_path) VALUES (?, ?, ?, ?)",
            (slack_user_id, host, username, home_path),
        )
        db.commit()
        db.close()
        slack_client.chat_postMessage(channel=channel, text=f"✅ Bound `{host}` → user `{username}`, home `{home_path}`\nNow you can use: `/init-worker {host}:~/your-project`")

    def _resolve_path(slack_user_id, host, path):
        """Expand ~ in path using host_bindings"""
        if "~" not in path:
            return path
        db = get_db()
        row = db.execute(
            "SELECT home_path FROM host_bindings WHERE slack_user_id = ? AND host = ?",
            (slack_user_id, host),
        ).fetchone()
        db.close()
        if row:
            return path.replace("~", row["home_path"])
        return path

    def _handle_init_worker(channel, channel_name, user_id, cmd_text):
        """处理 /init-worker 命令"""
        if not cmd_text or ":" not in cmd_text:
            slack_client.chat_postMessage(
                channel=channel,
                text="Usage: `/init-worker host:path`\nExample: `/init-worker timan1:/home/yurenh2/graph-grape`",
            )
            return

        host, path = cmd_text.split(":", 1)
        host = host.strip()
        path = path.strip()
        path = _resolve_path(user_id, host, path)

        # 检查是否已有这个路径的 worker
        existing = _find_worker_by_path(host, path)
        if existing:
            # 已有 worker，直接绑定
            _bind_channel(channel, existing["id"], channel_name, created_by=user_id)
            slack_client.chat_postMessage(
                channel=channel,
                text=f"✅ Bound to existing worker `{existing['session_name']}` ({host}:{path})\nYou can now @mention me to chat with the worker.",
            )
            print(f"[Slack] Channel {channel} bound to existing worker {existing['session_name']}")
        else:
            # 需要创建新 worker
            # 生成 session 名
            session_name = f"worker-{uuid.uuid4().hex[:6]}"
            worker_id = _create_worker(session_name, host, path)
            _bind_channel(channel, worker_id, channel_name, created_by=user_id)

            # 下发创建命令到 lab
            cmd_id = uuid.uuid4().hex[:8]
            import json as _json
            db = get_db()
            db.execute(
                "INSERT INTO commands (id, target, action, params, status, created_at) VALUES (?, ?, ?, ?, 'pending', ?)",
                (cmd_id, session_name, "create_worker", _json.dumps({"host": host, "path": path, "session_name": session_name, "slack_channel": channel}), time.time()),
            )
            db.commit()
            db.close()

            slack_client.chat_postMessage(
                channel=channel,
                text=f"🚀 Creating new worker `{session_name}` on {host}:{path}\nCommand ID: `{cmd_id}`, please wait...",
            )
            print(f"[Slack] Creating new worker {session_name} at {host}:{path} for channel {channel}")

    # 启动 Socket Mode
    socket_client = SocketModeClient(
        app_token=SLACK_APP_TOKEN,
        web_client=slack_client,
    )
    socket_client.socket_mode_request_listeners.append(_handle_event)

    def _run():
        try:
            socket_client.connect()
            print("[Slack] Socket Mode connected")
            # 保持线程活着
            import select
            while True:
                select.select([], [], [], 60)
        except Exception as e:
            print(f"[Slack] Socket Mode error: {e}")

    t = threading.Thread(target=_run, daemon=True)
    t.start()
    global _slack_handler
    _slack_handler = slack_client
    print("[Slack] Listener started in background thread")


@asynccontextmanager
async def lifespan(app: FastAPI):
    init_db()
    task = asyncio.create_task(poller.start())
    hb_task = asyncio.create_task(_heartbeat_checker())
    sched_task = asyncio.create_task(_schedule_checker())
    _start_slack_listener()
    yield
    poller.stop()
    task.cancel()
    hb_task.cancel()
    sched_task.cancel()


# === FastAPI App ===

app = FastAPI(title="Claude Bridge Broker", lifespan=lifespan)


# --- Pydantic Models ---

class TaskCreate(BaseModel):
    content: str
    target: str = ""
    type: str = "task"
    dispatcher_id: str = ""  # 哪个 dispatcher 创建的


class TaskResult(BaseModel):
    result: str


class TelegramSend(BaseModel):
    message: str


# --- Message Endpoints (调度端用) ---

@app.get("/messages/new", dependencies=[Depends(verify_token)])
def get_new_messages():
    """获取未处理的用户消息"""
    db = get_db()
    rows = db.execute(
        "SELECT id, chat_id, text, timestamp FROM messages WHERE processed = 0 ORDER BY id"
    ).fetchall()
    # 标记为已处理
    if rows:
        ids = [r["id"] for r in rows]
        placeholders = ",".join("?" * len(ids))
        db.execute(f"UPDATE messages SET processed = 1 WHERE id IN ({placeholders})", ids)
        db.commit()
    db.close()
    return {
        "messages": [
            {"id": r["id"], "chat_id": r["chat_id"], "text": r["text"], "timestamp": r["timestamp"]}
            for r in rows
        ]
    }


# --- Task Endpoints ---

@app.post("/tasks", dependencies=[Depends(verify_token)])
def create_task(body: TaskCreate):
    """创建任务（调度端 → 实验室）"""
    task_id = uuid.uuid4().hex[:8]
    now = time.time()
    db = get_db()
    db.execute(
        "INSERT INTO tasks (id, target, type, content, status, dispatcher_id, created_at, updated_at) VALUES (?, ?, ?, ?, 'pending', ?, ?, ?)",
        (task_id, body.target, body.type, body.content, body.dispatcher_id, now, now),
    )
    db.commit()
    db.close()
    return {"id": task_id, "target": body.target, "type": body.type, "content": body.content, "status": "pending"}


@app.get("/tasks/pending", dependencies=[Depends(verify_token)])
def get_pending_tasks(target: str = ""):
    """获取待执行的任务。target 过滤目标 session，只返回匹配的或无指定目标的任务"""
    db = get_db()
    if target:
        rows = db.execute(
            "SELECT id, target, type, content, created_at FROM tasks WHERE status = 'pending' AND (target = ? OR target = '') ORDER BY created_at",
            (target,),
        ).fetchall()
    else:
        rows = db.execute(
            "SELECT id, target, type, content, created_at FROM tasks WHERE status = 'pending' ORDER BY created_at"
        ).fetchall()
    db.close()
    return {"tasks": [dict(r) for r in rows]}


@app.post("/tasks/{task_id}/claim", dependencies=[Depends(verify_token)])
def claim_task(task_id: str):
    """领取任务（实验室端）"""
    db = get_db()
    cur = db.execute(
        "UPDATE tasks SET status = 'running', updated_at = ? WHERE id = ? AND status = 'pending'",
        (time.time(), task_id),
    )
    db.commit()
    if cur.rowcount == 0:
        db.close()
        raise HTTPException(404, "Task not found or already claimed")
    # 更新 worker→dispatcher 映射
    row = db.execute("SELECT target, dispatcher_id FROM tasks WHERE id = ?", (task_id,)).fetchone()
    if row and row["dispatcher_id"]:
        _worker_dispatcher_map[row["target"]] = row["dispatcher_id"]
    db.close()
    return {"ok": True, "task_id": task_id}


@app.post("/tasks/{task_id}/result", dependencies=[Depends(verify_token)])
def submit_result(task_id: str, body: TaskResult):
    """提交任务结果（实验室端）"""
    db = get_db()
    cur = db.execute(
        "UPDATE tasks SET status = 'done', result = ?, updated_at = ? WHERE id = ? AND status = 'running'",
        (body.result, time.time(), task_id),
    )
    db.commit()
    if cur.rowcount == 0:
        db.close()
        raise HTTPException(404, "Task not found or not running")
    db.close()
    # 如果这个 task 关联了 Slack channel 且是 task 类型，发结果到 channel
    # message 类型不发（worker 已经通过 reply_to_slack 直接回复了）
    if _slack_handler:
        channel = _slack_task_channels.pop(task_id, None)
        # 查 task 类型
        db2 = get_db()
        task_row = db2.execute("SELECT type FROM tasks WHERE id = ?", (task_id,)).fetchone()
        task_type = task_row[0] if task_row else "task"
        db2.close()
        if channel and task_type == "task":
            result_preview = body.result[:3000] if body.result else "(无结果)"
            try:
                _slack_handler.chat_postMessage(
                    channel=channel,
                    text=f"✅ Task `{task_id}` completed:\n\n{result_preview}",
                )
                print(f"[Slack] Result for {task_id} sent to channel {channel}")
            except Exception as e:
                print(f"[Slack] Failed to send result to channel: {e}")
    # 路由到正确的 dispatcher
    db3 = get_db()
    task_info = db3.execute("SELECT dispatcher_id FROM tasks WHERE id = ?", (task_id,)).fetchone()
    db3.close()
    did = task_info["dispatcher_id"] if task_info else ""
    _notify_dispatcher(f"Task {task_id} completed. Use check_task_status to view result and report to user.", dispatcher_id=did)
    return {"ok": True}


@app.post("/tasks/{task_id}/fail", dependencies=[Depends(verify_token)])
def fail_task(task_id: str, body: TaskResult):
    """标记任务失败（实验室端）"""
    db = get_db()
    cur = db.execute(
        "UPDATE tasks SET status = 'failed', result = ?, updated_at = ? WHERE id = ?",
        (body.result, time.time(), task_id),
    )
    db.commit()
    db.close()
    db2 = get_db()
    task_info = db2.execute("SELECT dispatcher_id FROM tasks WHERE id = ?", (task_id,)).fetchone()
    db2.close()
    did = task_info["dispatcher_id"] if task_info else ""
    _notify_dispatcher(f"Task {task_id} failed. Use check_task_status to view reason and notify user.", dispatcher_id=did)
    return {"ok": True}


@app.get("/tasks/{task_id}", dependencies=[Depends(verify_token)])
def get_task(task_id: str):
    """查询任务状态"""
    db = get_db()
    row = db.execute("SELECT * FROM tasks WHERE id = ?", (task_id,)).fetchone()
    db.close()
    if not row:
        raise HTTPException(404, "Task not found")
    return dict(row)


@app.get("/tasks", dependencies=[Depends(verify_token)])
def list_tasks(status: str = None, limit: int = 20):
    """列出任务"""
    db = get_db()
    if status:
        rows = db.execute(
            "SELECT * FROM tasks WHERE status = ? ORDER BY created_at DESC LIMIT ?",
            (status, limit),
        ).fetchall()
    else:
        rows = db.execute(
            "SELECT * FROM tasks ORDER BY created_at DESC LIMIT ?", (limit,)
        ).fetchall()
    db.close()
    return {"tasks": [dict(r) for r in rows]}


# --- Workers list ---

@app.get("/workers", dependencies=[Depends(verify_token)])
def list_workers_api():
    """List all active workers with channel bindings"""
    db = get_db()
    rows = db.execute("""
        SELECT w.*, cb.channel_name
        FROM workers w
        LEFT JOIN channel_bindings cb ON cb.worker_id = w.id
        WHERE w.status = 'active'
        ORDER BY w.created_at
    """).fetchall()
    db.close()
    return {"workers": [dict(r) for r in rows]}


# --- Telegram Send (调度端用) ---

@app.post("/telegram/send", dependencies=[Depends(verify_token)])
async def send_telegram(body: TelegramSend):
    """发送 Telegram 消息给用户"""
    async with httpx.AsyncClient() as client:
        resp = await client.post(
            f"https://api.telegram.org/bot{TELEGRAM_TOKEN}/sendMessage",
            json={
                "chat_id": TELEGRAM_CHAT_ID,
                "text": body.message,
                # 不用 Markdown，避免特殊字符解析失败
            },
        )
    if resp.status_code != 200:
        raise HTTPException(502, f"Telegram API error: {resp.text}")
    return {"ok": True}


# --- File Transfer ---

@app.post("/files/upload", dependencies=[Depends(verify_token)])
async def upload_file(file: UploadFile = File(...), filename: str = Form("")):
    """上传文件到 broker 文件存储"""
    name = filename or file.filename or f"upload_{int(time.time())}"
    safe_name = f"{uuid.uuid4().hex[:8]}_{name}"
    save_path = os.path.join(FILES_DIR, safe_name)
    content = await file.read()
    with open(save_path, "wb") as f:
        f.write(content)
    return {"filename": safe_name, "size": len(content)}


@app.get("/files/{filename}")
async def download_file(filename: str, request: Request):
    """下载文件（需认证）"""
    await verify_token(request)
    file_path = os.path.join(FILES_DIR, filename)
    if not os.path.exists(file_path):
        raise HTTPException(404, "File not found")
    return FileResponse(file_path, filename=filename)


@app.get("/files", dependencies=[Depends(verify_token)])
def list_files():
    """列出所有文件"""
    files = []
    for f in os.listdir(FILES_DIR):
        fp = os.path.join(FILES_DIR, f)
        files.append({"name": f, "size": os.path.getsize(fp)})
    return {"files": sorted(files, key=lambda x: x["name"])}


@app.post("/telegram/send_file", dependencies=[Depends(verify_token)])
async def send_telegram_file(filename: str = Form(...), caption: str = Form("")):
    """通过 Telegram 给用户发送文件"""
    file_path = os.path.join(FILES_DIR, filename)
    if not os.path.exists(file_path):
        raise HTTPException(404, "File not found")
    async with httpx.AsyncClient(timeout=60) as client:
        with open(file_path, "rb") as f:
            resp = await client.post(
                f"https://api.telegram.org/bot{TELEGRAM_TOKEN}/sendDocument",
                data={"chat_id": TELEGRAM_CHAT_ID, "caption": caption or ""},
                files={"document": (filename, f)},
            )
    if resp.status_code != 200:
        raise HTTPException(502, f"Telegram API error: {resp.text}")
    return {"ok": True}


# --- Slack Send ---

class SlackSend(BaseModel):
    channel: str  # channel ID
    message: str
    thread_ts: str = ""  # 可选，回复到 thread


@app.post("/slack/send", dependencies=[Depends(verify_token)])
def send_slack_message(body: SlackSend):
    """发送 Slack 消息"""
    if not _slack_handler:
        raise HTTPException(503, "Slack not configured")
    kwargs = {"channel": body.channel, "text": body.message}
    if body.thread_ts:
        kwargs["thread_ts"] = body.thread_ts
    try:
        _slack_handler.chat_postMessage(**kwargs)
    except Exception as e:
        raise HTTPException(502, f"Slack API error: {e}")
    return {"ok": True}


@app.post("/slack/send_file", dependencies=[Depends(verify_token)])
def send_slack_file(filename: str = Form(...), channel: str = Form(...), caption: str = Form("")):
    """Send a file to a Slack channel"""
    if not _slack_handler:
        raise HTTPException(503, "Slack not configured")
    file_path = os.path.join(FILES_DIR, filename)
    if not os.path.exists(file_path):
        raise HTTPException(404, "File not found")
    try:
        _slack_handler.files_upload_v2(
            channel=channel,
            file=file_path,
            title=filename,
            initial_comment=caption or "",
        )
    except Exception as e:
        raise HTTPException(502, f"Slack API error: {e}")
    return {"ok": True}


# --- Log / Reply (实验室 → dispatcher) ---

class LogEntry(BaseModel):
    source: str  # claude / claude2 / claude3 / system
    message: str


@app.post("/log", dependencies=[Depends(verify_token)])
def post_log(body: LogEntry):
    """实验室 Claude 回传消息给 dispatcher"""
    # 从 worker→dispatcher 映射推断目标
    did = _worker_dispatcher_map.get(body.source, "")
    _notify_dispatcher(f"[{body.source}] {body.message}", dispatcher_id=did)
    return {"ok": True}


class SlackHistory(BaseModel):
    session: str  # worker session name, 用来查绑定的 channel
    limit: int = 20


@app.post("/slack/history", dependencies=[Depends(verify_token)])
def get_slack_history(body: SlackHistory):
    """获取 worker 绑定的 Slack channel 的消息历史"""
    if not _slack_handler:
        raise HTTPException(503, "Slack not configured")
    db = get_db()
    row = db.execute(
        "SELECT cb.channel_id FROM channel_bindings cb JOIN workers w ON cb.worker_id = w.id WHERE w.session_name = ?",
        (body.session,),
    ).fetchone()
    db.close()
    if not row:
        raise HTTPException(404, f"No Slack channel bound to session {body.session}")
    channel = row[0]
    try:
        resp = _slack_handler.conversations_history(channel=channel, limit=body.limit)
        messages = []
        for msg in resp.get("messages", []):
            user = msg.get("user", "bot")
            text = msg.get("text", "")
            ts = msg.get("ts", "")
            messages.append({"user": user, "text": text, "ts": ts})
        # 反转让旧消息在前
        messages.reverse()
        return {"messages": messages}
    except Exception as e:
        raise HTTPException(502, f"Slack API error: {e}")


class SlackReply(BaseModel):
    session: str  # worker session name
    message: str


@app.post("/slack/reply", dependencies=[Depends(verify_token)])
def slack_reply(body: SlackReply):
    """Worker 回复到关联的 Slack channel（通过 session→channel_binding 查找）"""
    if not _slack_handler:
        raise HTTPException(503, "Slack not configured")
    # 从 worker session 找到绑定的 channel
    db = get_db()
    row = db.execute(
        "SELECT cb.channel_id FROM channel_bindings cb JOIN workers w ON cb.worker_id = w.id WHERE w.session_name = ?",
        (body.session,),
    ).fetchone()
    db.close()
    if not row:
        raise HTTPException(404, f"No Slack channel bound to session {body.session}")
    channel = row[0]
    try:
        _slack_handler.chat_postMessage(channel=channel, text=body.message)
    except Exception as e:
        raise HTTPException(502, f"Slack API error: {e}")
    return {"ok": True}


# --- GPT-Pro Expert ---

OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
expert_requests: dict = {}  # id -> {status, question, answer}


class ExpertQuestion(BaseModel):
    question: str


@app.post("/expert/ask", dependencies=[Depends(verify_token)])
def ask_expert(body: ExpertQuestion):
    """提交问题给 GPT-Pro，后台异步处理"""
    req_id = uuid.uuid4().hex[:8]
    expert_requests[req_id] = {
        "status": "thinking",
        "question": body.question,
        "answer": "",
    }
    # 后台线程调 OpenAI API
    import threading
    threading.Thread(target=_call_gpt_pro, args=(req_id, body.question), daemon=True).start()
    return {"id": req_id, "status": "thinking"}


@app.get("/expert/{req_id}", dependencies=[Depends(verify_token)])
def get_expert(req_id: str):
    if req_id not in expert_requests:
        raise HTTPException(404, "Request not found")
    return expert_requests[req_id]


def _call_gpt_pro(req_id: str, question: str):
    """同步调用 OpenAI GPT-Pro（在后台线程中运行）"""
    import json
    import urllib.request

    try:
        req = urllib.request.Request(
            "https://api.openai.com/v1/responses",
            data=json.dumps({
                "model": "o3-pro",
                "input": question,
            }).encode(),
            headers={
                "Authorization": f"Bearer {OPENAI_API_KEY}",
                "Content-Type": "application/json",
            },
        )
        with urllib.request.urlopen(req, timeout=600) as resp:
            data = json.loads(resp.read())

        # 提取回复文本
        answer = ""
        for item in data.get("output", []):
            if item.get("type") == "message":
                for content in item.get("content", []):
                    if content.get("type") == "output_text":
                        answer += content.get("text", "")

        expert_requests[req_id]["status"] = "done"
        expert_requests[req_id]["answer"] = answer or "(empty response)"

    except Exception as e:
        expert_requests[req_id]["status"] = "error"
        expert_requests[req_id]["answer"] = str(e)

    # 通知 dispatcher
    status = expert_requests[req_id]["status"]
    _notify_dispatcher(
        f"[system] GPT-Pro reply ready (ID: {req_id}, status: {status}). Use get_expert_answer to view and forward to user."
    )


# --- Commands (lab server执行的系统命令) ---

class CommandCreate(BaseModel):
    target: str  # tmux session name
    action: str  # switch_project, restart, etc.
    params: dict = {}


@app.post("/commands", dependencies=[Depends(verify_token)])
def create_command(body: CommandCreate):
    """创建一个命令让 lab cron 执行"""
    cmd_id = uuid.uuid4().hex[:8]
    db = get_db()
    import json as _json
    db.execute(
        "INSERT INTO commands (id, target, action, params, status, created_at) VALUES (?, ?, ?, ?, 'pending', ?)",
        (cmd_id, body.target, body.action, _json.dumps(body.params), time.time()),
    )
    db.commit()
    db.close()
    return {"id": cmd_id, "target": body.target, "action": body.action}


@app.get("/commands/pending", dependencies=[Depends(verify_token)])
def get_pending_commands():
    """lab cron 轮询待执行命令"""
    db = get_db()
    rows = db.execute(
        "SELECT id, target, action, params FROM commands WHERE status = 'pending' ORDER BY created_at"
    ).fetchall()
    db.close()
    return {"commands": [dict(r) for r in rows]}


@app.post("/commands/{cmd_id}/done", dependencies=[Depends(verify_token)])
def complete_command(cmd_id: str, body: TaskResult):
    """标记命令完成"""
    db = get_db()
    db.execute(
        "UPDATE commands SET status = 'done', result = ? WHERE id = ?",
        (body.result, cmd_id),
    )
    # 查命令的 target session，找到绑定的 Slack channel 回报
    row = db.execute("SELECT target FROM commands WHERE id = ?", (cmd_id,)).fetchone()
    target_session = row[0] if row else ""
    slack_channel = None
    if target_session and _slack_handler:
        ch_row = db.execute(
            "SELECT cb.channel_id FROM channel_bindings cb JOIN workers w ON cb.worker_id = w.id WHERE w.session_name = ?",
            (target_session,),
        ).fetchone()
        if ch_row:
            slack_channel = ch_row[0]
    db.commit()
    db.close()

    result = body.result
    is_ok = result.startswith("OK")
    if slack_channel:
        emoji = "✅" if is_ok else "❌"
        try:
            _slack_handler.chat_postMessage(channel=slack_channel, text=f"{emoji} {result}")
        except Exception:
            pass
    _notify_dispatcher(f"[system] Command {cmd_id} completed: {result}")
    return {"ok": True}


@app.get("/commands/{cmd_id}", dependencies=[Depends(verify_token)])
def get_command(cmd_id: str):
    """查询命令状态"""
    db = get_db()
    row = db.execute("SELECT * FROM commands WHERE id = ?", (cmd_id,)).fetchone()
    db.close()
    if not row:
        raise HTTPException(404, "Command not found")
    return dict(row)


# --- Pending Context (for UserPromptSubmit hook) ---

@app.get("/context/pending", dependencies=[Depends(verify_token)])
def get_pending_context(dispatcher_id: str = ""):
    """Hook 调用：拉取待注入的 context，拉取后自动删除"""
    db = get_db()
    if dispatcher_id:
        rows = db.execute("SELECT id, text FROM pending_context WHERE dispatcher_id = ? OR dispatcher_id = '' ORDER BY id", (dispatcher_id,)).fetchall()
    else:
        rows = db.execute("SELECT id, text FROM pending_context ORDER BY id").fetchall()
    if rows:
        ids = [r["id"] for r in rows]
        placeholders = ",".join("?" * len(ids))
        db.execute(f"DELETE FROM pending_context WHERE id IN ({placeholders})", ids)
        db.commit()
    db.close()
    texts = [r["text"] for r in rows]
    return {"messages": texts}


# --- Schedules (定时任务) ---

class ScheduleCreate(BaseModel):
    action: str
    trigger_at: float = 0
    delay_seconds: float = 0
    repeat_seconds: float = 0
    dispatcher_id: str = ""


@app.post("/schedules", dependencies=[Depends(verify_token)])
def create_schedule(body: ScheduleCreate):
    """创建定时任务"""
    schedule_id = uuid.uuid4().hex[:8]
    now = time.time()
    if body.trigger_at > 0:
        trigger = body.trigger_at
    elif body.delay_seconds > 0:
        trigger = now + body.delay_seconds
    else:
        trigger = now
    db = get_db()
    db.execute(
        "INSERT INTO schedules (id, action, trigger_at, repeat_seconds, dispatcher_id, status, created_at) VALUES (?, ?, ?, ?, ?, 'active', ?)",
        (schedule_id, body.action, trigger, body.repeat_seconds, body.dispatcher_id, now),
    )
    db.commit()
    db.close()
    import datetime
    trigger_str = datetime.datetime.fromtimestamp(trigger).strftime("%Y-%m-%d %H:%M:%S")
    return {"id": schedule_id, "trigger_at": trigger_str, "repeat_seconds": body.repeat_seconds}


@app.get("/schedules", dependencies=[Depends(verify_token)])
def list_schedules():
    """列出所有定时任务"""
    db = get_db()
    rows = db.execute("SELECT * FROM schedules WHERE status = 'active' ORDER BY trigger_at").fetchall()
    db.close()
    return {"schedules": [dict(r) for r in rows]}


@app.delete("/schedules/{schedule_id}", dependencies=[Depends(verify_token)])
def cancel_schedule(schedule_id: str):
    """取消定时任务"""
    db = get_db()
    db.execute("UPDATE schedules SET status = 'cancelled' WHERE id = ?", (schedule_id,))
    db.commit()
    db.close()
    return {"ok": True}


async def _schedule_checker():
    """后台协程：每 30 秒检查定时任务，到时间就触发"""
    while True:
        await asyncio.sleep(30)
        now = time.time()
        db = get_db()
        rows = db.execute(
            "SELECT id, action, trigger_at, repeat_seconds, dispatcher_id FROM schedules WHERE status = 'active' AND trigger_at <= ?",
            (now,),
        ).fetchall()
        for r in rows:
            sid, action, trigger_at, repeat = r["id"], r["action"], r["trigger_at"], r["repeat_seconds"]
            did = r["dispatcher_id"] if "dispatcher_id" in r.keys() else ""
            _notify_dispatcher(f"[scheduled {sid}] {action}", dispatcher_id=did)
            if repeat > 0:
                # 循环任务：更新下次触发时间
                next_trigger = trigger_at + repeat
                # 如果落后太多，跳到下一个未来时间点
                while next_trigger <= now:
                    next_trigger += repeat
                db.execute("UPDATE schedules SET trigger_at = ? WHERE id = ?", (next_trigger, sid))
            else:
                # 一次性任务：标记完成
                db.execute("UPDATE schedules SET status = 'done' WHERE id = ?", (sid,))
        db.commit()
        db.close()


# --- Heartbeat ---

heartbeats: dict = {}  # session -> {host, last_seen, alerted}


class Heartbeat(BaseModel):
    session: str
    host: str = ""


@app.post("/heartbeat", dependencies=[Depends(verify_token)])
def post_heartbeat(body: Heartbeat):
    """Worker 心跳上报"""
    heartbeats[body.session] = {
        "host": body.host,
        "last_seen": time.time(),
        "alerted": False,
    }
    return {"ok": True}


@app.get("/heartbeat", dependencies=[Depends(verify_token)])
def get_heartbeats():
    """查看所有 worker 心跳状态"""
    now = time.time()
    result = {}
    for session, info in heartbeats.items():
        age = now - info["last_seen"]
        result[session] = {
            "host": info["host"],
            "last_seen": info["last_seen"],
            "age_seconds": int(age),
            "alive": age < HEARTBEAT_TIMEOUT,
        }
    return result


async def _heartbeat_checker():
    """后台协程：定期检查心跳，超时则 Telegram 通知用户"""
    while True:
        await asyncio.sleep(60)  # 每分钟检查一次
        now = time.time()
        for session, info in heartbeats.items():
            age = now - info["last_seen"]
            if age > HEARTBEAT_TIMEOUT and not info.get("alerted"):
                # 超时！通知用户
                info["alerted"] = True
                host = info.get("host", "unknown")
                msg = f"Worker {session} ({host}) heartbeat timeout ({int(age)}s). May be offline. Please check and restart."
                try:
                    async with httpx.AsyncClient() as client:
                        await client.post(
                            f"https://api.telegram.org/bot{TELEGRAM_TOKEN}/sendMessage",
                            json={"chat_id": TELEGRAM_CHAT_ID, "text": msg},
                        )
                    print(f"[Heartbeat] ALERT: {session} timeout, notified user")
                except Exception as e:
                    print(f"[Heartbeat] Failed to send alert: {e}")
                # 也通知 dispatcher
                _notify_dispatcher(f"[system] Worker {session} ({host}) heartbeat timeout, may be offline")


# --- Health Check ---

@app.get("/health")
def health():
    return {"status": "ok", "time": time.time()}