bot.py

#!/usr/bin/env python3
"""
Claude Code Telegram Bot — Remote control Claude Code from your phone.

Lightweight single-file bot that bridges Telegram messages to `claude -p`
with real-time streaming output, tool call visibility, and session management.
"""

import asyncio
import json
import logging
import os
import shutil
import sys
import time
import uuid
from pathlib import Path

from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import (
    Application,
    CallbackQueryHandler,
    CommandHandler,
    ContextTypes,
    MessageHandler,
    filters,
)
from telegram.constants import ChatAction

# --- Config ---

CONFIG_PATH = Path(__file__).parent / "config.json"
LOG_DIR = Path(__file__).parent / "logs"
LOG_DIR.mkdir(exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler(LOG_DIR / "bot.log", encoding="utf-8"),
        logging.StreamHandler(),
    ],
)
logger = logging.getLogger(__name__)


def find_claude_cli() -> str:
    """Auto-detect claude CLI path."""
    # Check common locations
    candidates = [
        shutil.which("claude"),
        os.path.expanduser("~/.local/bin/claude"),
        os.path.expanduser("~/.claude/local/claude"),
        "/usr/local/bin/claude",
        "/opt/homebrew/bin/claude",
    ]
    for path in candidates:
        if path and os.path.isfile(path) and os.access(path, os.X_OK):
            return path
    return "claude"  # fallback, hope it's on PATH


def load_config() -> dict:
    if not CONFIG_PATH.exists():
        logger.error(f"Config file not found: {CONFIG_PATH}")
        logger.error("Copy config.example.json to config.json and fill in your details.")
        sys.exit(1)
    with open(CONFIG_PATH, "r") as f:
        return json.load(f)


CONFIG = load_config()
BOT_TOKEN = CONFIG["bot_token"]
ALLOWED_USERS = set(CONFIG["allowed_user_ids"])
CLAUDE_PATH = CONFIG.get("claude_path") or find_claude_cli()
WORKING_DIR = CONFIG.get("working_directory", os.getcwd())
TIMEOUT = CONFIG.get("timeout_seconds", 600)
PERMISSION_MODE = CONFIG.get("permission_mode", "default")
SYSTEM_PROMPT = CONFIG.get("system_prompt", "Keep responses concise. Only output useful content.")
MAX_MSG_LEN = 4096
STREAM_UPDATE_INTERVAL = CONFIG.get("stream_update_interval", 2.0)

# --- State ---

active_sessions: dict[int, str | None] = {}
running_tasks: dict[int, bool] = {}
last_exec_time: str | None = None


# --- Auth ---

def auth_check(func):
    async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE):
        user = update.effective_user
        if not user or user.id not in ALLOWED_USERS:
            logger.warning(f"Unauthorized: {user}")
            return
        return await func(update, context)
    return wrapper


# --- Helpers ---

def build_clean_env() -> dict[str, str]:
    """Build a clean environment for claude subprocess (prevents nesting errors)."""
    env = {}
    whitelist = {
        "PATH", "HOME", "USER", "LOGNAME", "SHELL",
        "LANG", "LC_ALL", "LC_CTYPE",
        "TMPDIR", "TEMP", "TMP", "TERM",
        "SSH_AUTH_SOCK", "XDG_CONFIG_HOME", "XDG_DATA_HOME",
    }
    for k, v in os.environ.items():
        if k in whitelist:
            env[k] = v
    # Ensure claude is findable
    claude_dir = os.path.dirname(os.path.abspath(CLAUDE_PATH))
    current_path = env.get("PATH", "/usr/bin:/bin")
    if claude_dir not in current_path:
        env["PATH"] = f"{claude_dir}:{current_path}"
    env["HOME"] = os.path.expanduser("~")
    env["LANG"] = env.get("LANG", "en_US.UTF-8")
    return env


async def safe_edit_message(msg, text: str):
    """Edit a message, ignoring 'not modified' errors."""
    try:
        await msg.edit_text(text)
    except Exception as e:
        if "message is not modified" not in str(e).lower():
            try:
                await msg.edit_text(text[:MAX_MSG_LEN])
            except Exception:
                pass


async def send_long_message(msg_or_update, text: str, bot=None, chat_id=None):
    """Send text, splitting into chunks if needed."""
    if not text:
        text = "(empty response)"
    chunks = split_text(text)
    for i, chunk in enumerate(chunks):
        prefix = f"[{i+1}/{len(chunks)}]\n" if len(chunks) > 1 else ""
        content = prefix + chunk
        try:
            if hasattr(msg_or_update, 'reply_text'):
                await msg_or_update.reply_text(content)
            elif bot and chat_id:
                await bot.send_message(chat_id=chat_id, text=content)
        except Exception as e:
            logger.error(f"Failed to send chunk: {e}")


def split_text(text: str, max_len: int = MAX_MSG_LEN - 100) -> list[str]:
    if len(text) <= max_len:
        return [text]
    chunks = []
    current = ""
    for line in text.split("\n"):
        if len(current) + len(line) + 1 > max_len:
            if current:
                chunks.append(current)
                current = ""
            while len(line) > max_len:
                chunks.append(line[:max_len])
                line = line[max_len:]
            current = line
        else:
            current = current + "\n" + line if current else line
    if current:
        chunks.append(current)
    return chunks if chunks else [text]


# --- Streaming Claude execution ---

async def run_claude_streaming(
    prompt: str,
    status_msg,
    bot,
    chat_id: int,
    session_id: str | None = None,
    continue_last: bool = False,
    resume_session: str | None = None,
):
    """Run claude -p with stream-json, updating Telegram message in real-time."""
    global last_exec_time

    cmd = [
        CLAUDE_PATH, "-p",
        "--output-format", "stream-json", "--verbose",
        "--append-system-prompt", SYSTEM_PROMPT,
    ]

    if continue_last:
        cmd.append("-c")
    if resume_session:
        cmd.extend(["--resume", resume_session])
    elif session_id:
        cmd.extend(["--session-id", session_id])
    if PERMISSION_MODE and PERMISSION_MODE != "default":
        cmd.extend(["--permission-mode", PERMISSION_MODE])

    cmd.append(prompt)

    logger.info(f"Streaming: prompt={len(prompt)} chars, session={session_id}")

    try:
        proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=WORKING_DIR,
            env=build_clean_env(),
        )
    except FileNotFoundError:
        await safe_edit_message(status_msg, f"Claude CLI not found: {CLAUDE_PATH}")
        return

    full_text = ""
    tool_events = []
    result_info = None
    last_update_time = 0.0
    display_text = ""

    async def update_display(force: bool = False):
        nonlocal last_update_time, display_text
        now = time.time()
        if not force and now - last_update_time < STREAM_UPDATE_INTERVAL:
            return
        parts = []
        if tool_events:
            tool_lines = []
            for name, status in tool_events:
                icon = "\U0001f527" if status == "running" else "\u2705"
                tool_lines.append(f"{icon} {name}")
            parts.append("\n".join(tool_lines))
            parts.append("---")
        if full_text:
            text_preview = full_text
            overhead = sum(len(p) + 1 for p in parts) + 50
            max_text = MAX_MSG_LEN - overhead
            if len(text_preview) > max_text:
                text_preview = "...\n" + text_preview[-max_text:]
            parts.append(text_preview)
        if not parts:
            parts.append("Claude is thinking...")
        new_display = "\n".join(parts)
        if new_display != display_text:
            display_text = new_display
            await safe_edit_message(status_msg, display_text)
            last_update_time = now

    try:
        typing_task = asyncio.create_task(_typing_loop(bot, chat_id, proc))
        deadline = time.time() + TIMEOUT

        while True:
            if time.time() > deadline:
                proc.kill()
                await safe_edit_message(status_msg, f"Timeout ({TIMEOUT}s)")
                typing_task.cancel()
                return
            try:
                line = await asyncio.wait_for(
                    proc.stdout.readline(), timeout=min(30, deadline - time.time())
                )
            except asyncio.TimeoutError:
                await update_display(force=True)
                continue

            if not line:
                break
            line_str = line.decode("utf-8", errors="replace").strip()
            if not line_str:
                continue

            try:
                event = json.loads(line_str)
            except json.JSONDecodeError:
                full_text += line_str + "\n"
                await update_display()
                continue

            etype = event.get("type", "")

            if etype == "assistant" and "message" in event:
                content = event["message"].get("content", [])
                for block in content:
                    if block.get("type") == "text":
                        full_text = block.get("text", "")
                    elif block.get("type") == "tool_use":
                        name = block.get("name", "unknown")
                        tool_events = [(n, "done" if s == "running" else s) for n, s in tool_events]
                        tool_events.append((name, "running"))
                await update_display()

            elif etype == "content_block_delta":
                delta = event.get("delta", {})
                if delta.get("type") == "text_delta":
                    full_text += delta.get("text", "")
                    await update_display()

            elif etype == "content_block_start":
                block = event.get("content_block", {})
                if block.get("type") == "tool_use":
                    name = block.get("name", "unknown")
                    tool_events = [(n, "done" if s == "running" else s) for n, s in tool_events]
                    tool_events.append((name, "running"))
                    await update_display(force=True)

            elif etype == "result":
                tool_events = [(n, "done") for n, _ in tool_events]
                result_info = event
                await update_display(force=True)

        typing_task.cancel()
        await proc.wait()

    except Exception as e:
        logger.exception("Error during streaming")
        await safe_edit_message(status_msg, f"Error: {e}")
        return

    last_exec_time = time.strftime("%Y-%m-%d %H:%M:%S")

    stderr_data = await proc.stderr.read()
    stderr_text = stderr_data.decode("utf-8", errors="replace").strip() if stderr_data else ""

    if proc.returncode != 0 and not full_text:
        error_text = f"Claude exit code: {proc.returncode}"
        if stderr_text:
            error_text += f"\n{stderr_text}"
        await safe_edit_message(status_msg, error_text)
        return

    if result_info:
        usage = result_info.get("usage", {})
        input_t = usage.get("input_tokens", 0)
        output_t = usage.get("output_tokens", 0)
        cost = result_info.get("cost_usd") or result_info.get("total_cost_usd")
        logger.info(f"Tokens: {input_t} in / {output_t} out" + (f" | ${cost:.4f}" if cost else ""))
        sid = result_info.get("session_id")
        if sid:
            logger.info(f"Session: {sid}")

    final_text = full_text if full_text else "(no output)"

    if len(final_text) <= MAX_MSG_LEN:
        await safe_edit_message(status_msg, final_text)
    else:
        try:
            await status_msg.delete()
        except Exception:
            pass
        await send_long_message(None, final_text, bot=bot, chat_id=chat_id)


async def _typing_loop(bot, chat_id: int, proc):
    try:
        while proc.returncode is None:
            try:
                await bot.send_chat_action(chat_id=chat_id, action=ChatAction.TYPING)
            except Exception:
                pass
            await asyncio.sleep(4)
    except asyncio.CancelledError:
        pass


# --- Handlers ---

@auth_check
async def cmd_start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    await update.message.reply_text(
        "Bot online! Send any text to execute via Claude Code.\n\n"
        "Commands:\n"
        "/status \u2014 View status\n"
        "/chat \u2014 Start multi-turn conversation\n"
        "/endchat \u2014 End multi-turn conversation\n"
        "/continue \u2014 Continue last terminal session\n"
        "/resume <id> \u2014 Resume specific session\n"
        "/mode \u2014 Switch permission mode\n"
        "/reload \u2014 Reload config"
    )


@auth_check
async def cmd_status(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = update.effective_user.id
    session = active_sessions.get(user_id)
    is_running = running_tasks.get(user_id, False)
    await update.message.reply_text(
        f"Mode: {'multi-turn' if session else 'single-shot'}\n"
        f"Session: {session or 'N/A'}\n"
        f"Task: {'running' if is_running else 'idle'}\n"
        f"Last run: {last_exec_time or 'never'}\n"
        f"Permission: {PERMISSION_MODE}\n"
        f"Timeout: {TIMEOUT}s"
    )


@auth_check
async def cmd_chat(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = update.effective_user.id
    session_id = str(uuid.uuid4())
    active_sessions[user_id] = session_id
    await update.message.reply_text(f"Multi-turn started (session: {session_id[:8]}...)\nSend /endchat to end.")


@auth_check
async def cmd_endchat(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = update.effective_user.id
    if user_id in active_sessions:
        active_sessions.pop(user_id)
        await update.message.reply_text("Multi-turn ended.")
    else:
        await update.message.reply_text("No active multi-turn session.")


@auth_check
async def cmd_continue(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = update.effective_user.id
    if running_tasks.get(user_id):
        await update.message.reply_text("A task is already running.")
        return
    prompt = " ".join(context.args) if context.args else "continue"
    status_msg = await update.message.reply_text("Continuing last session...")
    running_tasks[user_id] = True
    try:
        await run_claude_streaming(
            prompt, status_msg, context.bot, update.effective_chat.id,
            continue_last=True,
        )
    finally:
        running_tasks[user_id] = False


@auth_check
async def cmd_resume(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_id = update.effective_user.id
    if running_tasks.get(user_id):
        await update.message.reply_text("A task is already running.")
        return
    if not context.args:
        await update.message.reply_text("Usage: /resume <session_id> [prompt]")
        return
    resume_id = context.args[0]
    prompt = " ".join(context.args[1:]) if len(context.args) > 1 else "continue"
    status_msg = await update.message.reply_text(f"Resuming {resume_id[:8]}...")
    running_tasks[user_id] = True
    try:
        await run_claude_streaming(
            prompt, status_msg, context.bot, update.effective_chat.id,
            resume_session=resume_id,
        )
    finally:
        running_tasks[user_id] = False


@auth_check
async def cmd_mode(update: Update, context: ContextTypes.DEFAULT_TYPE):
    global PERMISSION_MODE
    modes = ["default", "acceptEdits", "plan", "bypassPermissions"]
    buttons = [[InlineKeyboardButton(m, callback_data=f"mode:{m}")] for m in modes]
    await update.message.reply_text(
        f"Current: {PERMISSION_MODE}\nSelect new mode:",
        reply_markup=InlineKeyboardMarkup(buttons),
    )


async def callback_mode(update: Update, context: ContextTypes.DEFAULT_TYPE):
    query = update.callback_query
    if not query or not query.from_user or query.from_user.id not in ALLOWED_USERS:
        return
    await query.answer()
    global PERMISSION_MODE
    mode = query.data.replace("mode:", "")
    PERMISSION_MODE = mode
    await query.edit_message_text(f"Permission mode: {mode}")


@auth_check
async def cmd_reload(update: Update, context: ContextTypes.DEFAULT_TYPE):
    global CONFIG, ALLOWED_USERS, CLAUDE_PATH, WORKING_DIR, TIMEOUT, PERMISSION_MODE, SYSTEM_PROMPT
    try:
        CONFIG = load_config()
        ALLOWED_USERS = set(CONFIG["allowed_user_ids"])
        CLAUDE_PATH = CONFIG.get("claude_path") or find_claude_cli()
        WORKING_DIR = CONFIG.get("working_directory", os.getcwd())
        TIMEOUT = CONFIG.get("timeout_seconds", 600)
        PERMISSION_MODE = CONFIG.get("permission_mode", "default")
        SYSTEM_PROMPT = CONFIG.get("system_prompt", "Keep responses concise. Only output useful content.")
        await update.message.reply_text("Config reloaded.")
    except Exception as e:
        await update.message.reply_text(f"Failed: {e}")


@auth_check
async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
    text = update.message.text
    if not text:
        return
    user_id = update.effective_user.id
    if running_tasks.get(user_id):
        await update.message.reply_text("A task is already running.")
        return
    session_id = active_sessions.get(user_id)
    status_msg = await update.message.reply_text("Claude is thinking...")
    running_tasks[user_id] = True
    try:
        await run_claude_streaming(
            text, status_msg, context.bot, update.effective_chat.id,
            session_id=session_id,
        )
    finally:
        running_tasks[user_id] = False


@auth_check
async def handle_document(update: Update, context: ContextTypes.DEFAULT_TYPE):
    doc = update.message.document
    if not doc:
        return
    file = await context.bot.get_file(doc.file_id)
    save_dir = os.path.join(WORKING_DIR, "inbox")
    os.makedirs(save_dir, exist_ok=True)
    save_path = os.path.join(save_dir, doc.file_name)
    await file.download_to_drive(save_path)
    await update.message.reply_text(f"File saved: inbox/{doc.file_name}")

    if update.message.caption:
        user_id = update.effective_user.id
        if running_tasks.get(user_id):
            await update.message.reply_text("A task is already running.")
            return
        prompt = f"User uploaded file to inbox/{doc.file_name}. Instruction: {update.message.caption}"
        status_msg = await update.message.reply_text("Processing file...")
        running_tasks[user_id] = True
        try:
            await run_claude_streaming(
                prompt, status_msg, context.bot, update.effective_chat.id,
                session_id=active_sessions.get(user_id),
            )
        finally:
            running_tasks[user_id] = False


# --- Main ---

def main():
    if BOT_TOKEN == "YOUR_BOT_TOKEN_HERE":
        logger.error("Set bot_token in config.json!")
        sys.exit(1)
    if not ALLOWED_USERS:
        logger.error("Set allowed_user_ids in config.json!")
        sys.exit(1)

    logger.info(f"Starting bot | claude={CLAUDE_PATH} | cwd={WORKING_DIR}")

    app = Application.builder().token(BOT_TOKEN).build()
    app.add_handler(CommandHandler("start", cmd_start))
    app.add_handler(CommandHandler("status", cmd_status))
    app.add_handler(CommandHandler("chat", cmd_chat))
    app.add_handler(CommandHandler("endchat", cmd_endchat))
    app.add_handler(CommandHandler("continue", cmd_continue))
    app.add_handler(CommandHandler("resume", cmd_resume))
    app.add_handler(CommandHandler("mode", cmd_mode))
    app.add_handler(CommandHandler("reload", cmd_reload))
    app.add_handler(CallbackQueryHandler(callback_mode, pattern=r"^mode:"))
    app.add_handler(MessageHandler(filters.Document.ALL, handle_document))
    app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))

    logger.info("Bot started, polling...")
    app.run_polling(drop_pending_updates=True)


if __name__ == "__main__":
    main()