bot.py

#!/usr/bin/env python3
"""
Homebase Support Agent -- AI-powered Discord support bot for Tezos Homebase.

Listens in a designated channel for messages from users with a specific role.
Each user gets their own Discord thread with isolated conversation history.
Claude (via the ATN bridge) answers questions using tool-assisted lookups
against local documentation repos.

Usage:
    python bot.py

Configuration is read from environment variables (see .env.example).
"""
import asyncio
import fnmatch
import logging
import os
import re
import sys
import time
from collections import defaultdict
from pathlib import Path

import discord
from dotenv import load_dotenv

load_dotenv()

# ---------------------------------------------------------------------------
# ATN bridge provider -- path configurable via env
# ---------------------------------------------------------------------------
ATN_PATH = os.environ.get("ATN_PATH", "")
if ATN_PATH:
    sys.path.insert(0, ATN_PATH)

from atn.providers.bridge import BridgeProvider  # noqa: E402

# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------
logging.basicConfig(
    level=os.environ.get("LOG_LEVEL", "INFO").upper(),
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logging.getLogger("discord").setLevel(logging.WARNING)
logging.getLogger("discord.gateway").setLevel(logging.WARNING)
logging.getLogger("asyncio").setLevel(logging.WARNING)
log = logging.getLogger("homebase")

# ---------------------------------------------------------------------------
# Config (all from environment)
# ---------------------------------------------------------------------------
TOKEN = os.environ["DISCORD_BOT_TOKEN"]
CHANNEL_ID = int(os.environ["DISCORD_CHANNEL_ID"])
REQUIRED_ROLE = os.environ.get("DISCORD_REQUIRED_ROLE", "webapp-user")
ESCALATION_USER_ID = int(os.environ.get("ESCALATION_USER_ID", "0"))
MODEL = os.environ.get("MODEL", "sonnet")

MAX_HISTORY_TURNS = int(os.environ.get("MAX_HISTORY_TURNS", "50"))
MAX_TURNS = int(os.environ.get("MAX_TURNS", "20"))
MAX_TURNS_PER_DAY = int(os.environ.get("MAX_TURNS_PER_DAY", "20"))
RATE_WINDOW_SECS = int(os.environ.get("RATE_WINDOW_SECS", "86400"))
MAX_FILE_CHARS = int(os.environ.get("MAX_FILE_CHARS", "12000"))

DISCORD_MSG_LIMIT = 2000

# Live status embed
STATUS_UPDATE_INTERVAL = 1.0
EMBED_COLOR = 0x5865F2        # Blurple
EMBED_COLOR_ERROR = 0xED4245  # Red

# Documentation repos
DOCS_ROOT = Path(os.environ.get("DOCS_ROOT", str(Path(__file__).resolve().parent / "docs")))
REPO_NAMES = [s.strip() for s in os.environ.get("DOC_REPOS", "").split(",") if s.strip()]
REPOS: dict[str, Path] = {name: DOCS_ROOT / name for name in REPO_NAMES}

# ---------------------------------------------------------------------------
# System prompt
# ---------------------------------------------------------------------------
SYSTEM_PROMPT = os.environ.get("SYSTEM_PROMPT", """\
You are a support agent for Tezos Homebase (tezos-homebase.io), a web \
application for creating and managing DAOs on the Tezos blockchain. It is \
built on the BaseDAO smart contract framework.

There are three DAO templates: Treasury, Registry, and Lambda. The governance \
cycle has two phases: proposal period and voting period. Token holders freeze \
tokens to vote.

You have tools to search and read documentation from local repos. \
Use these tools when you need specifics about how something works -- contract \
error codes, governance parameters, DAO creation flow, etc. Don't guess \
at details you can look up.

You also have an **escalate** tool. Use it when:
- The user's issue requires human intervention (account problems, stuck \
transactions, suspected bugs you can't diagnose from docs alone).
- You've exhausted what you can help with and the user is still stuck.
- The user explicitly asks to speak to a human.
Do NOT escalate for questions you can answer from the docs.

Keep responses concise. This is a Discord support thread.

For bugs, direct users to: https://github.com/dOrgTech/homebase-app/issues
You cannot perform transactions or access wallets.
""")

# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------
_repo_enum = list(REPOS.keys()) + ["all"]

TOOLS = [
    {
        "name": "search_docs",
        "description": (
            "Search across documentation and source code for a query string. "
            "Returns matching lines with file paths. Use this to find relevant "
            "files before reading them."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Text or regex pattern to search for.",
                },
                "repo": {
                    "type": "string",
                    "enum": _repo_enum,
                    "description": "Which repo to search. Default: all.",
                },
                "file_pattern": {
                    "type": "string",
                    "description": (
                        "Optional glob to filter files, e.g. '*.md' for docs "
                        "only, '*.ts' for TypeScript. Default: all files."
                    ),
                },
            },
            "required": ["query"],
        },
    },
    {
        "name": "read_file",
        "description": (
            "Read a specific file from a documentation repo. Returns the file "
            "contents (truncated if very large). Use after search_docs to read "
            "files you've identified as relevant."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "repo": {
                    "type": "string",
                    "enum": list(REPOS.keys()),
                    "description": "Which repo the file is in.",
                },
                "path": {
                    "type": "string",
                    "description": "Path relative to the repo root.",
                },
            },
            "required": ["repo", "path"],
        },
    },
    {
        "name": "list_files",
        "description": (
            "List files in a directory of a documentation repo. Useful for "
            "discovering what documentation or source files exist."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "repo": {
                    "type": "string",
                    "enum": list(REPOS.keys()),
                    "description": "Which repo to list.",
                },
                "path": {
                    "type": "string",
                    "description": (
                        "Directory path relative to repo root. "
                        "Use '.' or '' for the root."
                    ),
                },
            },
            "required": ["repo"],
        },
    },
    {
        "name": "escalate",
        "description": (
            "Escalate the issue to a human support engineer. Use this when "
            "you cannot resolve the user's problem from documentation alone, "
            "when the issue requires human intervention, or when the user "
            "explicitly asks to talk to a person. Provide a brief summary "
            "of the issue and what you've already tried."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "reason": {
                    "type": "string",
                    "description": (
                        "Brief summary of the issue for the human engineer."
                    ),
                },
            },
            "required": ["reason"],
        },
    },
]

# ---------------------------------------------------------------------------
# Tool implementations
# ---------------------------------------------------------------------------

_SKIP_EXTS = {
    ".png", ".jpg", ".jpeg", ".gif", ".svg", ".ico", ".woff", ".woff2",
    ".ttf", ".eot", ".mp4", ".mp3", ".zip", ".gz", ".tar", ".lock",
    ".map", ".min.js", ".min.css", ".pyc", ".pyo", ".exe", ".dll",
    ".so", ".dylib", ".class", ".jar", ".bin", ".dat",
}
_SKIP_DIRS = {".git", "node_modules", "__pycache__", ".next", "dist", "build"}
_MAX_MATCHES_PER_REPO = 30
_MAX_LINE_LEN = 200


async def tool_executor(name: str, args: dict) -> dict:
    """Execute a tool call and return the result."""
    log.debug("TOOL CALL: %s(%s)", name, args)
    try:
        if name == "search_docs":
            result = _search_docs(
                args["query"],
                args.get("repo", "all"),
                args.get("file_pattern"),
            )
        elif name == "read_file":
            result = _read_file(args["repo"], args["path"])
        elif name == "list_files":
            result = _list_files(args["repo"], args.get("path", "."))
        else:
            result = {"error": f"Unknown tool: {name}"}
        log.debug("TOOL RESULT: %s -> %s", name, str(result)[:200])
        return result
    except Exception as e:
        log.exception("Tool %s failed", name)
        return {"error": str(e)}


def _search_docs(query: str, repo: str, file_pattern: str | None) -> dict:
    """Pure-Python recursive search across documentation repos."""
    repos_to_search = list(REPOS.keys()) if repo == "all" else [repo]
    all_matches: list[str] = []

    try:
        pattern = re.compile(query, re.IGNORECASE)
    except re.error:
        pattern = re.compile(re.escape(query), re.IGNORECASE)

    for repo_name in repos_to_search:
        repo_path = REPOS.get(repo_name)
        if not repo_path or not repo_path.exists():
            continue

        matches_this_repo = 0
        for filepath in repo_path.rglob("*"):
            if matches_this_repo >= _MAX_MATCHES_PER_REPO:
                break
            if filepath.is_dir():
                continue
            if any(p in _SKIP_DIRS for p in filepath.parts):
                continue
            if filepath.suffix.lower() in _SKIP_EXTS:
                continue
            if file_pattern and not fnmatch.fnmatch(filepath.name, file_pattern):
                continue

            try:
                text = filepath.read_text(encoding="utf-8", errors="ignore")
            except (OSError, PermissionError):
                continue

            hits_in_file = 0
            for line_no, line in enumerate(text.split("\n"), start=1):
                if hits_in_file >= 5:
                    break
                if pattern.search(line):
                    rel = str(filepath.relative_to(DOCS_ROOT)).replace("\\", "/")
                    snippet = line.strip()[:_MAX_LINE_LEN]
                    all_matches.append(f"{rel}:{line_no}: {snippet}")
                    hits_in_file += 1
                    matches_this_repo += 1
                    if matches_this_repo >= _MAX_MATCHES_PER_REPO:
                        break

    if not all_matches:
        return {"results": "No matches found."}
    return {"results": "\n".join(all_matches)}


def _read_file(repo: str, path: str) -> dict:
    repo_path = REPOS.get(repo)
    if not repo_path:
        return {"error": f"Unknown repo: {repo}"}

    file_path = repo_path / path
    try:
        file_path.resolve().relative_to(repo_path.resolve())
    except ValueError:
        return {"error": "Path traversal not allowed."}

    if not file_path.is_file():
        return {"error": f"File not found: {path}"}

    try:
        content = file_path.read_text(encoding="utf-8", errors="replace")
    except Exception as e:
        return {"error": f"Could not read file: {e}"}

    if len(content) > MAX_FILE_CHARS:
        content = content[:MAX_FILE_CHARS] + f"\n\n[truncated at {MAX_FILE_CHARS} chars]"
    return {"content": content}


def _list_files(repo: str, path: str) -> dict:
    repo_path = REPOS.get(repo)
    if not repo_path:
        return {"error": f"Unknown repo: {repo}"}

    dir_path = repo_path / (path or ".")
    try:
        dir_path.resolve().relative_to(repo_path.resolve())
    except ValueError:
        return {"error": "Path traversal not allowed."}

    if not dir_path.is_dir():
        return {"error": f"Directory not found: {path}"}

    entries: list[str] = []
    for item in sorted(dir_path.iterdir()):
        if item.name.startswith("."):
            continue
        suffix = "/" if item.is_dir() else ""
        entries.append(f"{item.name}{suffix}")
    return {"entries": entries}


async def _escalate(channel: discord.abc.Messageable, reason: str) -> dict:
    """Ping the human support engineer in the thread."""
    if not ESCALATION_USER_ID:
        return {"error": "No escalation user configured."}
    mention = f"<@{ESCALATION_USER_ID}>"
    escalation_msg = f"{mention} **Escalation** -- {reason}"
    try:
        await channel.send(escalation_msg)
        log.info("Escalated to %s: %s", ESCALATION_USER_ID, reason[:80])
        return {"status": "escalated", "message": "A human engineer has been notified in this thread."}
    except Exception as e:
        log.exception("Failed to send escalation")
        return {"error": f"Failed to notify engineer: {e}"}


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def split_message(text: str, limit: int = DISCORD_MSG_LIMIT) -> list[str]:
    """Split text into chunks that fit within Discord's message limit."""
    if len(text) <= limit:
        return [text]

    chunks: list[str] = []
    while text:
        if len(text) <= limit:
            chunks.append(text)
            break
        cut = text.rfind("\n", 0, limit)
        if cut == -1 or cut < limit // 2:
            cut = text.rfind(" ", 0, limit)
        if cut == -1 or cut < limit // 2:
            cut = limit
        chunks.append(text[:cut])
        text = text[cut:].lstrip("\n")
    return chunks


# ---------------------------------------------------------------------------
# Live status embed
# ---------------------------------------------------------------------------

TOOL_LABELS = {
    "search_docs": "Searching docs",
    "read_file": "Reading file",
    "list_files": "Listing files",
    "escalate": "Escalating to human",
}


class StatusEmbed:
    """Live-updating Discord embed that shows what the bot is doing."""

    def __init__(self, channel: discord.abc.Messageable) -> None:
        self._channel = channel
        self._message: discord.Message | None = None
        self._tool_calls: list[str] = []
        self._current_tool: str | None = None
        self._current_detail: str | None = None
        self._last_update: float = 0.0

    async def send_initial(self) -> None:
        embed = discord.Embed(title="Thinking...", color=EMBED_COLOR)
        try:
            self._message = await self._channel.send(embed=embed)
        except Exception:
            log.debug("Failed to send status embed", exc_info=True)

    async def set_tool(self, name: str, detail: str | None = None) -> None:
        label = TOOL_LABELS.get(name, name)
        self._tool_calls.append(label)
        self._current_tool = label
        self._current_detail = detail
        await self._maybe_update()

    async def clear_tool(self) -> None:
        self._current_tool = None
        self._current_detail = None
        await self._maybe_update()

    async def finalize(self, text: str, tool_summary: bool = True) -> None:
        """Replace the embed with the final response text."""
        if self._tool_calls and tool_summary:
            seen: set[str] = set()
            unique: list[str] = []
            for t in self._tool_calls:
                if t not in seen:
                    seen.add(t)
                    unique.append(t)
            summary = ", ".join(unique)
            final = f"-# {summary}\n{text}"
        else:
            final = text

        if self._message:
            chunks = split_message(final)
            try:
                await self._message.edit(content=chunks[0], embed=None)
                for chunk in chunks[1:]:
                    await self._channel.send(chunk)
                return
            except Exception:
                log.debug("Failed to edit status into final", exc_info=True)
                try:
                    await self._message.delete()
                except Exception:
                    pass

        for chunk in split_message(final):
            await self._channel.send(chunk)

    async def _maybe_update(self) -> None:
        now = time.monotonic()
        if now - self._last_update < STATUS_UPDATE_INTERVAL:
            return
        self._last_update = now
        await self._update()

    async def _update(self) -> None:
        if not self._message:
            return

        if self._current_tool:
            title = f"Working... {self._current_tool}"
            desc = self._current_detail or None
        else:
            title = "Thinking..."
            desc = None

        embed = discord.Embed(title=title, color=EMBED_COLOR)
        if desc:
            embed.description = desc
        if self._tool_calls:
            display = self._tool_calls[-8:]
            embed.add_field(
                name="Steps so far",
                value=", ".join(display),
                inline=False,
            )

        try:
            await self._message.edit(embed=embed)
        except Exception:
            log.debug("Failed to update status embed", exc_info=True)


# ---------------------------------------------------------------------------
# Rate limiter
# ---------------------------------------------------------------------------

class RateLimiter:
    """Rolling-window turn counter per user."""

    def __init__(self, max_turns: int = MAX_TURNS_PER_DAY,
                 window_secs: int = RATE_WINDOW_SECS):
        self.max_turns = max_turns
        self.window_secs = window_secs
        self._usage: dict[str, list[float]] = defaultdict(list)

    def _prune(self, user_id: str) -> None:
        cutoff = time.time() - self.window_secs
        self._usage[user_id] = [
            t for t in self._usage[user_id] if t > cutoff
        ]

    def remaining(self, user_id: str) -> int:
        self._prune(user_id)
        return max(0, self.max_turns - len(self._usage[user_id]))

    def record(self, user_id: str) -> None:
        self._usage[user_id].append(time.time())


# ---------------------------------------------------------------------------
# Per-thread conversation state
# ---------------------------------------------------------------------------

class ThreadState:
    """Conversation history for a single support thread."""

    def __init__(self) -> None:
        self.history: list[dict[str, str]] = []

    def add_user(self, name: str, content: str) -> None:
        self.history.append({"role": "user", "name": name, "content": content})
        if len(self.history) > MAX_HISTORY_TURNS:
            del self.history[: len(self.history) - MAX_HISTORY_TURNS]

    def add_assistant(self, content: str) -> None:
        self.history.append({"role": "assistant", "content": content})

    def format(self) -> str:
        parts: list[str] = []
        for turn in self.history:
            if turn["role"] == "user":
                parts.append(f"[{turn['name']}]: {turn['content']}")
            else:
                parts.append(f"[Support Agent]: {turn['content']}")
        return "\n\n".join(parts)


# ---------------------------------------------------------------------------
# Bot
# ---------------------------------------------------------------------------

intents = discord.Intents.default()
intents.message_content = True
intents.members = True
client = discord.Client(intents=intents)

provider = BridgeProvider(model=MODEL)
rate_limiter = RateLimiter()
threads: dict[int, ThreadState] = {}
thread_locks: dict[int, asyncio.Lock] = {}
_processed_messages: set[int] = set()


def has_role(member: discord.Member, role_name: str) -> bool:
    return any(r.name == role_name for r in member.roles)


def get_thread_state(thread_id: int) -> ThreadState:
    if thread_id not in threads:
        threads[thread_id] = ThreadState()
    return threads[thread_id]


def get_thread_lock(thread_id: int) -> asyncio.Lock:
    if thread_id not in thread_locks:
        thread_locks[thread_id] = asyncio.Lock()
    return thread_locks[thread_id]


@client.event
async def on_ready():
    log.info("Connected as %s (channel: %s, role: %s, repos: %s)",
             client.user, CHANNEL_ID, REQUIRED_ROLE, list(REPOS.keys()))


@client.event
async def on_message(message: discord.Message):
    if message.author.bot:
        return

    # Deduplicate: Discord can fire on_message twice for thread starters
    if message.id in _processed_messages:
        return
    _processed_messages.add(message.id)
    if len(_processed_messages) > 1000:
        _processed_messages.clear()

    if not isinstance(message.author, discord.Member):
        return
    if not has_role(message.author, REQUIRED_ROLE):
        return

    channel = message.channel
    user_id = str(message.author.id)
    username = message.author.display_name or message.author.name
    user_text = message.content.strip()
    if not user_text:
        return

    # --- New question in the support channel -> create a thread ---
    if channel.id == CHANNEL_ID:
        remaining = rate_limiter.remaining(user_id)
        if remaining <= 0:
            await message.reply(
                "You've reached the daily limit for support questions. "
                "Please try again tomorrow."
            )
            return

        thread_name = f"{username}: {user_text[:50]}"
        thread = await message.create_thread(name=thread_name)
        await _handle_turn(thread, message, user_id, username, user_text,
                           thread.id)
        return

    # --- Follow-up inside an existing support thread ---
    if isinstance(channel, discord.Thread) and channel.parent_id == CHANNEL_ID:
        remaining = rate_limiter.remaining(user_id)
        if remaining <= 0:
            await channel.send(
                "You've reached the daily limit for support questions. "
                "Please try again tomorrow."
            )
            return

        await _handle_turn(channel, message, user_id, username, user_text,
                           channel.id)
        return


async def _handle_turn(
    channel: discord.Thread | discord.abc.Messageable,
    message: discord.Message,
    user_id: str,
    username: str,
    user_text: str,
    thread_id: int,
) -> None:
    """Process a single conversational turn."""
    state = get_thread_state(thread_id)
    lock = get_thread_lock(thread_id)

    log.info("[thread:%s] [%s] %s", thread_id, username, user_text[:80])

    state.add_user(username, user_text)
    formatted = state.format()

    status = StatusEmbed(channel)

    async def _executor(name: str, args: dict) -> dict:
        detail = _tool_detail(name, args)
        await status.set_tool(name, detail)

        if name == "escalate":
            result = await _escalate(channel, args.get("reason", ""))
        else:
            result = await tool_executor(name, args)

        await status.clear_tool()
        return result

    async with lock:
        await status.send_initial()
        try:
            response = await provider.send_orchestrate(
                message=formatted,
                system=SYSTEM_PROMPT,
                model=MODEL,
                tools=TOOLS,
                max_turns=MAX_TURNS,
                tool_executor=_executor,
            )
            reply = response.text.strip() if response.text else ""
            log.info("REPLY (%d chars): %s", len(reply), reply[:120])

            if not reply:
                log.warning("Empty reply from bridge (stop_reason: %s)",
                            getattr(response, "stop_reason", "unknown"))
                reply = (
                    "I investigated your question but wasn't able to put together "
                    "a complete answer. Could you try rephrasing, or would you "
                    "like me to escalate this to a human?"
                )
        except Exception:
            log.exception("Bridge call failed")
            reply = (
                "Sorry, I ran into a problem processing your question. "
                "Please try again in a moment."
            )

    rate_limiter.record(user_id)
    state.add_assistant(reply)

    remaining = rate_limiter.remaining(user_id)
    if remaining == 0:
        reply += "\n\n*You've reached your daily question limit. Your allocation resets in 24 hours.*"
    elif remaining <= 3:
        reply += f"\n\n*({remaining} questions remaining today)*"

    await status.finalize(reply)


def _tool_detail(name: str, args: dict) -> str | None:
    """Build a short human-readable detail string for a tool call."""
    if name == "search_docs":
        q = args.get("query", "")
        repo = args.get("repo", "all")
        return f'`{q}` in {repo}'
    if name == "read_file":
        return f'`{args.get("path", "")}`'
    if name == "list_files":
        return f'`{args.get("repo", "")}/{args.get("path", ".")}`'
    return None


if __name__ == "__main__":
    try:
        client.run(TOKEN)
    except KeyboardInterrupt:
        log.info("Shutting down")