replica-omnisciente/scripts/lifestream/lifestream_userbot.py

#!/usr/bin/env python3
"""Lifestream Userbot — monitors Telegram Saved Messages via Telethon.

Uses your Telegram USER account (not a bot) to watch Saved Messages
and ingest them into the lifestream SQLite database.

First run requires interactive authentication (phone + code).
After that, the session is persisted in data/session.session.

Requirements:
    pip install telethon aiosqlite

Environment variables (in .env):
    TELEGRAM_API_ID        — from https://my.telegram.org
    TELEGRAM_API_HASH      — from https://my.telegram.org
    LIFESTREAM_DB          — path to SQLite DB (default: ./data/lifestream.db)
    MEDIA_PATH             — path to save media (default: ./data/media)
"""
import asyncio
import json
import os
import re
import sys
from datetime import datetime, timedelta, timezone
from pathlib import Path

import aiosqlite
from telethon import TelegramClient, events
from telethon.tl.types import (
    MessageMediaPhoto,
    MessageMediaDocument,
    MessageMediaWebPage,
)

# ─── Config ──────────────────────────────────────────────────────────

BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "data"
MEDIA_DIR = Path(os.environ.get("MEDIA_PATH", str(DATA_DIR / "media")))
DB_PATH = Path(os.environ.get("LIFESTREAM_DB", str(DATA_DIR / "lifestream.db")))
SESSION_PATH = str(DATA_DIR / "session")

API_ID = int(os.environ.get("TELEGRAM_API_ID", "0"))
API_HASH = os.environ.get("TELEGRAM_API_HASH", "")

if not API_ID or not API_HASH:
    print("ERROR: TELEGRAM_API_ID and TELEGRAM_API_HASH are required.")
    print("Get them at https://my.telegram.org → API development tools")
    sys.exit(1)

# Saved Messages is "self" — the user's own chat
SAVED_MESSAGES_PEER = "self"


# ─── Database ────────────────────────────────────────────────────────

async def init_db(db_path: Path) -> aiosqlite.Connection:
    """Open the lifestream DB and ensure the entries table exists."""
    db_path.parent.mkdir(parents=True, exist_ok=True)
    db = await aiosqlite.connect(str(db_path))
    await db.execute("""
        CREATE TABLE IF NOT EXISTS entries (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            source TEXT NOT NULL DEFAULT 'telegram',
            source_id TEXT NOT NULL UNIQUE,
            timestamp TEXT NOT NULL,
            content_type TEXT NOT NULL,
            raw_text TEXT,
            urls TEXT,
            articles TEXT,
            documents TEXT,
            classification TEXT,
            chat_id TEXT,
            chat_title TEXT,
            from_name TEXT,
            from_username TEXT,
            reply_to_message_id TEXT,
            is_forwarded INTEGER DEFAULT 0,
            forward_from TEXT,
            media_group_id TEXT,
            created_at TEXT DEFAULT (datetime('now'))
        )
    """)
    await db.execute("CREATE INDEX IF NOT EXISTS idx_entries_timestamp ON entries(timestamp DESC)")
    await db.execute("CREATE INDEX IF NOT EXISTS idx_entries_source_id ON entries(source_id)")
    await db.commit()
    return db


async def entry_exists(db: aiosqlite.Connection, source_id: str) -> bool:
    """Check if an entry already exists (idempotent)."""
    async with db.execute(
        "SELECT 1 FROM entries WHERE source_id = ?", (source_id,)
    ) as cur:
        return (await cur.fetchone()) is not None


async def insert_entry(db: aiosqlite.Connection, entry: dict):
    """Insert a lifestream entry."""
    await db.execute(
        """INSERT OR IGNORE INTO entries
           (source, source_id, timestamp, content_type, raw_text, urls,
            documents, chat_id, chat_title, from_name, from_username,
            reply_to_message_id, is_forwarded, forward_from, media_group_id)
           VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
        (
            entry.get("source", "telegram"),
            entry["source_id"],
            entry["timestamp"],
            entry["content_type"],
            entry.get("raw_text"),
            json.dumps(entry.get("urls", [])),
            json.dumps(entry.get("documents", [])),
            entry.get("chat_id"),
            entry.get("chat_title"),
            entry.get("from_name"),
            entry.get("from_username"),
            entry.get("reply_to_message_id"),
            1 if entry.get("is_forwarded") else 0,
            entry.get("forward_from"),
            entry.get("media_group_id"),
        ),
    )
    await db.commit()


# ─── Media download ──────────────────────────────────────────────────

async def download_media(client: TelegramClient, msg, media_type: str) -> str | None:
    """Download media from a message. Returns the local file path."""
    MEDIA_DIR.mkdir(parents=True, exist_ok=True)
    ext_map = {"photo": ".jpg", "video": ".mp4", "voice": ".ogg", "audio": ".mp3", "document": ""}
    ext = ext_map.get(media_type, "")
    filename = f"{msg.id}{ext}"
    out_path = MEDIA_DIR / filename
    if out_path.exists():
        return str(out_path)
    try:
        await client.download_media(msg, file=str(out_path))
        return str(out_path) if out_path.exists() else None
    except Exception as e:
        print(f"[Media] Download failed for {msg.id}: {e}")
        return None


# ─── Message classification ──────────────────────────────────────────

def classify_message(msg) -> dict:
    """Extract content type and metadata from a Telegram message."""
    content_type = "text"
    media_type = None

    if msg.photo:
        content_type, media_type = "image", "photo"
    elif msg.video:
        content_type, media_type = "video", "video"
    elif msg.voice:
        content_type, media_type = "voice", "voice"
    elif msg.audio:
        content_type, media_type = "audio", "audio"
    elif msg.document:
        content_type, media_type = "document", "document"
    elif msg.sticker:
        content_type, media_type = "sticker", "sticker"
    elif msg.geo:
        content_type, media_type = "location", "location"
    elif msg.contact:
        content_type, media_type = "contact", "contact"
    elif msg.gif:
        content_type, media_type = "animation", "animation"
    elif msg.video_note:
        content_type, media_type = "video_note", "video_note"
    elif msg.poll:
        content_type, media_type = "poll", "poll"

    return {"content_type": content_type, "media_type": media_type}


def extract_urls(text: str | None) -> list[str]:
    """Extract URLs from message text."""
    if not text:
        return []
    return re.findall(r"https?://\S+", text)


# ─── Message handler ─────────────────────────────────────────────────

async def handle_message(client: TelegramClient, db: aiosqlite.Connection, msg):
    """Process a single message from Saved Messages."""
    source_id = str(msg.id)

    if await entry_exists(db, source_id):
        return  # already ingested

    info = classify_message(msg)
    text = msg.text or msg.message or ""
    media_file = None

    if info["media_type"]:
        media_file = await download_media(client, msg, info["media_type"])

    # Build forward info
    forward_from = None
    is_forwarded = False
    if msg.forward:
        is_forwarded = True
        if msg.forward.from_id:
            forward_from = str(msg.forward.from_id)

    entry = {
        "source": "telegram",
        "source_id": source_id,
        "timestamp": datetime.fromtimestamp(msg.date.timestamp(), tz=timezone.utc).isoformat(),
        "content_type": info["content_type"],
        "raw_text": text if text else None,
        "urls": extract_urls(text),
        "documents": [media_file] if media_file else [],
        "chat_id": "saved_messages",
        "chat_title": "Saved Messages",
        "from_name": "self",
        "from_username": None,
        "reply_to_message_id": str(msg.reply_to.reply_to_msg_id) if msg.reply_to else None,
        "is_forwarded": is_forwarded,
        "forward_from": forward_from,
        "media_group_id": str(msg.grouped_id) if msg.grouped_id else None,
    }

    await insert_entry(db, entry)
    media = text[:80] if text else "(media)"
    print(f"[Ingest] {info['content_type']}: {media}")


# ─── Historical sync ─────────────────────────────────────────────────

async def sync_historical(client: TelegramClient, db: aiosqlite.Connection, limit: int = 500):
    """Fetch recent Saved Messages and ingest any missing ones."""
    print(f"[Sync] Fetching last {limit} Saved Messages...")
    count = 0
    async for msg in client.iter_messages(SAVED_MESSAGES_PEER, limit=limit):
        source_id = str(msg.id)
        if not await entry_exists(db, source_id):
            await handle_message(client, db, msg)
            count += 1
    print(f"[Sync] Ingested {count} new messages")


# ─── Polling fallback ────────────────────────────────────────────────

POLL_INTERVAL = 60  # seconds

async def poll_self_messages(client: TelegramClient, db: aiosqlite.Connection):
    """Poll Saved Messages because Telethon events.NewMessage(chats='self') does not
    fire for messages sent by the user itself."""
    print("[Poll] Starting self-message polling loop")
    last_check = datetime.now(timezone.utc) - timedelta(seconds=POLL_INTERVAL)
    while True:
        try:
            since = last_check - timedelta(seconds=5)
            async for msg in client.iter_messages(SAVED_MESSAGES_PEER, limit=200):
                msg_time = msg.date.replace(tzinfo=timezone.utc) if msg.date.tzinfo is None else msg.date
                if msg_time > since:
                    await handle_message(client, db, msg)
            last_check = datetime.now(timezone.utc)
        except Exception as e:
            print(f"[Poll] error: {e}")
        await asyncio.sleep(POLL_INTERVAL)


# ─── Main ────────────────────────────────────────────────────────────

async def main():
    # Init DB
    db = await init_db(DB_PATH)

    # Count existing entries
    async with db.execute("SELECT COUNT(*) FROM entries") as cur:
        row = await cur.fetchone()
        total = row[0] if row else 0
    print(f"[Lifestream] DB has {total} entries")

    # Init Telethon client
    client = TelegramClient(SESSION_PATH, API_ID, API_HASH)
    await client.start()
    me = await client.get_me()
    print(f"[Lifestream] Logged in as {me.first_name} (ID: {me.id})")

    # Historical sync on startup
    await sync_historical(client, db)

    # Live listener for new Saved Messages (works for some forwards / external events)
    @client.on(events.NewMessage(chats=SAVED_MESSAGES_PEER))
    async def on_new_message(event):
        await handle_message(client, db, event.message)

    # Start polling task as a background coroutine
    poll_task = asyncio.create_task(poll_self_messages(client, db))

    print("[Lifestream] Listening for new Saved Messages...")
    try:
        await client.run_until_disconnected()
    finally:
        poll_task.cancel()
        try:
            await poll_task
        except asyncio.CancelledError:
            pass


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