#!/usr/bin/env python3
"""
Telegram <-> Claude Code Bridge — обезличенный пример.

Это упрощённая копия моего production-моста (там ~1200 строк с кастомной
логикой: handoff между сессиями, plan-маркеры, медиа-retention, разбор
сложных Telegram-entities). Здесь оставлены только ключевые паттерны —
~250 строк, чтобы можно было отдать этот файл Claude Code и сказать:
«адаптируй под моего агента» — он соберёт за час.

Ключевые паттерны (то, ради чего вообще стоит делать мост):

1. Per-agent isolation. Каждому агенту своя папка-проект, свой bot-токен,
   своя SQLite-запись. Агенты не пересекаются и не путают контекст.

2. Session resumption через `claude -p --resume <session_id>`. Без этого
   каждое сообщение боту = новая сессия = модель не помнит предыдущее.
   А с `--resume` это полноценный диалог с памятью.

3. Подписка вместо API. `claude` CLI берёт токен из env-переменной
   CLAUDE_CODE_OAUTH_TOKEN — это OAuth-токен подписки Pro/Max, а не API-ключ.
   Получить его: на машине, где `claude /login` уже выполнен, посмотреть
   ~/.claude/.credentials.json или через `claude auth token`.

4. Длинные ответы режутся на 4000-символьные куски (Telegram-лимит).

Что НЕ показано (намеренно — у вас будет своя реализация):
- aiogram middleware, FSM, фильтры
- TLS, прокси, обработка media
- Многоаккаунтные параллельные claude-сессии
- Контекст-watcher (предупреждение о приближении к лимиту окна)

Запуск:
    pip install aiogram
    export CLAUDE_CODE_OAUTH_TOKEN="sk-ant-oat01-..."
    export TELEGRAM_OWNER_ID="123456789"
    export AGENT_BOT_TOKEN_MAIN="..."
    python bridge_example.py
"""

import asyncio
import logging
import os
import sqlite3
import subprocess
import time
from dataclasses import dataclass, field
from pathlib import Path

from aiogram import Bot, Dispatcher
from aiogram.types import Message

# ----------------------------------------------------------------------------
# Config — всё личное вынесено в env-переменные
# ----------------------------------------------------------------------------

OWNER_ID = int(os.environ["TELEGRAM_OWNER_ID"])  # ваш Telegram user_id
MODEL = os.environ.get("CLAUDE_MODEL", "claude-sonnet-4-6")
TIMEOUT_SECONDS = int(os.environ.get("CLAUDE_TIMEOUT", "600"))
TYPING_INTERVAL = 4  # пинг "печатает..." каждые N секунд
TG_MAX_CHARS = 4000  # лимит одного сообщения Telegram

BRIDGE_DIR = Path(__file__).parent
DB_PATH = BRIDGE_DIR / "sessions.db"
LOG_PATH = BRIDGE_DIR / "bridge.log"

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(name)s] %(message)s",
    handlers=[logging.FileHandler(LOG_PATH), logging.StreamHandler()],
)
log = logging.getLogger("bridge")


# ----------------------------------------------------------------------------
# Описание агента
# ----------------------------------------------------------------------------


@dataclass
class AgentConfig:
    agent_id: str           # уникальный id, например "main"
    name: str               # отображаемое имя для логов
    bot_token: str          # токен Telegram-бота этого агента
    workspace: str          # путь к папке агента (cwd для claude -p)
    dm_allow: list[int] = field(default_factory=lambda: [OWNER_ID])


# Здесь объявляете своих агентов. Имена/токены — из env, никаких хардкодов.
AGENTS: dict[str, AgentConfig] = {
    "main": AgentConfig(
        agent_id="main",
        name="Main",
        bot_token=os.environ["AGENT_BOT_TOKEN_MAIN"],
        workspace=os.path.expanduser("~/agents/main"),
    ),
    # пример второго агента — раскомментируйте и заведите бота
    # "scout": AgentConfig(
    #     agent_id="scout",
    #     name="Scout",
    #     bot_token=os.environ["AGENT_BOT_TOKEN_SCOUT"],
    #     workspace=os.path.expanduser("~/agents/scout"),
    # ),
}


# ----------------------------------------------------------------------------
# SQLite — храним session_id для каждой пары (агент, чат)
# ----------------------------------------------------------------------------


def init_db() -> sqlite3.Connection:
    db = sqlite3.connect(DB_PATH)
    db.execute("""
        CREATE TABLE IF NOT EXISTS sessions (
            agent_id TEXT,
            chat_id INTEGER,
            session_id TEXT,
            updated_at REAL,
            PRIMARY KEY (agent_id, chat_id)
        )
    """)
    db.commit()
    return db


db = init_db()


def save_session(agent_id: str, chat_id: int, session_id: str) -> None:
    db.execute(
        "INSERT OR REPLACE INTO sessions VALUES (?, ?, ?, ?)",
        (agent_id, chat_id, session_id, time.time()),
    )
    db.commit()


def load_session(agent_id: str, chat_id: int) -> str | None:
    row = db.execute(
        "SELECT session_id FROM sessions WHERE agent_id=? AND chat_id=?",
        (agent_id, chat_id),
    ).fetchone()
    return row[0] if row else None


# ----------------------------------------------------------------------------
# Главная функция — спросить у claude -p
# ----------------------------------------------------------------------------


async def ask_claude(agent: AgentConfig, chat_id: int, prompt: str) -> str:
    """Запускает claude -p в папке агента, возвращает ответ.

    Используется --output-format json чтобы получить session_id и сохранить
    его. На следующем запросе используем --resume <session_id> — модель
    помнит предыдущий диалог.
    """
    session_id = load_session(agent.agent_id, chat_id)
    cmd = [
        "claude", "-p",
        "--output-format", "json",
        "--model", MODEL,
    ]
    if session_id:
        cmd += ["--resume", session_id]

    log.info("[%s/%s] %s -> %s", agent.agent_id, chat_id,
             session_id or "(new)", prompt[:80])

    proc = await asyncio.create_subprocess_exec(
        *cmd,
        cwd=agent.workspace,
        stdin=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        env={**os.environ},  # пробрасываем CLAUDE_CODE_OAUTH_TOKEN
    )

    try:
        stdout, stderr = await asyncio.wait_for(
            proc.communicate(input=prompt.encode("utf-8")),
            timeout=TIMEOUT_SECONDS,
        )
    except asyncio.TimeoutError:
        proc.kill()
        return f"⏱ Превышен таймаут {TIMEOUT_SECONDS}с. Попробуй задать вопрос проще."

    if proc.returncode != 0:
        log.error("claude exit %s, stderr: %s", proc.returncode, stderr[:500])
        return f"⚠️ Ошибка вызова Claude Code (exit {proc.returncode})."

    try:
        import json
        data = json.loads(stdout.decode("utf-8"))
        new_session_id = data.get("session_id")
        if new_session_id:
            save_session(agent.agent_id, chat_id, new_session_id)
        return data.get("result") or "(пустой ответ)"
    except Exception as e:
        log.exception("parse error")
        return f"⚠️ Не смог разобрать ответ Claude Code: {e}"


# ----------------------------------------------------------------------------
# Telegram-обработчики
# ----------------------------------------------------------------------------


def split_long(text: str, limit: int = TG_MAX_CHARS) -> list[str]:
    """Делит длинный текст на куски по limit символов, по границам строк."""
    if len(text) <= limit:
        return [text]
    chunks, current = [], ""
    for line in text.split("\n"):
        if len(current) + len(line) + 1 > limit:
            chunks.append(current)
            current = line
        else:
            current = (current + "\n" + line) if current else line
    if current:
        chunks.append(current)
    return chunks


async def typing_loop(bot: Bot, chat_id: int, stop: asyncio.Event) -> None:
    """Шлёт "печатает..." пока stop не выставлен — Telegram забывает за 5с."""
    while not stop.is_set():
        try:
            await bot.send_chat_action(chat_id, "typing")
        except Exception:
            pass
        try:
            await asyncio.wait_for(stop.wait(), timeout=TYPING_INTERVAL)
        except asyncio.TimeoutError:
            pass


def make_handler(agent: AgentConfig):
    async def on_message(msg: Message) -> None:
        # ACL — только владелец и явно разрешённые
        if msg.from_user is None or msg.from_user.id not in agent.dm_allow:
            return
        if not msg.text:
            await msg.reply("Пока умею только текст.")
            return

        bot: Bot = msg.bot
        stop_typing = asyncio.Event()
        typing_task = asyncio.create_task(typing_loop(bot, msg.chat.id, stop_typing))
        try:
            answer = await ask_claude(agent, msg.chat.id, msg.text)
        finally:
            stop_typing.set()
            await typing_task

        for chunk in split_long(answer):
            await msg.answer(chunk)

    return on_message


# ----------------------------------------------------------------------------
# Запуск
# ----------------------------------------------------------------------------


async def main() -> None:
    if "CLAUDE_CODE_OAUTH_TOKEN" not in os.environ:
        raise SystemExit("CLAUDE_CODE_OAUTH_TOKEN не установлен — войдите через `claude /login`")

    bots = []
    for agent in AGENTS.values():
        bot = Bot(token=agent.bot_token)
        dp = Dispatcher()
        dp.message()(make_handler(agent))
        bots.append(asyncio.create_task(dp.start_polling(bot)))
        log.info("Started bot for agent %s (workspace=%s)", agent.name, agent.workspace)

    await asyncio.gather(*bots)


if __name__ == "__main__":
    asyncio.run(main())
