#!/usr/bin/env python3
"""Odysseus ↔ Telegram Bot Bridge.

Bridges Telegram messages to the Odysseus chat API, enabling conversational AI
through Telegram. Supports per-user sessions, web search toggle, and streaming.
"""

import os
import sys
import re
import json
import time
import logging
import asyncio
import httpx
from typing import Optional

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [TG-BOT] %(levelname)s %(message)s",
)
logger = logging.getLogger("telegram_bot")

# ── Config ──────────────────────────────────────────────────────────────────
TELEGRAM_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN", "")
ODYSSEUS_URL = os.getenv("ODYSSEUS_INTERNAL_URL", "http://127.0.0.1:7860")
ALLOWED_USERS = os.getenv("TELEGRAM_ALLOWED_USERS", "").split(",")
ALLOWED_USERS = [u.strip() for u in ALLOWED_USERS if u.strip()]

# Per-user sessions stored in memory (reset on restart)
user_sessions: dict[str, str] = {}

TG_API = f"https://api.telegram.org/bot{TELEGRAM_TOKEN}"


def get_api_token() -> str:
    """Get Odysseus API token from env or file (written by auto_configure)."""
    token = os.getenv("ODYSSEUS_API_TOKEN", "")
    if token:
        return token
    try:
        with open("/tmp/tg_api_token", "r") as f:
            return f.read().strip()
    except FileNotFoundError:
        return ""


async def tg_request(method: str, data: dict, client: httpx.AsyncClient) -> dict:
    """Make a Telegram Bot API request."""
    resp = await client.post(f"{TG_API}/{method}", json=data, timeout=30)
    result = resp.json()
    if not result.get("ok"):
        logger.error(f"Telegram API error on {method}: {result}")
    return result


async def send_message(chat_id: int, text: str, client: httpx.AsyncClient,
                       reply_to: Optional[int] = None, parse_mode: str = "Markdown"):
    """Send a message, splitting if > 4096 chars."""
    chunks = [text[i:i+4000] for i in range(0, len(text), 4000)]
    for chunk in chunks:
        data = {"chat_id": chat_id, "text": chunk}
        if parse_mode:
            data["parse_mode"] = parse_mode
        if reply_to:
            data["reply_to_message_id"] = reply_to
        try:
            await tg_request("sendMessage", data, client)
        except Exception:
            # Retry without parse_mode if markdown fails
            data.pop("parse_mode", None)
            try:
                await tg_request("sendMessage", data, client)
            except Exception as e:
                logger.error(f"Failed to send message: {e}")


async def send_typing(chat_id: int, client: httpx.AsyncClient):
    """Send typing indicator."""
    try:
        await tg_request("sendChatAction", {"chat_id": chat_id, "action": "typing"}, client)
    except Exception:
        pass


async def get_or_create_session(user_id: str, client: httpx.AsyncClient) -> str:
    """Get existing session or create a new one for this Telegram user."""
    api_token = get_api_token()
    headers = {"Authorization": f"Bearer {api_token}"} if api_token else {}

    if user_id in user_sessions:
        return user_sessions[user_id]

    # Create a new Odysseus session
    try:
        resp = await client.post(
            f"{ODYSSEUS_URL}/api/session",
            json={
                "name": f"Telegram-{user_id}",
                "endpoint_url": "",
                "model": "",
            },
            headers=headers,
            timeout=15,
        )
        if resp.status_code == 200:
            data = resp.json()
            sid = data.get("id", "")
            if sid:
                user_sessions[user_id] = sid
                logger.info(f"Created session {sid} for user {user_id}")
                return sid
    except Exception as e:
        logger.error(f"Failed to create session: {e}")

    # Fallback: try to use an existing session
    try:
        resp = await client.get(
            f"{ODYSSEUS_URL}/api/sessions",
            headers=headers,
            timeout=10,
        )
        if resp.status_code == 200:
            sessions = resp.json()
            if isinstance(sessions, list):
                for s in sessions:
                    if f"Telegram-{user_id}" in s.get("name", ""):
                        sid = s["id"]
                        user_sessions[user_id] = sid
                        return sid
                if sessions:
                    sid = sessions[0]["id"]
                    user_sessions[user_id] = sid
                    return sid
    except Exception as e:
        logger.error(f"Failed to list sessions: {e}")

    return ""


async def chat_with_odysseus(message: str, session_id: str, client: httpx.AsyncClient,
                              use_web: bool = False) -> str:
    """Send a message to Odysseus and get the response."""
    api_token = get_api_token()
    headers = {"Content-Type": "application/json"}
    if api_token:
        headers["Authorization"] = f"Bearer {api_token}"

    try:
        resp = await client.post(
            f"{ODYSSEUS_URL}/api/chat",
            json={
                "message": message,
                "session": session_id,
                "use_web": use_web,
                "attachments": [],
            },
            headers=headers,
            timeout=180,
        )
        if resp.status_code == 200:
            data = resp.json()
            return data.get("response", "⚠️ No response from model")
        else:
            error = resp.text[:500]
            logger.error(f"Chat API error {resp.status_code}: {error}")
            return f"⚠️ Error ({resp.status_code}): {error}"
    except httpx.TimeoutException:
        return "⏳ Response timed out. The model is taking too long."
    except Exception as e:
        logger.error(f"Chat error: {e}")
        return f"⚠️ Error: {str(e)}"


def clean_response(text: str) -> str:
    """Clean up Odysseus response for Telegram display."""
    # Remove ```tool_name ... ``` blocks (tool execution output)
    text = re.sub(r'```(?:bash|python|web_search|web_fetch|read_file|write_file|edit_file|grep|glob|ls|create_document|edit_document|manage_memory|manage_tasks|trigger_research)\n.*?\n```', '', text, flags=re.DOTALL)
    # Remove [doc:...] references
    text = re.sub(r'\[doc:[^\]]+\]', '', text)
    # Remove <tool_result>...</tool_result> blocks
    text = re.sub(r'<tool_result>.*?</tool_result>', '', text, flags=re.DOTALL)
    # Clean up excess whitespace
    text = re.sub(r'\n{3,}', '\n\n', text).strip()
    return text


async def handle_message(update: dict, client: httpx.AsyncClient):
    """Process an incoming Telegram message."""
    msg = update.get("message", {})
    if not msg:
        return

    chat_id = msg.get("chat", {}).get("id")
    user_id = str(msg.get("from", {}).get("id", ""))
    username = msg.get("from", {}).get("username", "unknown")
    text = msg.get("text", "").strip()
    message_id = msg.get("message_id")

    if not text or not chat_id:
        return

    # Check allowed users
    if ALLOWED_USERS and user_id not in ALLOWED_USERS:
        await send_message(chat_id, "⛔ You are not authorized to use this bot.", client)
        logger.warning(f"Unauthorized user: {user_id} ({username})")
        return

    logger.info(f"Message from {username} ({user_id}): {text[:100]}")

    # Handle commands
    if text.startswith("/"):
        cmd = text.split()[0].lower().split("@")[0]

        if cmd == "/start":
            await send_message(
                chat_id,
                "🚀 *Welcome to Odysseus AI!*\n\n"
                "Send me any message and I'll respond using AI.\n\n"
                "*Commands:*\n"
                "• /new — Start a new conversation\n"
                "• /web `query` — Search the web\n"
                "• /research `topic` — Deep research\n"
                "• /help — Show this help\n\n"
                "Powered by NVIDIA Nemotron 🧠",
                client,
            )
            return

        elif cmd == "/new":
            user_sessions.pop(user_id, None)
            await send_message(chat_id, "🔄 New conversation started!", client)
            return

        elif cmd == "/help":
            await send_message(
                chat_id,
                "*Odysseus AI Bot* 🏛️\n\n"
                "Just send any message to chat with AI.\n\n"
                "*Agent Capabilities:*\n"
                "• 🧠 Persistent memory across chats\n"
                "• 🌐 Web search & fetching\n"
                "• 💻 Code execution (Python & Bash)\n"
                "• 📄 Document creation & editing\n"
                "• 🔬 Deep research mode\n"
                "• 📁 File operations\n\n"
                "*Commands:*\n"
                "• /new — New conversation\n"
                "• /web `query` — Web search\n"
                "• /research `topic` — Deep research\n"
                "• /help — This message",
                client,
            )
            return

        elif cmd == "/web":
            query = text[len("/web"):].strip()
            if not query:
                await send_message(chat_id, "Usage: /web `your query`", client)
                return
            text = f"Search the web for: {query}"

        elif cmd == "/research":
            topic = text[len("/research"):].strip()
            if not topic:
                await send_message(chat_id, "Usage: /research `your topic`", client)
                return
            text = f"Do deep research on: {topic}"

    # Send typing indicator
    await send_typing(chat_id, client)

    # Get or create session
    session_id = await get_or_create_session(user_id, client)
    if not session_id:
        await send_message(
            chat_id,
            "⚠️ Could not create a chat session. Please try again later.",
            client,
        )
        return

    # Determine if web search should be used
    use_web = any(text.lower().startswith(p) for p in [
        "search the web", "ابحث", "/web", "search for"
    ])

    # Keep sending typing while waiting
    typing_task = asyncio.create_task(_keep_typing(chat_id, client))

    try:
        # Chat with Odysseus
        response = await chat_with_odysseus(text, session_id, client, use_web=use_web)
        response = clean_response(response)

        if not response:
            response = "✅ Done! (The result was shown in a document/tool output)"

        await send_message(chat_id, response, client, reply_to=message_id)
    finally:
        typing_task.cancel()


async def _keep_typing(chat_id: int, client: httpx.AsyncClient):
    """Keep sending typing indicator every 5 seconds."""
    try:
        while True:
            await send_typing(chat_id, client)
            await asyncio.sleep(5)
    except asyncio.CancelledError:
        pass


async def poll_updates(client: httpx.AsyncClient):
    """Long-poll Telegram for updates."""
    offset = 0
    logger.info("Starting Telegram bot polling...")

    while True:
        try:
            resp = await client.post(
                f"{TG_API}/getUpdates",
                json={"offset": offset, "timeout": 30, "allowed_updates": ["message"]},
                timeout=35,
            )
            data = resp.json()

            if data.get("ok") and data.get("result"):
                for update in data["result"]:
                    offset = update["update_id"] + 1
                    try:
                        await handle_message(update, client)
                    except Exception as e:
                        logger.error(f"Error handling update: {e}", exc_info=True)

        except httpx.TimeoutException:
            continue
        except Exception as e:
            logger.error(f"Polling error: {e}")
            await asyncio.sleep(5)


async def main():
    if not TELEGRAM_TOKEN:
        logger.error("TELEGRAM_BOT_TOKEN not set — Telegram bot disabled")
        return

    logger.info(f"Telegram bot starting (allowed users: {ALLOWED_USERS or 'all'})")

    # Wait for auto_configure to finish and write the token
    for i in range(60):
        token = get_api_token()
        if token:
            logger.info(f"Got API token (prefix: {token[:8]}...)")
            break
        logger.info(f"Waiting for API token... ({i+1}/60)")
        await asyncio.sleep(2)

    if not get_api_token():
        logger.error("No API token available — Telegram bot cannot authenticate!")
        return

    async with httpx.AsyncClient() as client:
        # Wait for Odysseus to be ready
        for i in range(60):
            try:
                resp = await client.get(f"{ODYSSEUS_URL}/api/auth/status", timeout=5)
                if resp.status_code == 200:
                    logger.info("Odysseus is ready!")
                    break
            except Exception:
                pass
            if i % 10 == 0:
                logger.info(f"Waiting for Odysseus... ({i+1}/60)")
            await asyncio.sleep(2)

        # Delete any existing webhook so polling works
        await tg_request("deleteWebhook", {"drop_pending_updates": False}, client)

        # Set bot commands
        await tg_request("setMyCommands", {
            "commands": [
                {"command": "new", "description": "Start new conversation"},
                {"command": "web", "description": "Search the web"},
                {"command": "research", "description": "Deep research on a topic"},
                {"command": "help", "description": "Show help"},
            ]
        }, client)

        logger.info("Telegram bot is live! Polling for messages...")

        # Start polling
        await poll_updates(client)


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