core/telegram_bot.py

"""
QuantForge Telegram Notifier
────────────────────────────
Sends auth-expiry alerts, listens for cookie replies, and fires
submission notifications. Pure httpx — no extra dependencies.
"""
import asyncio
import httpx
import logging
import time
from typing import Optional

logger = logging.getLogger("TelegramBot")

_COOKIE_HINTS = ("t=", "session=", "csrftoken=", "sessionid=")

def _looks_like_cookie(text: str) -> bool:
    """Heuristic: does this text look like a WQ session cookie string?"""
    if not text or len(text) < 15:
        return False
    # Named cookie (t=xxx; session=xxx …)
    if any(hint in text for hint in _COOKIE_HINTS):
        return True
    # Raw long token — no spaces, mostly alnum + - _ . =
    stripped = text.replace("-", "").replace("_", "").replace(".", "").replace("=", "").replace(";", "")
    if len(text) > 30 and " " not in text and stripped.isalnum():
        return True
    return False

class TelegramNotifier:
    """
    Async Telegram bot client using HTTPX.
    """
    def __init__(self, token: str, chat_id: str, proxy_url: Optional[str] = None) -> None:
        self.token = token.strip()
        self.chat_id = str(chat_id).strip()
        self.proxy_url = proxy_url.strip() if (proxy_url and proxy_url.strip()) else None
        
        base_domain = self.proxy_url if self.proxy_url else "https://api.telegram.org"
        if base_domain.endswith("/"):
            base_domain = base_domain[:-1]
            
        self._base = f"{base_domain}/bot{self.token}"
        self._client: Optional[httpx.AsyncClient] = None
        self._last_update_id: int = 0

        # Biometric state variables
        self._auth_active = False
        self._biometric_url = None
        self._biometric_client = None
        self._biometric_cancelled = False
        self._successful_cookie = None
        self._polling_failed = False

    async def _get_client(self) -> httpx.AsyncClient:
        if self._client is None or self._client.is_closed:
            # Bypass SSL certificate checks for api.telegram.org in container sandbox
            self._client = httpx.AsyncClient(
                verify=False,
                timeout=httpx.Timeout(35.0)
            )
        return self._client

    async def close(self) -> None:
        if self._client and not self._client.is_closed:
            await self._client.aclose()
        self._cleanup_biometric()

    async def send(self, text: str, parse_mode: str = "HTML", reply_markup: Optional[dict] = None) -> bool:
        """Send a message to the configured chat. Returns True on success."""
        try:
            client = await self._get_client()
            payload = {"chat_id": self.chat_id, "text": text, "parse_mode": parse_mode}
            if reply_markup is not None:
                payload["reply_markup"] = reply_markup
            resp = await client.post(f"{self._base}/sendMessage", json=payload)
            if resp.status_code == 200:
                return True
            body = resp.text
            logger.error(f"Telegram sendMessage failed {resp.status_code}: {body[:200]}")
            return False
        except Exception as exc:
            import traceback
            logger.error(f"Telegram send error: {exc}\n{traceback.format_exc()}")
            return False

    async def _answer_callback_query(self, callback_query_id: str) -> None:
        try:
            client = await self._get_client()
            payload = {"callback_query_id": callback_query_id}
            await client.post(f"{self._base}/answerCallbackQuery", json=payload)
        except Exception as e:
            logger.warning(f"Failed to answer callback query: {e}")

    async def send_auth_expired_alert(self) -> None:
        msg = (
            "<b>🔴 QuantForge — Session Expired</b>\n\n"
            "The WorldQuant BRAIN session is paused. Choose an option to authenticate:\n\n"
            "1️⃣ <b>Biometric Flow:</b> Click <b>Authenticate Now</b> below (or reply <b>OK</b>) to get a face scan link.\n"
            "2️⃣ <b>One-Tap Sync:</b> Use your configured bookmarklet on mobile or Chrome extension on desktop.\n"
            "3️⃣ <b>Manual Pasting:</b> Reply directly with your fresh cookie string (starting with <code>t=</code>)."
        )
        reply_markup = {
            "inline_keyboard": [
                [
                    {"text": "🔐 Authenticate Now", "callback_data": "auth_now"}
                ]
            ]
        }
        await self.send(msg, reply_markup=reply_markup)
        logger.info("Telegram: Auth-expired alert sent.")

    async def send_session_restored(self) -> None:
        await self.send(
            "<b>✅ QuantForge — Session Restored</b>\n\n"
            "New cookie accepted. Miner is resuming operations. 🚀"
        )
        logger.info("Telegram: Session-restored confirmation sent.")

    async def send_auth_failed_permanently(self) -> None:
        await self.send(
            "<b>⚠️ QuantForge — Auth Recovery Timed Out</b>\n\n"
            "No cookie was received within 30 minutes.\n"
            "The miner will keep retrying every 10 minutes.\n"
            "Please send your new cookie whenever you're ready."
        )

    async def send_submission_alert(
        self, alpha_id: str, sharpe: float, fitness: float, region: str = "USA"
    ) -> None:
        msg = (
            f"<b>🎯 Alpha Submitted to BRAIN!</b>\n\n"
            f"• Region: <code>{region}</code>\n"
            f"• Sharpe: <code>{sharpe:.4f}</code>\n"
            f"• Fitness: <code>{fitness:.4f}</code>\n\n"
            "Check your WorldQuant BRAIN dashboard 💰"
        )
        await self.send(msg)

    async def send_startup(self) -> None:
        await self.send(
            "<b>🟢 QuantForge Miner — Started</b>\n\n"
            "System is live and authenticated. Mining alphas... 🔬"
        )

    async def _drain_pending_updates(self) -> None:
        """Advance offset past all currently queued messages so we only
        listen for messages sent *after* this call."""
        try:
            client = await self._get_client()
            resp = await client.get(
                f"{self._base}/getUpdates",
                params={"offset": -1, "limit": 1, "timeout": 1},
                timeout=5.0
            )
            if resp.status_code == 200:
                data = resp.json()
                updates = data.get("result", [])
                if updates:
                    self._last_update_id = updates[-1]["update_id"]
        except Exception:
            pass

    def _cleanup_biometric(self) -> None:
        self._auth_active = False
        self._biometric_url = None
        self._biometric_cancelled = False
        self._successful_cookie = None
        self._polling_failed = False
        if self._biometric_client:
            client = self._biometric_client
            self._biometric_client = None
            try:
                asyncio.create_task(client.aclose())
            except RuntimeError:
                pass

    def _cancel_biometric(self) -> None:
        self._biometric_cancelled = True
        if hasattr(self, "_polling_task") and self._polling_task and not self._polling_task.done():
            self._polling_task.cancel()
        self._cleanup_biometric()

    async def _run_biometric_polling(self) -> None:
        import os
        import base64
        from config import WQ_MFA_POLL_INTERVAL, WQ_MFA_TIMEOUT
        
        poll_interval = WQ_MFA_POLL_INTERVAL
        timeout = WQ_MFA_TIMEOUT
        
        username = os.environ.get("WQ_USERNAME", "").strip()
        password = os.environ.get("WQ_PASSWORD", "").strip()
        credentials = f"{username}:{password}"
        encoded_credentials = base64.b64encode(credentials.encode()).decode()
        auth_headers = {
            "Authorization": f"Basic {encoded_credentials}",
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
        }
        
        # The inquiry URL is the ORIGINAL /authentication/persona?inquiry=inq_xxx
        # returned in the Location header of the initial 401. Polling this specific
        # URL (rather than /authentication) checks the status of the original scan.
        poll_url = getattr(self, "_biometric_url", None)
        
        logger.info(
            f"Telegram Bot: Starting biometric polling loop "
            f"(interval={poll_interval}s, timeout={timeout}s, "
            f"username={'SET' if username else 'MISSING'}, "
            f"poll_url={poll_url})"
        )
        
        if not poll_url:
            logger.error("Telegram Bot: No biometric URL stored — cannot poll. Aborting.")
            self._polling_failed = True
            await self.send("⚠️ <b>Internal error: biometric URL not set.</b> Please try again.")
            return
        
        # Log initial cookie jar state
        initial_cookies = dict(self._biometric_client.cookies)
        logger.info(f"Telegram Bot: Initial cookie jar keys: {list(initial_cookies.keys())}")
        
        deadline = time.monotonic() + timeout
        poll_count = 0
        
        try:
            while time.monotonic() < deadline:
                if self._biometric_cancelled:
                    logger.info("Telegram Bot: Biometric polling cancelled by user request.")
                    return
                
                await asyncio.sleep(poll_interval)
                
                if self._biometric_cancelled:
                    return
                
                poll_count += 1
                try:
                    # Log exactly what we're sending
                    current_cookies = dict(self._biometric_client.cookies)
                    logger.info(
                        f"Telegram Bot: Poll #{poll_count} — POST {poll_url} "
                        f"| Cookie jar keys: {list(current_cookies.keys())} "
                        f"| Auth header present: {'Authorization' in auth_headers}"
                    )
                    
                    # POST to the original inquiry URL — WQ returns 201 once the Persona scan is complete
                    resp = await self._biometric_client.post(
                        poll_url,
                        headers=auth_headers
                    )
                    
                    # Log the complete response for diagnosis
                    resp_body = ""
                    try:
                        resp_body = resp.text[:500]
                    except Exception:
                        resp_body = "<unreadable>"
                    
                    resp_headers_log = dict(resp.headers) if hasattr(resp, "headers") else {}
                    logger.info(
                        f"Telegram Bot: Poll #{poll_count} response — "
                        f"Status: {resp.status_code} | "
                        f"Body: {resp_body!r} | "
                        f"Location: {resp_headers_log.get('location', resp_headers_log.get('Location', 'none'))} | "
                        f"WWW-Auth: {resp_headers_log.get('www-authenticate', resp_headers_log.get('WWW-Authenticate', 'none'))}"
                    )
                    
                    # Log updated cookie jar after each poll
                    after_cookies = dict(self._biometric_client.cookies)
                    logger.info(f"Telegram Bot: Poll #{poll_count} cookie jar after request: {list(after_cookies.keys())}")
                    
                    if resp.status_code == 201:
                        logger.info("Telegram Bot: Biometric authentication successful! Got 201.")
                        
                        cookies_dict = {}
                        for k, v in self._biometric_client.cookies.items():
                            cookies_dict[k] = v
                        
                        logger.info(f"Telegram Bot: Captured {len(cookies_dict)} cookies on success: {list(cookies_dict.keys())}")
                        
                        cookie_parts = [f"{k}={v}" for k, v in cookies_dict.items()]
                        cookie_string = "; ".join(cookie_parts)
                        
                        if cookie_string:
                            self._successful_cookie = cookie_string
                            await self.send("✅ <b>Biometric verification successful!</b> Session restored. Resuming miner loop.")
                            return
                        else:
                            logger.error("Telegram Bot: Got 201 but no session cookies found in jar.")
                            self._polling_failed = True
                            await self.send("⚠️ <b>Biometric succeeded but no session cookies were returned.</b>")
                            return
                    
                    # Check if the scan is pending:
                    # 202 = accepted/pending, 401 = still awaiting scan,
                    # 403 with INQUIRY_INCOMPLETE body = user has not finished scan yet.
                    is_pending = False
                    if resp.status_code in (202, 401):
                        is_pending = True
                    elif resp.status_code == 403 and "INQUIRY_INCOMPLETE" in resp_body:
                        is_pending = True

                    if is_pending:
                        pending_location = resp_headers_log.get("location", resp_headers_log.get("Location", "none"))
                        logger.info(
                            f"Telegram Bot: Poll #{poll_count} — Scan still pending ({resp.status_code}). "
                            f"Location: {pending_location}. Waiting..."
                        )
                        continue
                    
                    elif resp.status_code == 429:
                        # Throttled mid-polling — wait for Retry-After, then continue
                        retry_after = int(resp_headers_log.get("retry-after", resp_headers_log.get("Retry-After", 30)))
                        logger.warning(
                            f"Telegram Bot: Poll #{poll_count} — Throttled (429). "
                            f"Waiting {retry_after}s before continuing..."
                        )
                        await asyncio.sleep(retry_after)
                        continue
                    
                    else:
                        logger.error(
                            f"Telegram Bot: Poll #{poll_count} — Unexpected status {resp.status_code}. "
                            f"Full body: {resp_body!r}"
                        )
                        self._polling_failed = True
                        await self.send(f"❌ <b>Biometric verification failed.</b> (Status {resp.status_code})\n<pre>{resp_body[:200]}</pre>")
                        return
                
                except Exception as poll_err:
                    logger.warning(f"Telegram Bot: Poll #{poll_count} exception: {type(poll_err).__name__}: {poll_err}", exc_info=True)
            
            logger.warning(f"Telegram Bot: Biometric polling timed out after {poll_count} polls.")
            self._polling_failed = True
            await self.send("⏰ <b>Biometric authentication timed out.</b> Please try again.")
        except Exception as e:
            logger.error(f"Telegram Bot: Exception in polling task: {type(e).__name__}: {e}", exc_info=True)
            self._polling_failed = True


    async def _initiate_biometric_auth(self) -> None:
        import os
        import base64
        import urllib.parse
        
        username = os.environ.get("WQ_USERNAME", "").strip()
        password = os.environ.get("WQ_PASSWORD", "").strip()
        if not username or not password:
            await self.send("⚠️ <b>Credentials Missing:</b> WQ_USERNAME and WQ_PASSWORD must be configured in your environment to authenticate.")
            return

        await self.send("⏳ <b>Initiating login to generate verification link...</b>")
        
        self._cleanup_biometric()
        
        credentials = f"{username}:{password}"
        encoded_credentials = base64.b64encode(credentials.encode()).decode()
        
        default_headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
        }
        
        self._biometric_client = httpx.AsyncClient(
            verify=False,
            timeout=httpx.Timeout(10.0),
            headers=default_headers
        )
        
        auth_headers = {
            "Authorization": f"Basic {encoded_credentials}"
        }
        
        try:
            resp = await self._biometric_client.post(
                "https://api.worldquantbrain.com/authentication",
                headers=auth_headers
            )
            
            # Handle rate limiting — WQ throttles biometric attempts across sessions
            if resp.status_code == 429:
                retry_after = int(resp.headers.get("Retry-After", 60))
                body_text = ""
                try:
                    body_text = resp.text[:200]
                except Exception:
                    pass
                logger.warning(f"Telegram Bot: WQ auth throttled (429). Retry-After: {retry_after}s. Body: {body_text!r}")
                await self.send(
                    f"⏱ <b>WorldQuant is throttling auth requests.</b>\n"
                    f"Too many attempts were made recently. Waiting <b>{retry_after}s</b> before retrying automatically…"
                )
                await asyncio.sleep(retry_after)
                logger.info("Telegram Bot: Retrying authentication after throttle backoff...")
                resp = await self._biometric_client.post(
                    "https://api.worldquantbrain.com/authentication",
                    headers=auth_headers
                )
                logger.info(f"Telegram Bot: Retry response status: {resp.status_code}")
                if resp.status_code == 429:
                    logger.error("Telegram Bot: Still throttled after backoff. Aborting.")
                    await self.send("❌ <b>Still throttled by WorldQuant.</b> Please wait a few minutes and try again.")
                    return
            
            if resp.status_code == 201:
                logger.info("Telegram Bot: Basic Auth succeeded directly without biometrics.")
                cookie_parts = []
                for k, v in self._biometric_client.cookies.items():
                    cookie_parts.append(f"{k}={v}")
                cookie_string = "; ".join(cookie_parts)
                if cookie_string:
                    self._successful_cookie = cookie_string
                    await self.send("✅ <b>Login succeeded immediately!</b> No biometric scan was required.")
                else:
                    await self.send("⚠️ <b>Login succeeded but no cookies were returned.</b>")
                return
                
            elif resp.status_code == 401:
                www_auth = resp.headers.get("WWW-Authenticate")
                location = resp.headers.get("Location")
                
                if www_auth == "persona" and location:
                    logger.info("Telegram Bot: Biometric challenge detected. Location: %s", location)
                    biometric_url = urllib.parse.urljoin("https://api.worldquantbrain.com/authentication", location)
                    self._biometric_url = biometric_url
                    
                    # Capture and parse challenge cookies
                    self._biometric_cookies = ""
                    set_cookies = []
                    if hasattr(resp.headers, "getlist"):
                        set_cookies = resp.headers.getlist("set-cookie")
                    elif hasattr(resp.headers, "get"):
                        cookie_val = resp.headers.get("set-cookie")
                        if cookie_val:
                            set_cookies = [cookie_val] if isinstance(cookie_val, str) else cookie_val
                    cookies_dict = {}
                    for cookie_str in set_cookies:
                        parts = cookie_str.split(";")[0].split("=")
                        if len(parts) == 2:
                            cookies_dict[parts[0].strip()] = parts[1].strip()
                    # Also include client jar cookies
                    for k, v in self._biometric_client.cookies.items():
                        cookies_dict[k] = v
                    if cookies_dict:
                        self._biometric_cookies = "; ".join([f"{k}={v}" for k, v in cookies_dict.items()])
                        logger.info(f"Telegram Bot: Captured challenge cookies: {self._biometric_cookies}")
                    
                    self._auth_active = True
                    self._biometric_cancelled = False
                    self._polling_failed = False
                    self._successful_cookie = None
                    self._polling_task = asyncio.create_task(self._run_biometric_polling())
                    
                    reply_markup = {
                        "inline_keyboard": [
                            [
                                {"text": "❌ Cancel", "callback_data": "cancel_auth"}
                            ]
                        ]
                    }
                    await self.send(
                        f"🔗 <b>Biometric Verification Required:</b>\n\n"
                        f"<a href=\"{biometric_url}\">Click here to start your Face Scan</a>\n\n"
                        f"⏳ Once completed, the session will restore automatically.\n"
                        f"<i>(You can click Cancel to abort)</i>",
                        reply_markup=reply_markup
                    )
                else:
                    logger.warning("Telegram Bot: 401 response without persona challenge: %s", resp.text)
                    await self.send("❌ <b>Authentication failed:</b> Incorrect WorldQuant username or password.")
            else:
                resp_body = ""
                try:
                    resp_body = resp.text[:200]
                except Exception:
                    pass
                logger.error(f"Telegram Bot: WQ auth returned status {resp.status_code}. Body: {resp_body!r}")
                await self.send(f"❌ <b>WorldQuant Auth Error:</b> Received status code {resp.status_code}\n<pre>{resp_body}</pre>")
        except Exception as e:
            logger.exception(f"Telegram Bot: Exception during authentication initiation: {e}")
            await self.send(f"❌ <b>Connection Error:</b> Failed to contact WorldQuant Brain: {e}")

    async def poll_for_cookie(self, timeout_seconds: int = 1800, wake_event: asyncio.Event = None) -> Optional[str]:
        """
        Long-polls Telegram for a message that looks like a WQ session cookie.
        Also wakes up immediately if wake_event is set (by /login-cookie HTTP endpoint).
        Returns the cookie string on success, None on timeout.
        Default timeout: 30 minutes.
        """
        import os
        logger.info("Telegram: Listening for cookie reply (30-minute window)…")
        await self._drain_pending_updates()

        client = await self._get_client()
        deadline = time.monotonic() + timeout_seconds
        warned = False

        while time.monotonic() < deadline:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                break

            # --- Check if /login-cookie HTTP endpoint already delivered a cookie ---
            if wake_event and wake_event.is_set():
                wake_event.clear()
                fresh_cookie = os.environ.get("WQ_SESSION_COOKIE", "").strip()
                if fresh_cookie:
                    logger.info("Telegram: Wake event fired by /login-cookie endpoint. Using HTTP-delivered cookie.")
                    self._cancel_biometric()
                    return fresh_cookie

            # --- Check if biometric polling task completed successfully ---
            if self._successful_cookie:
                cookie = self._successful_cookie
                self._cleanup_biometric()
                return cookie

            # Warn at 10-minute mark
            if remaining < 600 and not warned:
                warned = True
                await self.send(
                    "<b>⏰ Reminder:</b> 10 minutes remaining to send your new cookie."
                )

            try:
                poll_timeout_param = min(20, int(remaining))
                if self._auth_active:
                    poll_timeout_param = min(3, poll_timeout_param)

                import json
                params = {
                    "offset": self._last_update_id + 1,
                    "timeout": poll_timeout_param,
                    "allowed_updates": json.dumps(["message", "callback_query"]),
                }
                poll_timeout = float(poll_timeout_param + 5)
                resp = await client.get(
                    f"{self._base}/getUpdates",
                    params=params,
                    timeout=poll_timeout
                )
                if resp.status_code != 200:
                    await asyncio.sleep(5)
                    continue

                data = resp.json()
                for update in data.get("result", []):
                    self._last_update_id = max(
                        self._last_update_id, update["update_id"]
                    )
                    
                    # 1. Handle callback query (button clicks)
                    callback_query = update.get("callback_query", {})
                    if callback_query:
                        chat_id = str(callback_query.get("message", {}).get("chat", {}).get("id", ""))
                        if chat_id != self.chat_id:
                            continue
                            
                        query_id = callback_query.get("id", "")
                        await self._answer_callback_query(query_id)
                        
                        data_payload = callback_query.get("data", "")
                        if data_payload == "auth_now":
                            if self._auth_active:
                                await self.send("⚠️ <b>Auth in progress:</b> An active biometric login session is already running.")
                            else:
                                asyncio.create_task(self._initiate_biometric_auth())
                        elif data_payload == "cancel_auth":
                            if self._auth_active:
                                self._cancel_biometric()
                                await self.send("❌ <b>Authentication cancelled.</b>")
                            else:
                                await self.send("⚠️ No active authentication session to cancel.")
                        continue

                    # 2. Handle message
                    msg = update.get("message", {})
                    if str(msg.get("chat", {}).get("id", "")) != self.chat_id:
                        continue

                    text = (msg.get("text") or "").strip()
                    if not text:
                        continue
                        
                    # Check if user typed "/cancel"
                    if text == "/cancel":
                        if self._auth_active:
                            self._cancel_biometric()
                            await self.send("❌ <b>Authentication cancelled.</b>")
                        else:
                            await self.send("⚠️ No active authentication session to cancel.")
                        continue
                        
                    # Check if user typed "/force-cookie"
                    if text == "/force-cookie":
                        self._cancel_biometric()
                        await self.send(
                            "❌ <b>Biometric auth aborted.</b>\n\n"
                            "Please paste your fresh cookie starting with <code>t=</code> directly here."
                        )
                        continue
                        
                    # Check if user replied with "OK" or "/auth"
                    if text.upper() in ("OK", "/AUTH"):
                        if self._auth_active:
                            await self.send("⚠️ <b>Auth in progress:</b> An active biometric login session is already running.")
                        else:
                            asyncio.create_task(self._initiate_biometric_auth())
                        continue

                    # Check if text is a pasted session cookie
                    if _looks_like_cookie(text):
                        logger.info("Telegram: Valid cookie reply received. Cancelling any active biometric poll.")
                        self._cancel_biometric()
                        return text

                    # User sent something else — tell them what we need
                    await self.send(
                        "🤔 That doesn't look like a session cookie.\n"
                        "Please paste the <code>t=</code> value, click <b>Authenticate Now</b>, or reply <code>OK</code>."
                    )

            except httpx.TimeoutException:
                continue
            except Exception as exc:
                logger.warning(f"Telegram poll error: {exc}")
                await asyncio.sleep(10)

        # Timed out
        self._cancel_biometric()
        await self.send_auth_failed_permanently()
        logger.warning("Telegram: Cookie poll timed out.")
        return None