tg-channel-analysis/bot/handlers/analyze.py

import logging
import re

from aiogram import F, Router
from aiogram.filters import Command
from aiogram.types import (
    CallbackQuery,
    FSInputFile,
    InlineKeyboardButton,
    InlineKeyboardMarkup,
    LabeledPrice,
    Message,
)
from telethon import TelegramClient
from telethon.errors import (
    ChannelInvalidError,
    ChannelPrivateError,
    FloodWaitError,
    UsernameInvalidError,
    UsernameNotOccupiedError,
)

from bot.config import settings
from bot.db import usage_repo, user_repo
from bot.i18n import Lang, t
from bot.models import AnalysisSession, Depth, FocusArea
from bot.services.analyzer import analyze_channel
from bot.services.chunker import chunk_messages
from bot.services.fetcher import fetch_channel_messages
from bot.services.formatter import split_report
from bot.services.report_saver import save_report

log = logging.getLogger(__name__)
router = Router()

_sessions: dict[int, AnalysisSession] = {}


def _extract_channel(text: str) -> str | None:
    text = text.strip()
    m = re.match(r"@(\w+)", text)
    if m:
        return m.group(1)
    m = re.match(r"https?://t\.me/(\w+)", text)
    if m:
        return m.group(1)
    if re.match(r"^\w+$", text):
        return text
    return None


def _depth_keyboard(lang: Lang) -> InlineKeyboardMarkup:
    return InlineKeyboardMarkup(inline_keyboard=[
        [InlineKeyboardButton(text=t("depth_basic", lang), callback_data="depth:basic")],
        [InlineKeyboardButton(text=t("depth_standard", lang), callback_data="depth:standard")],
        [InlineKeyboardButton(text=t("depth_full", lang), callback_data="depth:full")],
    ])


def _focus_keyboard(lang: Lang, selected: set[str]) -> InlineKeyboardMarkup:
    rows = []
    for area in FocusArea:
        check = "✅" if area.value in selected else "☐"
        rows.append([InlineKeyboardButton(
            text=f"{check} {t(f'focus_{area.value}', lang)}",
            callback_data=f"focus:{area.value}",
        )])
    rows.append([InlineKeyboardButton(text=f"✅ {t('done', lang)}", callback_data="focus:done")])
    return InlineKeyboardMarkup(inline_keyboard=rows)


def _model_keyboard() -> InlineKeyboardMarkup:
    rows = []
    for model_id in settings.models_list:
        label = model_id.split("/")[-1] if "/" in model_id else model_id
        rows.append([InlineKeyboardButton(text=label, callback_data=f"model:{model_id}")])
    return InlineKeyboardMarkup(inline_keyboard=rows)


# Step 0: /analyze @channel
@router.message(Command("analyze"))
async def cmd_analyze(message: Message, lang: Lang = Lang.EN, **_: object) -> None:
    args = (message.text or "").split(maxsplit=1)
    if len(args) < 2:
        await message.answer(t("provide_channel", lang), parse_mode="HTML")
        return

    channel = _extract_channel(args[1])
    if not channel:
        await message.answer(t("bad_channel", lang), parse_mode="HTML")
        return

    _sessions[message.from_user.id] = AnalysisSession(channel=channel)
    await message.answer(
        t("choose_depth", lang, channel=channel),
        parse_mode="HTML",
        reply_markup=_depth_keyboard(lang),
    )


# Step 1: Depth selected
@router.callback_query(F.data.startswith("depth:"))
async def on_depth(callback: CallbackQuery, lang: Lang = Lang.EN, **_: object) -> None:
    await callback.answer()
    user_id = callback.from_user.id
    session = _sessions.get(user_id)
    if not session:
        await callback.message.answer(t("session_expired", lang))
        return

    depth_val = callback.data.split(":", 1)[1]
    session.depth = Depth(depth_val)

    await callback.message.edit_text(
        t("choose_focus", lang),
        parse_mode="HTML",
        reply_markup=_focus_keyboard(lang, set()),
    )


# Step 2: Focus area toggle
@router.callback_query(F.data.startswith("focus:"))
async def on_focus(callback: CallbackQuery, lang: Lang = Lang.EN, **_: object) -> None:
    user_id = callback.from_user.id
    session = _sessions.get(user_id)
    if not session:
        await callback.answer()
        await callback.message.answer(t("session_expired", lang))
        return

    value = callback.data.split(":", 1)[1]

    if value == "done":
        if not session.focus_areas:
            await callback.answer(t("no_focus_selected", lang), show_alert=True)
            return
        await callback.answer()
        # Show model selection
        if len(settings.models_list) == 1:
            # Skip model selection if only one available
            session.model_id = settings.models_list[0]
            await _check_payment_and_run(callback.message, user_id, lang)
        else:
            await callback.message.edit_text(
                t("choose_model", lang),
                parse_mode="HTML",
                reply_markup=_model_keyboard(),
            )
        return

    await callback.answer()
    area = FocusArea(value)
    if area in session.focus_areas:
        session.focus_areas.remove(area)
    else:
        session.focus_areas.append(area)

    selected = {a.value for a in session.focus_areas}
    await callback.message.edit_reply_markup(
        reply_markup=_focus_keyboard(lang, selected),
    )


# Step 3: Model selected
@router.callback_query(F.data.startswith("model:"))
async def on_model(callback: CallbackQuery, lang: Lang = Lang.EN, **_: object) -> None:
    await callback.answer()
    user_id = callback.from_user.id
    session = _sessions.get(user_id)
    if not session:
        await callback.message.answer(t("session_expired", lang))
        return

    session.model_id = callback.data.split(":", 1)[1]
    await _check_payment_and_run(callback.message, user_id, lang)


# Step 4: Payment check + run
async def _check_payment_and_run(message: Message, user_id: int, lang: Lang) -> None:
    session = _sessions.get(user_id)
    if not session:
        return

    free_used = await user_repo.get_free_used(user_id)
    price_map = {
        Depth.BASIC: settings.price_basic,
        Depth.STANDARD: settings.price_standard,
        Depth.FULL: settings.price_full,
    }
    price = price_map[session.depth]

    if free_used < settings.free_analyses:
        await user_repo.increment_free(user_id)
        await message.edit_text(
            t("free_analysis", lang, used=free_used + 1, max=settings.free_analyses),
            parse_mode="HTML",
        )
        await _run_analysis(message, user_id, lang, stars_paid=0)
    else:
        # Record analysis first to get ID for payload
        analysis_id = await usage_repo.record_analysis(
            telegram_id=user_id,
            channel=session.channel,
            depth=session.depth.value,
            focus_areas=[a.value for a in session.focus_areas],
            model_id=session.model_id,
            stars_paid=price,
        )
        await message.answer_invoice(
            title=t("invoice_title", lang, depth=session.depth.label),
            description=t("invoice_description", lang, channel=session.channel, depth=session.depth.label),
            payload=str(analysis_id),
            currency="XTR",
            prices=[LabeledPrice(label="Analysis", amount=price)],
        )


# Payment callback triggers analysis
@router.message(lambda m: m.successful_payment is not None)
async def on_payment_run(message: Message, lang: Lang = Lang.EN, **_: object) -> None:
    user_id = message.from_user.id
    session = _sessions.get(user_id)
    if not session:
        return
    payment = message.successful_payment
    await usage_repo.record_payment(
        telegram_id=user_id,
        telegram_payment_id=payment.telegram_payment_charge_id,
        stars_amount=payment.total_amount,
        analysis_id=int(payment.invoice_payload) if payment.invoice_payload.isdigit() else None,
    )
    await _run_analysis(message, user_id, lang, stars_paid=payment.total_amount)


# Step 5: Run pipeline
async def _run_analysis(message: Message, user_id: int, lang: Lang, stars_paid: int) -> None:
    session = _sessions.pop(user_id, None)
    if not session:
        return

    telethon_client: TelegramClient | None = message.bot.__dict__.get("_telethon_client")
    if not telethon_client:
        await message.answer(t("internal_error", lang))
        return

    status_msg = await message.answer(
        t("analysis_starting", lang, channel=session.channel, status=t("fetching", lang)),
        parse_mode="HTML",
    )

    async def update_status(text: str) -> None:
        try:
            await status_msg.edit_text(
                t("analysis_starting", lang, channel=session.channel, status=text),
                parse_mode="HTML",
            )
        except Exception:
            pass

    try:
        messages, stats = await fetch_channel_messages(telethon_client, session.channel)
    except (ChannelPrivateError, ChannelInvalidError):
        await status_msg.edit_text(t("channel_private", lang))
        return
    except (UsernameInvalidError, UsernameNotOccupiedError):
        await status_msg.edit_text(t("channel_not_found", lang))
        return
    except FloodWaitError as e:
        await status_msg.edit_text(t("flood_wait", lang, s=e.seconds))
        return
    except Exception as e:
        log.exception("Failed to fetch channel %s", session.channel)
        await status_msg.edit_text(t("fetch_failed", lang, e=e))
        return

    if not messages:
        await status_msg.edit_text(t("no_messages", lang))
        return

    await update_status(t("fetched_n", lang, n=len(messages)))
    chunks = chunk_messages(messages)
    await update_status(t("chunked", lang, n=len(messages), c=len(chunks)))

    try:
        report = await analyze_channel(
            depth=session.depth,
            focus_areas=session.focus_areas,
            lang=lang.value,
            model_id=session.model_id,
            chunks=chunks,
            channel_title=stats["title"],
            channel_username=stats.get("username"),
            subscribers=stats.get("subscribers"),
            msg_count=len(messages),
            on_progress=update_status,
        )
    except Exception as e:
        log.exception("Analysis failed for %s", session.channel)
        await status_msg.edit_text(t("analysis_failed", lang, e=e))
        return

    # Save report
    report_path = await save_report(user_id, session.channel, report)

    # Record in DB
    await usage_repo.record_analysis(
        telegram_id=user_id,
        channel=session.channel,
        depth=session.depth.value,
        focus_areas=[a.value for a in session.focus_areas],
        model_id=session.model_id,
        stars_paid=stars_paid,
        report_path=report_path,
    )

    await update_status(t("sending_report", lang))

    parts = split_report(report)
    for part in parts:
        try:
            await message.answer(part, parse_mode="HTML")
        except Exception:
            await message.answer(part)

    # Offer file download
    await message.answer_document(
        FSInputFile(report_path, filename=f"{session.channel}_report.md"),
        caption=t("download_report", lang),
    )

    try:
        await status_msg.delete()
    except Exception:
        pass