#!/usr/bin/env python3
"""
Telegram MCP Proxy — stdio MCP сервер для Claude Code.
========================================================
Тонкий прокси: регистрируется в демоне (daemon.py),
проксирует инструменты через HTTP API демона.

При первом вызове любого tool:
  1. Подключается к демону на localhost:DAEMON_PORT
  2. Регистрирует новую сессию (получает свой тред в Telegram)
  3. Запускает фоновый polling входящих сообщений от демона

Real-time мониторинг:
  - MCP Sampling: при получении сообщения proxy запрашивает ответ у Claude
  - Fallback: tool monitor_chat() блокируется до прихода сообщения

Конфиг env:
  DAEMON_PORT — порт демона (по умолчанию 8765)
  SESSION_NAME — имя сессии для треда (опционально)
"""

import io
import os
import sys
import asyncio
import time
import signal
import traceback

import httpx
from mcp.server.fastmcp import FastMCP, Context
from mcp.types import SamplingMessage, TextContent

# --- Логирование в файл с ротацией ---
LOG_FILE = "/tmp/tg_mcp_proxy.log"
LOG_MAX_SIZE = 5 * 1024 * 1024  # 5 MB


def log(msg: str):
    """Лог в файл с ротацией — stderr может быть недоступен в Claude Code."""
    ts = time.strftime("%H:%M:%S")
    line = f"[{ts}] {msg}\n"
    try:
        if os.path.exists(LOG_FILE) and os.path.getsize(LOG_FILE) > LOG_MAX_SIZE:
            backup = LOG_FILE + ".1"
            if os.path.exists(backup):
                os.remove(backup)
            os.rename(LOG_FILE, backup)
        with open(LOG_FILE, "a") as f:
            f.write(line)
    except Exception:
        pass
    print(f"[proxy] {msg}", file=sys.stderr)

log(f"=== Прокси запускается, PID={os.getpid()} ===")
log(f"Python: {sys.executable} {sys.version}")
log(f"argv: {sys.argv}")
log(f"cwd: {os.getcwd()}")

DAEMON_PORT = int(os.environ.get("DAEMON_PORT", "8765"))
DAEMON_URL = f"http://127.0.0.1:{DAEMON_PORT}"
def _build_session_name() -> str:
    """Имя сессии: название проекта (из cwd) + номер сессии."""
    if os.environ.get("SESSION_NAME"):
        return os.environ["SESSION_NAME"]
    project = os.path.basename(os.getcwd()) or "unknown"
    return f"{project} #{os.getpid()} — {time.strftime('%H:%M %d.%m')}"

SESSION_NAME = _build_session_name()
POLL_INTERVAL = 2

# Состояние прокси
session_id: str | None = None
thread_id: int | None = None
_inbox: asyncio.Queue | None = None
_registered = False
_poll_task: asyncio.Task | None = None
_http_client: httpx.AsyncClient | None = None
_session = None  # ServerSession — захватывается из первого tool call
_sampling_supported = False  # True если Claude Code поддерживает sampling


async def daemon_request(method: str, path: str, params: dict | None = None) -> dict:
    """HTTP запрос к демону. Использует общий клиент."""
    global _http_client
    if _http_client is None:
        _http_client = httpx.AsyncClient(timeout=60)

    if method == "GET":
        resp = await _http_client.get(f"{DAEMON_URL}{path}", params=params)
    else:
        resp = await _http_client.post(f"{DAEMON_URL}{path}", json=params or {})
    resp.raise_for_status()
    return resp.json()


async def ensure_registered():
    """Lazy-регистрация в демоне при первом вызове tool."""
    global session_id, thread_id, _registered

    if _registered:
        return

    log("Регистрация в демоне...")
    try:
        await daemon_request("GET", "/api/health")
        log("Демон доступен")
    except Exception as e:
        log(f"Демон недоступен: {e}")
        raise RuntimeError(
            "Демон не запущен! "
            "Запустите: python -m telegram_mcp_bridge.daemon"
        )

    result = await daemon_request("POST", "/api/register", {"name": SESSION_NAME})
    session_id = result["session_id"]
    thread_id = result["thread_id"]
    _registered = True
    log(f"Зарегистрирован: sid={session_id}, тред={thread_id}")


def _capture_session(ctx: Context):
    """Захватить ServerSession из Context для sampling."""
    global _session, _sampling_supported
    if _session is None:
        _session = ctx.session
        # Проверяем поддержку sampling через client capabilities
        try:
            caps = getattr(_session, '_client_params', None)
            if caps:
                log(f"Client capabilities: {caps}")
            init_result = getattr(_session, '_initialization_state', None)
            log(f"Session init state: {init_result}")
            client_caps = getattr(_session, '_client_capabilities', None)
            if client_caps:
                log(f"Client capabilities obj: {client_caps}")
                sampling = getattr(client_caps, 'sampling', None)
                _sampling_supported = sampling is not None
                log(f"Sampling supported: {_sampling_supported}")
            else:
                log("Client capabilities не найдены, пробуем sampling при первом сообщении")
                _sampling_supported = True  # Попробуем, ошибку поймаем
        except Exception as e:
            log(f"Ошибка проверки capabilities: {e}")
            _sampling_supported = True  # Попробуем
        log(f"Session захвачен: {type(_session).__name__}, sampling={_sampling_supported}")


async def unregister():
    """Отключение от демона при завершении."""
    if _registered and session_id:
        try:
            await daemon_request("POST", "/api/unregister", {"sid": session_id})
            log("Отключён от демона")
        except Exception:
            pass
    if _http_client:
        await _http_client.aclose()


async def _handle_telegram_message(msg: dict):
    """Через MCP Sampling запросить ответ у Claude и отправить в Telegram."""
    global _sampling_supported

    text = msg.get("text", "")
    sender = msg.get("from", "user")
    if not text or _session is None or not _sampling_supported:
        return False

    try:
        log(f"Sampling запрос: [{sender}]: {text[:80]}...")
        result = await _session.create_message(
            messages=[
                SamplingMessage(
                    role="user",
                    content=TextContent(type="text", text=f"[Telegram от {sender}]: {text}")
                )
            ],
            max_tokens=4096,
            system_prompt=(
                "Ты помощник разработчика. Отвечай кратко на русском. "
                "Это сообщение из Telegram-треда твоей рабочей сессии."
            ),
            include_context="allServers",
        )

        # Извлечь текст ответа
        if hasattr(result.content, 'text'):
            response_text = result.content.text
        elif isinstance(result.content, list):
            response_text = " ".join(
                c.text for c in result.content if hasattr(c, 'text')
            )
        else:
            response_text = str(result.content)

        await daemon_request("POST", "/api/send", {
            "sid": session_id,
            "text": response_text
        })
        log(f"Sampling ответ отправлен: {response_text[:80]}...")
        return True

    except Exception as e:
        error_msg = str(e)
        log(f"Sampling ошибка: {error_msg}")
        if "not supported" in error_msg.lower() or "not implemented" in error_msg.lower():
            _sampling_supported = False
            log("Sampling отключён — клиент не поддерживает")
        return False


async def poll_incoming():
    """Фоновый опрос демона на входящие сообщения от пользователя."""
    global _inbox
    _inbox = asyncio.Queue()

    while True:
        if not _registered:
            await asyncio.sleep(1)
            continue
        try:
            result = await daemon_request("GET", f"/api/poll/{session_id}", {"timeout": "25"})
            for msg in result.get("messages", []):
                # Попробовать sampling (автоответ)
                await _handle_telegram_message(msg)
                # Всегда кладём в inbox (для check_inbox / monitor_chat / send_and_wait)
                await _inbox.put(msg)
        except asyncio.CancelledError:
            break
        except Exception as e:
            log(f"Ошибка polling: {e}")
            await asyncio.sleep(5)


def _ensure_polling():
    """Запуск polling task в текущем event loop (lazy)."""
    global _poll_task
    if _poll_task is None or _poll_task.done():
        _poll_task = asyncio.get_running_loop().create_task(poll_incoming())


# NO lifespan — lazy init on first tool call to avoid blocking MCP handshake

def _sync_unregister():
    """Синхронная отписка от демона — вызывается при завершении."""
    if not _registered or not session_id:
        return
    try:
        with httpx.Client(timeout=5) as c:
            c.post(f"{DAEMON_URL}/api/unregister", json={"sid": session_id})
        log("Отключён от демона (sync)")
    except Exception as e:
        log(f"Ошибка отключения: {e}")


def _signal_handler(signum, frame):
    log(f"Получен сигнал {signum}, завершаемся...")
    _sync_unregister()
    sys.exit(0)

signal.signal(signal.SIGTERM, _signal_handler)
signal.signal(signal.SIGINT, _signal_handler)

log("Создаём FastMCP без lifespan")
mcp = FastMCP("telegram-bridge")


# --- MCP Tools ---

@mcp.tool()
async def send_and_wait(message: str, ctx: Context) -> str:
    """
    Отправляет сообщение пользователю в Telegram и ждёт его ответа.
    Используй когда нужно задать вопрос, показать результат и получить фидбек,
    или уточнить требования перед продолжением работы.

    Args:
        message: Текст сообщения для отправки пользователю
    Returns:
        Текст ответа пользователя
    """
    log(f"send_and_wait вызван: {message[:50]}...")
    _capture_session(ctx)
    await ensure_registered()
    _ensure_polling()

    await daemon_request("POST", "/api/send", {"sid": session_id, "text": message})

    try:
        msg = await asyncio.wait_for(_inbox.get(), timeout=300)
        return msg.get("text", "[пустое сообщение]")
    except asyncio.TimeoutError:
        return "[Таймаут: ответ не получен]"


@mcp.tool()
async def send_message(message: str, ctx: Context) -> str:
    """
    Отправляет сообщение пользователю в Telegram БЕЗ ожидания ответа.
    Используй для уведомлений, отчётов о прогрессе, или финальных результатов.

    Args:
        message: Текст сообщения для отправки
    Returns:
        Подтверждение отправки
    """
    log(f"send_message вызван: {message[:50]}...")
    _capture_session(ctx)
    await ensure_registered()
    _ensure_polling()

    result = await daemon_request("POST", "/api/send", {"sid": session_id, "text": message})
    return f"Сообщение отправлено (id: {result['message_id']})"


@mcp.tool()
async def send_and_wait_with_options(message: str, options: list[str], ctx: Context) -> str:
    """
    Отправляет сообщение с кнопками выбора и ждёт ответа.
    Используй когда нужно предложить пользователю варианты действий.

    Args:
        message: Текст вопроса
        options: Список вариантов для кнопок (например: ["Да", "Нет", "Пропустить"])
    Returns:
        Текст выбранного варианта или произвольный ответ пользователя
    """
    log(f"send_and_wait_with_options вызван: {message[:50]}...")
    _capture_session(ctx)
    await ensure_registered()
    _ensure_polling()

    await daemon_request("POST", "/api/send_with_options", {
        "sid": session_id,
        "text": message,
        "options": options,
    })

    try:
        msg = await asyncio.wait_for(_inbox.get(), timeout=300)
        return msg.get("text", "[пустое сообщение]")
    except asyncio.TimeoutError:
        return "[Таймаут: ответ не получен]"


@mcp.tool()
async def send_file(file_path: str, caption: str = "", ctx: Context = None) -> str:
    """
    Отправляет файл в Telegram-тред сессии.
    Изображения (png/jpg/gif/webp) отправляются как фото, остальное — как документ.

    Args:
        file_path: Абсолютный путь к файлу на диске
        caption: Подпись к файлу (опционально)
    Returns:
        Подтверждение отправки
    """
    log(f"send_file вызван: {file_path}")
    if ctx:
        _capture_session(ctx)
    await ensure_registered()
    _ensure_polling()

    if not os.path.isfile(file_path):
        return f"Ошибка: файл не найден: {file_path}"

    filename = os.path.basename(file_path)
    with open(file_path, "rb") as f:
        file_bytes = f.read()

    # multipart upload к демону
    global _http_client
    if _http_client is None:
        _http_client = httpx.AsyncClient(timeout=120)

    resp = await _http_client.post(
        f"{DAEMON_URL}/api/send_file",
        data={"sid": session_id, "caption": caption},
        files={"file": (filename, io.BytesIO(file_bytes), "application/octet-stream")},
    )
    resp.raise_for_status()
    result = resp.json()
    return f"Файл отправлен (id: {result['message_id']})"


@mcp.tool()
async def check_inbox(ctx: Context) -> str:
    """
    Проверяет входящие сообщения от пользователя из Telegram.
    Пользователь может писать в свой тред в любое время.
    Вызывай эту функцию периодически или когда пользователь
    упоминает что отправил что-то в Telegram.

    Returns:
        Список непрочитанных сообщений или "Нет новых сообщений"
    """
    log("check_inbox вызван")
    _capture_session(ctx)
    await ensure_registered()
    _ensure_polling()

    messages = []
    if _inbox is not None:
        while not _inbox.empty():
            try:
                msg = _inbox.get_nowait()
                messages.append(msg)
            except asyncio.QueueEmpty:
                break

    if not messages:
        return "Нет новых сообщений"

    lines = []
    for msg in messages:
        sender = msg.get("from", "?")
        text = msg.get("text", "")
        lines.append(f"[{sender}]: {text}")
    return "\n".join(lines)


@mcp.tool()
async def monitor_chat(ctx: Context) -> str:
    """
    Мониторинг Telegram-треда в реальном времени.
    Блокируется до получения следующего сообщения (таймаут 30 мин).

    Используй для непрерывного общения через Telegram:
    1. Вызови monitor_chat() — ждёт сообщения
    2. Обработай полученное сообщение
    3. Ответь через send_message()
    4. Вызови monitor_chat() снова

    Returns:
        Текст полученного сообщения или таймаут
    """
    log("monitor_chat вызван")
    _capture_session(ctx)
    await ensure_registered()
    _ensure_polling()

    try:
        msg = await asyncio.wait_for(_inbox.get(), timeout=1800)
        sender = msg.get("from", "?")
        text = msg.get("text", "")
        log(f"monitor_chat получил: [{sender}]: {text[:50]}...")
        return (
            f"[Telegram от {sender}]: {text}\n\n"
            "Ответь через send_message(), затем вызови monitor_chat() для продолжения."
        )
    except asyncio.TimeoutError:
        return (
            "[Таймаут 30 мин — нет новых сообщений]\n"
            "Вызови monitor_chat() снова если хочешь продолжить мониторинг."
        )


@mcp.tool()
async def list_sessions(ctx: Context) -> str:
    """
    Показывает все активные сессии Claude Code, подключённые к демону.
    """
    log("list_sessions вызван")
    _capture_session(ctx)
    await ensure_registered()

    result = await daemon_request("GET", "/api/sessions")
    if not result:
        return "Нет активных сессий"

    lines = []
    for s in result:
        marker = " (текущая)" if s["session_id"] == session_id else ""
        pending = f", ожидает: {s['pending_messages']}" if s["pending_messages"] else ""
        lines.append(f"- {s['name']}, тред={s['thread_id']}{pending}{marker}")
    return "\n".join(lines)


def main():
    log("Запуск mcp.run(transport='stdio')")
    try:
        mcp.run(transport="stdio")
    except Exception as e:
        log(f"FATAL: {e}\n{traceback.format_exc()}")
        raise
    finally:
        _sync_unregister()
        log("Процесс завершается")


if __name__ == "__main__":
    main()
