steward_mirror/steward/bot/telegram.py

"""Telegram bot interface for Steward."""

import logging
from collections import defaultdict
from typing import Any

from telegram import Update
from telegram.constants import ParseMode
from telegram.ext import (
    Application,
    CommandHandler,
    ContextTypes,
    MessageHandler,
    filters,
)

from steward.config import Settings
from steward.llm.client import LLMClient
from steward.memory.thread_store import ThreadMemoryStore, ThreadSummary
from steward.proposals.generator import Proposal, ProposalGenerator
from steward.tools.client import ToolClient

logger = logging.getLogger(__name__)

# Per-user conversation history for non-threaded messages (capped at _MAX_HISTORY turns).
_history: dict[int, list[dict[str, Any]]] = defaultdict(list)
_MAX_HISTORY = 20

# Per-thread conversation history: (chat_id, thread_id) → full history (unbounded).
# Messages belonging to a Telegram message thread are kept in their entirety here
# until explicitly flushed by the /flush command.
_thread_history: dict[tuple[int, int], list[dict[str, Any]]] = defaultdict(list)

_FLUSH_SYSTEM_PROMPT = (
    "You are Steward. The following is a complete Telegram message thread conversation. "
    "Produce a concise but comprehensive summary that captures:\n"
    "- The main topics discussed\n"
    "- Key decisions or conclusions reached\n"
    "- Any outstanding actions or open questions\n"
    "- Important context that would help recall this conversation later\n\n"
    "Be precise. Omit pleasantries."
)

_TAGS_SYSTEM_PROMPT = (
    "You are a keyword tagger for a knowledge base. "
    "Extract 5–8 short, lowercase keyword tags from the following conversation summary. "
    "Tags should represent the main topics, entities, and concepts discussed. "
    "Return ONLY a comma-separated list of tags with no other text or punctuation. "
    "Example output: api design, authentication, database schema, user roles, caching"
)

_KB_CONTEXT_HEADER = (
    "The following are relevant past conversation summaries from your knowledge base. "
    "Use them as background context if they relate to the current question, "
    "but do not repeat their contents unless directly asked."
)


def _thread_key(update: Update) -> tuple[int, int] | None:
    """Return the (chat_id, thread_id) key if the message is part of a thread, else None."""
    msg = update.message
    chat = update.effective_chat
    if msg is None or chat is None:
        return None
    thread_id = msg.message_thread_id
    if thread_id is None:
        return None
    return (chat.id, thread_id)


def _is_allowed(user_id: int, settings: Settings) -> bool:
    """Return True if the user is in the allow-list (or no list is configured)."""
    if not settings.telegram_allowed_user_ids:
        return True
    allowed = user_id in settings.telegram_allowed_user_ids
    if not allowed:
        logger.info("Ignoring update from unauthorized user %s", user_id)
    return allowed


def _is_group_enabled(chat_id: int, settings: Settings) -> bool:
    """Return True if the group/channel is in the configured list (or no list is configured)."""
    if not settings.telegram_group_ids:
        return True
    return chat_id in settings.telegram_group_ids


def _is_chat_enabled(chat: Any, settings: Settings) -> bool:
    """Return True if the chat is allowed.

    Private chats are governed only by the user allow-list. Group/channel allow-listing
    applies only to non-private chats.
    """
    if getattr(chat, "type", None) == "private":
        return True

    allowed = _is_group_enabled(chat.id, settings)
    if not allowed:
        logger.info(
            "Ignoring update in unauthorized chat %s (type=%s); configured group IDs: %s",
            chat.id,
            getattr(chat, "type", None),
            settings.telegram_group_ids,
        )
    return allowed


async def _send_long(update: Update, text: str) -> None:
    """Send text, splitting across messages if it exceeds Telegram's 4096-char limit."""
    limit = 4096
    for i in range(0, len(text), limit):
        await update.message.reply_text(  # type: ignore[union-attr]
            text[i : i + limit],
            parse_mode=ParseMode.MARKDOWN,
        )


async def start_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    """Handle /start."""
    settings: Settings = context.bot_data["settings"]
    user = update.effective_user
    if user is None or not _is_allowed(user.id, settings):
        return

    chat = update.effective_chat
    if chat is None or not _is_chat_enabled(chat, settings):
        return

    await update.message.reply_text(  # type: ignore[union-attr]
        "Hello, I'm *Steward* \U0001f916\n\n"
        "I'm your AI-assisted personal operations platform.\n"
        "Talk to me naturally, or use:\n"
        "/help – show available commands\n"
        "/clear – reset conversation history\n"
        "/flush – summarise and archive this thread's memory\n"
        "/recall [query] – retrieve archived thread summaries\n"
        "/analyse – run a manual API analysis right now",
        parse_mode=ParseMode.MARKDOWN,
    )


async def help_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    """Handle /help."""
    settings: Settings = context.bot_data["settings"]
    user = update.effective_user
    if user is None or not _is_allowed(user.id, settings):
        return

    chat = update.effective_chat
    if chat is None or not _is_chat_enabled(chat, settings):
        return

    await update.message.reply_text(  # type: ignore[union-attr]
        "*Steward commands*\n\n"
        "/start – greeting\n"
        "/help – this message\n"
        "/clear – reset conversation history for this context\n"
        "/flush – summarise the current thread, store the summary, and compress memory\n"
        "         _(only available inside a message thread)_\n"
        "/recall [query] – show this thread's summary, list all summaries, or search by keyword\n"
        "/analyse – trigger an immediate API analysis and proposal",
        parse_mode=ParseMode.MARKDOWN,
    )


async def clear_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    """Handle /clear – wipe conversation history for this context.

    Inside a message thread: clears the thread's unbounded history.
    Outside a thread: clears the per-user capped history.
    """
    settings: Settings = context.bot_data["settings"]
    user = update.effective_user
    if user is None or not _is_allowed(user.id, settings):
        return

    chat = update.effective_chat
    if chat is None or not _is_chat_enabled(chat, settings):
        return

    key = _thread_key(update)
    if key is not None:
        _thread_history[key].clear()
        await update.message.reply_text(  # type: ignore[union-attr]
            "Thread conversation history cleared."
        )
    else:
        _history[user.id].clear()
        await update.message.reply_text(  # type: ignore[union-attr]
            "Conversation history cleared."
        )


async def flush_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    """Handle /flush – summarise thread memory, persist it, and compress in-memory history.

    Steps:
    1. Verify the command is issued inside a message thread.
    2. Summarise the full thread history via the LLM.
    3. Persist the summary in the ThreadMemoryStore (keyed by chat_id + thread_id).
    4. Replace the in-memory thread history with a single compressed context message
       so conversation can continue with the summary as background.
    """
    settings: Settings = context.bot_data["settings"]
    llm: LLMClient = context.bot_data["llm"]
    store: ThreadMemoryStore = context.bot_data["thread_store"]
    user = update.effective_user

    if user is None or not _is_allowed(user.id, settings):
        return

    chat = update.effective_chat
    if chat is None or not _is_chat_enabled(chat, settings):
        return

    key = _thread_key(update)
    if key is None:
        await update.message.reply_text(  # type: ignore[union-attr]
            "\u26a0\ufe0f /flush can only be used inside a message thread."
        )
        return

    chat_id, thread_id = key
    history = _thread_history[key]

    if not history:
        await update.message.reply_text(  # type: ignore[union-attr]
            "This thread has no conversation history to flush."
        )
        return

    await update.message.reply_text(  # type: ignore[union-attr]
        "\U0001f4be Summarising thread memory\u2026"
    )

    # Build a readable transcript for the LLM to summarise
    transcript_lines = []
    for msg in history:
        role_label = "User" if msg["role"] == "user" else "Steward"
        transcript_lines.append(f"{role_label}: {msg['content']}")
    transcript = "\n".join(transcript_lines)

    summary_text = await llm.chat(
        f"Thread transcript:\n\n{transcript}",
        system_prompt=_FLUSH_SYSTEM_PROMPT,
    )

    # Extract keyword tags for knowledge-base indexing (second LLM call, lightweight)
    tags_raw = await llm.chat(
        f"Summary to tag:\n\n{summary_text}",
        system_prompt=_TAGS_SYSTEM_PROMPT,
    )
    tags = [t.strip().lower() for t in tags_raw.split(",") if t.strip()][:8]

    message_count = sum(1 for m in history if m["role"] == "user")
    thread_summary = ThreadSummary(
        chat_id=chat_id,
        thread_id=thread_id,
        summary=summary_text,
        message_count=message_count,
        tags=tags,
    )
    store.save(thread_summary)

    # Compress: replace history with a single system-context entry so the thread
    # can continue with the summary as background knowledge.
    _thread_history[key] = [
        {"role": "system", "content": f"Summary of earlier conversation:\n{summary_text}"}
    ]

    await update.message.reply_text(  # type: ignore[union-attr]
        f"\u2705 Thread memory flushed and stored "
        f"(thread `{thread_id}`, {message_count} messages summarised).",
        parse_mode=ParseMode.MARKDOWN,
    )


async def recall_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    """Handle /recall [query] – retrieve stored thread summaries.

    With a query argument (e.g. ``/recall api design``): searches all stored
    summaries whose tags overlap with the query keywords and returns matches.

    Without a query:
    - Inside a thread: shows the stored summary for this thread (if any).
    - Outside a thread: lists all stored summaries (newest first).
    """
    settings: Settings = context.bot_data["settings"]
    store: ThreadMemoryStore = context.bot_data["thread_store"]
    user = update.effective_user

    if user is None or not _is_allowed(user.id, settings):
        return

    chat = update.effective_chat
    if chat is None or not _is_chat_enabled(chat, settings):
        return

    # If the user supplied a keyword query, search the knowledge base
    args: list[str] = context.args or []  # type: ignore[assignment]
    if args:
        query = " ".join(args).strip()
        results = store.search(query)
        if not results:
            await update.message.reply_text(  # type: ignore[union-attr]
                f"No memories found matching *{query}*. "
                "Try a different keyword or use /flush to add more summaries.",
                parse_mode=ParseMode.MARKDOWN,
            )
            return
        lines = [f"\U0001f50d *Knowledge base search: {query}*\n"]
        for s in results:
            date = s.flushed_at[:10]
            first_line = s.summary.split("\n")[0][:80]
            lines.append(f"\u2022 Thread `{s.thread_id}` ({date}): {first_line}\u2026")
            if s.tags:
                lines.append(f"  \U0001f3f7 {', '.join(s.tags)}")
        await _send_long(update, "\n".join(lines))
        return

    key = _thread_key(update)
    if key is not None:
        chat_id, thread_id = key
        stored = store.get(chat_id, thread_id)
        if stored is None:
            await update.message.reply_text(  # type: ignore[union-attr]
                "No stored summary for this thread yet. Use /flush to create one."
            )
        else:
            await _send_long(update, stored.format_for_telegram())
        return

    # Outside a thread: list all stored summaries
    all_summaries = store.all()
    if not all_summaries:
        await update.message.reply_text(  # type: ignore[union-attr]
            "No thread summaries stored yet. Use /flush inside a message thread."
        )
        return

    lines = ["\U0001f4da *Stored thread summaries*\n"]
    for s in all_summaries:
        date = s.flushed_at[:10]
        first_line = s.summary.split("\n")[0][:80]
        lines.append(f"\u2022 Thread `{s.thread_id}` ({date}): {first_line}\u2026")
        if s.tags:
            lines.append(f"  \U0001f3f7 {', '.join(s.tags)}")
    await _send_long(update, "\n".join(lines))


async def analyse_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    """Handle /analyse – run the proposal generator on demand."""
    settings: Settings = context.bot_data["settings"]
    llm: LLMClient = context.bot_data["llm"]
    user = update.effective_user
    if user is None or not _is_allowed(user.id, settings):
        return

    chat = update.effective_chat
    if chat is None or not _is_chat_enabled(chat, settings):
        return

    await update.message.reply_text("Running analysis, please wait...")  # type: ignore[union-attr]
    generator = ProposalGenerator(settings, llm)
    proposal = await generator.run()
    if proposal is None:
        await update.message.reply_text(  # type: ignore[union-attr]
            "Analysis could not be completed. Check that `ANALYSIS_TARGET_URL` is configured."
        )
        return

    await _send_long(update, proposal.format_for_telegram())


def _with_kb_context(
    history: list[dict[str, Any]],
    relevant: list[ThreadSummary],
) -> list[dict[str, Any]]:
    """Prepend relevant knowledge-base summaries as a transient system context message.

    The returned list is a *new* list — the original *history* is not mutated.
    The injected message is never appended to the stored history, so it does not
    permanently consume the context window.
    """
    if not relevant:
        return history
    snippets = [f"[Thread {s.thread_id}] {s.summary[:400]}" for s in relevant[:3]]
    kb_msg = _KB_CONTEXT_HEADER + "\n\n" + "\n\n---\n\n".join(snippets)
    return [{"role": "system", "content": kb_msg}, *history]


async def message_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    """Handle plain text messages – forward to LLM and reply.

    Thread messages: history is stored unbounded under the (chat_id, thread_id) key.
    Non-thread messages: history is capped at _MAX_HISTORY turns per user.

    Before each LLM call the knowledge base is searched for summaries whose tags
    overlap with keywords in the current message.  Any matches are injected as
    transient context — they are NOT stored in the rolling history, so they do
    not permanently consume the context window.

    If a :class:`~steward.tools.client.ToolClient` is registered in
    ``context.bot_data``, the LLM is invoked with tool calling support so it
    can take actions on the configured MCP/OpenAPI tool server.
    """
    settings: Settings = context.bot_data["settings"]
    llm: LLMClient = context.bot_data["llm"]
    store: ThreadMemoryStore = context.bot_data["thread_store"]
    tool_client: ToolClient | None = context.bot_data.get("tool_client")
    user = update.effective_user
    chat = update.effective_chat
    if user is None or not _is_allowed(user.id, settings):
        return
    if chat is None or not _is_chat_enabled(chat, settings):
        return

    text = update.message.text  # type: ignore[union-attr]
    if not text:
        return

    key = _thread_key(update)
    if key is not None:
        # Thread message: unbounded history
        history: list[dict[str, Any]] = _thread_history[key]
        call_history = _with_kb_context(history, store.search(text))
        if tool_client is not None:
            reply = await llm.chat_with_tools(text, tool_client, history=call_history)
        else:
            reply = await llm.chat(text, history=call_history)
        history.append({"role": "user", "content": text})
        history.append({"role": "assistant", "content": reply})
    else:
        # Non-thread message: capped history per user
        user_history: list[dict[str, Any]] = _history[user.id]
        call_history = _with_kb_context(user_history, store.search(text))
        if tool_client is not None:
            reply = await llm.chat_with_tools(text, tool_client, history=call_history)
        else:
            reply = await llm.chat(text, history=call_history)
        user_history.append({"role": "user", "content": text})
        user_history.append({"role": "assistant", "content": reply})
        if len(user_history) > _MAX_HISTORY * 2:
            _history[user.id] = user_history[-(_MAX_HISTORY * 2) :]

    await _send_long(update, reply)


def build_application(
    settings: Settings,
    llm: LLMClient,
    thread_store: ThreadMemoryStore | None = None,
    tool_client: ToolClient | None = None,
) -> Application:  # type: ignore[type-arg]
    """Build and return the Telegram Application."""
    app = Application.builder().token(settings.telegram_bot_token).build()
    app.bot_data["settings"] = settings
    app.bot_data["llm"] = llm
    app.bot_data["thread_store"] = thread_store or ThreadMemoryStore(settings.thread_memory_path)
    app.bot_data["tool_client"] = tool_client  # None when tools are not configured
    logger.info("Configured allowed users: %s", settings.telegram_allowed_user_ids)
    logger.info("Configured group IDs: %s", settings.telegram_group_ids)

    app.add_handler(CommandHandler("start", start_handler))
    app.add_handler(CommandHandler("help", help_handler))
    app.add_handler(CommandHandler("clear", clear_handler))
    app.add_handler(CommandHandler("flush", flush_handler))
    app.add_handler(CommandHandler("recall", recall_handler))
    app.add_handler(CommandHandler("analyse", analyse_handler))
    app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, message_handler))

    return app


async def send_proposal(app: Application, proposal: Proposal, user_ids: list[int]) -> None:  # type: ignore[type-arg]
    """Send a proposal message to all configured user IDs."""
    text = proposal.format_for_telegram()
    for uid in user_ids:
        try:
            await app.bot.send_message(
                chat_id=uid,
                text=text,
                parse_mode=ParseMode.MARKDOWN,
            )
        except Exception:
            logger.exception("Failed to send proposal to user %d", uid)