roa2web-service-auto/reports-app/telegram-bot/app/bot/email_handlers.py

"""
Telegram bot handlers for email-based authentication flow
"""
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import (
    ContextTypes,
    ConversationHandler,
    CommandHandler,
    MessageHandler,
    CallbackQueryHandler,
    filters
)
import logging
from datetime import datetime
import asyncio

from app.auth.email_auth import (
    is_valid_email_format,
    verify_email_in_oracle,
    generate_email_code,
    generate_session_token,
    check_rate_limit,
    clear_rate_limit
)
from app.utils.email_service import get_email_service
from app.db.operations import (
    create_email_auth_code,
    get_email_auth_code,
    get_pending_email_code,
    mark_email_code_used,
    increment_failed_attempts,
    delete_user_email_codes,
    is_user_authenticated,
    link_user_to_oracle,
    create_or_update_user
)
from app.api.client import get_backend_client

logger = logging.getLogger(__name__)

# Conversation states
AWAITING_EMAIL, AWAITING_CODE, AWAITING_PASSWORD = range(3)

# Constants
MAX_CODE_ATTEMPTS = 3


# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

async def edit_login_message(
    context: ContextTypes.DEFAULT_TYPE,
    chat_id: int,
    text: str,
    reply_markup=None,
    parse_mode="Markdown"
):
    """
    Helper function to edit the login message stored in context.
    If message_id is not stored, creates a new message instead.
    """
    message_id = context.user_data.get('login_message_id')

    if message_id:
        try:
            await context.bot.edit_message_text(
                chat_id=chat_id,
                message_id=message_id,
                text=text,
                reply_markup=reply_markup,
                parse_mode=parse_mode
            )
        except Exception as e:
            logger.warning(f"Could not edit message {message_id}: {e}")
            # Fallback: send new message and update ID
            msg = await context.bot.send_message(
                chat_id=chat_id,
                text=text,
                reply_markup=reply_markup,
                parse_mode=parse_mode
            )
            context.user_data['login_message_id'] = msg.message_id
    else:
        # No message ID stored - create new message
        msg = await context.bot.send_message(
            chat_id=chat_id,
            text=text,
            reply_markup=reply_markup,
            parse_mode=parse_mode
        )
        context.user_data['login_message_id'] = msg.message_id


async def delete_login_message(context: ContextTypes.DEFAULT_TYPE, chat_id: int):
    """Delete the login message and clear the message_id from context"""
    message_id = context.user_data.get('login_message_id')

    if message_id:
        try:
            await context.bot.delete_message(chat_id=chat_id, message_id=message_id)
        except Exception as e:
            logger.warning(f"Could not delete message {message_id}: {e}")

        # Clear from context
        context.user_data.pop('login_message_id', None)


# ============================================================================
# ENTRY POINTS: /login command and action:login button
# ============================================================================

async def login_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """
    Handler pentru /login command
    Oferă opțiuni de autentificare: Email sau Web App
    """
    user = update.effective_user

    # Check dacă e deja autentificat
    if await is_user_authenticated(user.id):
        await update.message.reply_text(
            "Ești deja autentificat.\n\n"
            "Folosește:\n"
            "• /companies - Vezi companiile tale\n"
            "• /help - Comenzi disponibile\n"
            "• /unlink - Deautentifică-te"
        )
        return ConversationHandler.END

    # Check rate limiting (3 requests per hour)
    if not await check_rate_limit(f"login_{user.id}", max_attempts=3, window_minutes=60):
        await update.message.reply_text(
            "Prea multe încercări de autentificare.\n\n"
            "Te rugăm să aștepți 60 de minute înainte de a încerca din nou."
        )
        return ConversationHandler.END

    # Afișează opțiuni de autentificare
    keyboard = [
        [InlineKeyboardButton("Login cu Email + Parolă", callback_data="email_login")],
        [InlineKeyboardButton("Login din Web App", callback_data="web_login_info")],
        [InlineKeyboardButton("Anulează", callback_data="cancel")]
    ]
    reply_markup = InlineKeyboardMarkup(keyboard)

    # CREATE message and SAVE message_id
    msg = await update.message.reply_text(
        "Alege metoda de autentificare:",
        reply_markup=reply_markup,
        parse_mode="Markdown"
    )

    # Save message ID for future edits
    context.user_data['login_message_id'] = msg.message_id

    return AWAITING_EMAIL


async def action_login_callback(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """
    Handler pentru butonul Login din meniu (action:login)
    Oferă opțiuni de autentificare: Email sau Web App
    """
    query = update.callback_query
    user = update.effective_user

    logger.info(f"[EMAIL_AUTH] action_login_callback triggered for user {user.id}")

    await query.answer()

    # Check dacă e deja autentificat
    if await is_user_authenticated(user.id):
        await query.edit_message_text(
            "Ești deja autentificat.\n\n"
            "Folosește:\n"
            "• /companies - Vezi companiile tale\n"
            "• /help - Comenzi disponibile\n"
            "• /unlink - Deautentifică-te"
        )
        return ConversationHandler.END

    # Check rate limiting (3 requests per hour)
    if not await check_rate_limit(f"login_{user.id}", max_attempts=3, window_minutes=60):
        await query.edit_message_text(
            "Prea multe încercări de autentificare.\n\n"
            "Te rugăm să aștepți 60 de minute înainte de a încerca din nou."
        )
        return ConversationHandler.END

    # Afișează opțiuni de autentificare
    keyboard = [
        [InlineKeyboardButton("Login cu Email + Parolă", callback_data="email_login")],
        [InlineKeyboardButton("Login din Web App", callback_data="web_login_info")],
        [InlineKeyboardButton("Anulează", callback_data="cancel")]
    ]
    reply_markup = InlineKeyboardMarkup(keyboard)

    # EDIT existing menu message and SAVE message_id
    await query.edit_message_text(
        "Alege metoda de autentificare:",
        reply_markup=reply_markup,
        parse_mode="Markdown"
    )

    # Save message ID for future edits
    context.user_data['login_message_id'] = query.message.message_id

    return AWAITING_EMAIL


# ============================================================================
# CALLBACK: Email Login
# ============================================================================

async def email_login_callback(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Callback pentru butonul 'Login cu Email'"""
    query = update.callback_query
    user = update.effective_user

    logger.info(f"[EMAIL_AUTH] email_login_callback triggered for user {user.id}")

    await query.answer()

    # IMPORTANT: Salvează message_id înainte de a edita
    context.user_data['login_message_id'] = query.message.message_id

    # EDIT same message - remove buttons, ask for email
    await query.edit_message_text(
        text="Introdu adresa de email ROA:",
        parse_mode="Markdown"
    )

    return AWAITING_EMAIL


async def web_login_info_callback(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Info despre web app login"""
    query = update.callback_query
    await query.answer()

    await query.edit_message_text(
        "**Login din Web App**\n\n"
        "Pentru această metodă:\n\n"
        "1. Accesează aplicația web ROA2WEB\n"
        "2. Autentifică-te cu username + parolă\n"
        "3. Apasă butonul \"Link Telegram\"\n"
        "4. Copiază codul generat (8 caractere)\n"
        "5. Trimite-mi codul: /start ABC123XY\n\n"
        "Vei fi autentificat automat.",
        parse_mode="Markdown"
    )

    # IMPORTANT: Salvează message_id pentru ca /start să poată edita același mesaj
    context.user_data['web_login_message_id'] = query.message.message_id

    return ConversationHandler.END


# ============================================================================
# STATE: AWAITING_EMAIL
# ============================================================================

async def receive_email(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Handler pentru primirea email-ului"""
    email = update.message.text.strip().lower()
    user_id = update.effective_user.id

    # ȘTERG mesajul utilizatorului imediat (chat curat)
    try:
        await update.message.delete()
    except Exception as e:
        logger.warning(f"Could not delete email message: {e}")

    # Validare format email
    if not is_valid_email_format(email):
        # Show error in main message
        await edit_login_message(
            context=context,
            chat_id=update.effective_chat.id,
            text="Email invalid\n\nIntrodu o adresă validă (nume@domeniu.ro)",
            reply_markup=InlineKeyboardMarkup([
                [InlineKeyboardButton("Anulează", callback_data="cancel")]
            ])
        )
        return AWAITING_EMAIL

    # Check for existing pending code
    existing_code = await get_pending_email_code(user_id)
    if existing_code:
        # Delete old pending code
        await delete_user_email_codes(user_id)
        logger.info(f"Deleted existing pending code for user {user_id}")

    # EDIT login message to show loading
    await edit_login_message(
        context=context,
        chat_id=update.effective_chat.id,
        text="Verificare email...",
        reply_markup=None
    )

    try:
        # Verifică email în Oracle
        username = await verify_email_in_oracle(email)

        # IMPORTANT: Generic response to prevent email enumeration
        # We always say "code sent" even if email doesn't exist

        if username:
            # Email exists - generate and send code
            code = generate_email_code()

            # Save code in database
            code_saved = await create_email_auth_code(
                code=code,
                email=email,
                username=username,
                telegram_user_id=user_id,
                expiry_minutes=5
            )

            if not code_saved:
                await edit_login_message(
                    context=context,
                    chat_id=update.effective_chat.id,
                    text="Eroare la salvarea codului.\n\nIncearcă din nou cu /login"
                )
                return ConversationHandler.END

            # Send email (async with retry)
            email_service = get_email_service()
            email_sent = await email_service.send_auth_code(email, code, username)

            if not email_sent:
                logger.error(f"Failed to send email to {email}")
                # Don't reveal this to user - they'll timeout naturally

        # Wait 1 second for better UX (looks like verification happened)
        await asyncio.sleep(1)

        # ALWAYS show this message (prevent enumeration)
        # EDIT same message with success + buttons
        await edit_login_message(
            context=context,
            chat_id=update.effective_chat.id,
            text=f"Cod trimis pe {email}\n\nIntrodu codul primit pe email:",
            reply_markup=InlineKeyboardMarkup([
                [InlineKeyboardButton("Retrimite Cod", callback_data=f"resend:{email}")],
                [InlineKeyboardButton("Anulează", callback_data="cancel")]
            ])
        )

        # Save email in context for resend functionality
        context.user_data['pending_email'] = email
        context.user_data['pending_username'] = username

        return AWAITING_CODE

    except Exception as e:
        logger.error(f"Error in receive_email: {e}", exc_info=True)
        await edit_login_message(
            context=context,
            chat_id=update.effective_chat.id,
            text="Eroare internă.\n\nIncearcă din nou mai târziu."
        )
        return ConversationHandler.END


# ============================================================================
# STATE: AWAITING_CODE
# ============================================================================

async def receive_code(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Handler pentru primirea codului din email"""
    code = update.message.text.strip()
    user_id = update.effective_user.id

    # ȘTERG mesajul utilizatorului imediat (chat curat)
    try:
        await update.message.delete()
    except Exception as e:
        logger.warning(f"Could not delete code message: {e}")

    # Validare format cod (6 digits)
    if not (code.isdigit() and len(code) == 6):
        await edit_login_message(
            context=context,
            chat_id=update.effective_chat.id,
            text="Cod invalid\n\nIntrodu cele 6 cifre din email.",
            reply_markup=InlineKeyboardMarkup([
                [InlineKeyboardButton("Retrimite Cod", callback_data=f"resend:{context.user_data.get('pending_email', '')}")],
                [InlineKeyboardButton("Anulează", callback_data="cancel")]
            ])
        )
        return AWAITING_CODE

    # Verifică cod în DB
    try:
        code_data = await get_email_auth_code(code)

        if not code_data:
            # EDIT login message to show error
            await edit_login_message(
                context=context,
                chat_id=update.effective_chat.id,
                text="Cod invalid sau expirat\n\nIncearcă din nou cu /login"
            )
            return ConversationHandler.END

        # Verificări de securitate

        # 1. Check if already used
        if code_data['used']:
            await edit_login_message(
                context=context,
                chat_id=update.effective_chat.id,
                text="Cod deja folosit\n\nIncearcă din nou cu /login"
            )
            return ConversationHandler.END

        # 2. Check if expired
        if datetime.now() > code_data['expires_at']:
            await edit_login_message(
                context=context,
                chat_id=update.effective_chat.id,
                text="Cod expirat\n\nIncearcă din nou cu /login"
            )
            return ConversationHandler.END

        # 3. Check if belongs to this user
        if code_data['telegram_user_id'] != user_id:
            logger.warning(
                f"User {user_id} tried to use code belonging to "
                f"user {code_data['telegram_user_id']}"
            )
            await edit_login_message(
                context=context,
                chat_id=update.effective_chat.id,
                text="Cod invalid"
            )
            return ConversationHandler.END

        # 4. Check failed attempts (max 3)
        if code_data['failed_attempts'] >= MAX_CODE_ATTEMPTS:
            await edit_login_message(
                context=context,
                chat_id=update.effective_chat.id,
                text="Prea multe încercări greșite\n\nIncearcă din nou cu /login"
            )
            return ConversationHandler.END

        # Cod valid - Marchează ca folosit
        await mark_email_code_used(code)

        # Salvează date verificate în context
        context.user_data['verified_username'] = code_data['oracle_username']
        context.user_data['verified_email'] = code_data['email']
        context.user_data['session_token'] = generate_session_token(
            user_id,
            code_data['email']
        )

        # EDIT same message - ask for password
        await edit_login_message(
            context=context,
            chat_id=update.effective_chat.id,
            text="Cod validat!\n\nIntroduci parola ROA:",
            reply_markup=None  # No buttons for security
        )

        return AWAITING_PASSWORD

    except Exception as e:
        logger.error(f"Error validating code: {e}", exc_info=True)
        await edit_login_message(
            context=context,
            chat_id=update.effective_chat.id,
            text="Eroare la validarea codului.\n\nIncearcă din nou."
        )
        return ConversationHandler.END


async def resend_code_callback(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Retrimite codul pe email"""
    query = update.callback_query
    await query.answer("Retrimitem codul...")

    # Extract email from callback data
    callback_data = query.data  # Format: "resend:email@example.com"
    if not callback_data.startswith("resend:"):
        await edit_login_message(
            context=context,
            chat_id=update.effective_chat.id,
            text="Eroare\n\nIncearcă din nou cu /login"
        )
        return ConversationHandler.END

    email = callback_data.split(":", 1)[1]
    user_id = update.effective_user.id

    # Check rate limiting for resend (max 2 per 10 minutes)
    if not await check_rate_limit(f"resend_{user_id}", max_attempts=2, window_minutes=10):
        await edit_login_message(
            context=context,
            chat_id=update.effective_chat.id,
            text="Prea multe solicitări\n\nAșteaptă 10 minute."
        )
        return ConversationHandler.END

    # Get username from context or re-verify
    username = context.user_data.get('pending_username')

    if not username:
        username = await verify_email_in_oracle(email)
        if not username:
            await edit_login_message(
                context=context,
                chat_id=update.effective_chat.id,
                text="Eroare\n\nIncearcă din nou cu /login"
            )
            return ConversationHandler.END

    # Delete old code and generate new one
    await delete_user_email_codes(user_id)

    code = generate_email_code()

    # Save new code
    await create_email_auth_code(
        code=code,
        email=email,
        username=username,
        telegram_user_id=user_id,
        expiry_minutes=5
    )

    # Send email
    email_service = get_email_service()
    await email_service.send_auth_code(email, code, username)

    # FIX BUG: EDIT message and KEEP buttons!
    await edit_login_message(
        context=context,
        chat_id=update.effective_chat.id,
        text=f"Cod retrimis pe {email}\n\nIntrodu codul primit pe email:",
        reply_markup=InlineKeyboardMarkup([
            [InlineKeyboardButton("Retrimite Cod", callback_data=f"resend:{email}")],
            [InlineKeyboardButton("Anulează", callback_data="cancel")]
        ])
    )

    return AWAITING_CODE


# ============================================================================
# STATE: AWAITING_PASSWORD
# ============================================================================

async def receive_password(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Handler pentru primirea parolei Oracle"""
    password = update.message.text.strip()
    user_id = update.effective_user.id

    # Șterge IMEDIAT mesajul cu parola (securitate)
    try:
        await update.message.delete()
        logger.info(f"Password message deleted for user {user_id}")
    except Exception as e:
        logger.warning(f"Could not delete password message: {e}")

    # Get verified data from context
    username = context.user_data.get('verified_username')
    email = context.user_data.get('verified_email')
    session_token = context.user_data.get('session_token')

    if not all([username, email, session_token]):
        await edit_login_message(
            context=context,
            chat_id=update.effective_chat.id,
            text="Sesiune expirată\n\nIncearcă din nou cu /login"
        )
        return ConversationHandler.END

    # EDIT login message to show loading
    await edit_login_message(
        context=context,
        chat_id=update.effective_chat.id,
        text="Verificare...",
        reply_markup=None
    )

    try:
        # Call backend endpoint pentru verificare parolă + JWT
        backend_client = get_backend_client()

        response = await backend_client.login_with_email(
            email=email,
            password=password,
            telegram_user_id=user_id,
            session_token=session_token
        )

        if not response.get('success'):
            await edit_login_message(
                context=context,
                chat_id=update.effective_chat.id,
                text="Credențiale invalide\n\nParolă incorectă sau cont inactiv.\n\nIncearcă din nou cu /login"
            )
            return ConversationHandler.END

        # Success - Salvează user în telegram_users
        # First create or update user record
        await create_or_update_user(
            telegram_user_id=user_id,
            username=update.effective_user.username,
            first_name=update.effective_user.first_name,
            last_name=update.effective_user.last_name
        )

        # Then link to Oracle
        from datetime import datetime, timedelta
        token_expires_at = datetime.now() + timedelta(minutes=30)  # Default expiry

        await link_user_to_oracle(
            telegram_user_id=user_id,
            oracle_username=response['username'],
            jwt_token=response['access_token'],
            jwt_refresh_token=response['refresh_token'],
            token_expires_at=token_expires_at
        )

        # Clear rate limits on successful auth
        clear_rate_limit(f"login_{user_id}")
        clear_rate_limit(f"resend_{user_id}")

        # Get session and active company BEFORE editing message
        from app.agent.session import get_session_manager
        from app.bot.menus import create_main_menu, pad_message_for_wide_buttons

        session_manager = get_session_manager()
        session = await session_manager.get_or_create_session(user_id)
        company = session.get_active_company()

        company_name = company['name'] if company else None
        company_cui = company.get('cui') if company else None

        # Create main menu keyboard
        keyboard = create_main_menu(
            company_name=company_name,
            company_cui=company_cui,
            is_authenticated=True,  # Now authenticated
            cache_enabled=True  # Default enabled
        )

        # Menu message with company info
        companies_count = len(response.get('companies', []))

        if company_name:
            menu_text = f"{company_name}"
        else:
            menu_text = f"Companii disponibile: {companies_count}\n\nSelectează o companie pentru a continua"

        menu_message = pad_message_for_wide_buttons(menu_text)

        # EDIT login message to show menu (no deletion, direct edit)
        await edit_login_message(
            context=context,
            chat_id=update.effective_chat.id,
            text=menu_message,
            reply_markup=keyboard
        )

        # Clear sensitive data from context
        context.user_data.clear()

        logger.info(f"User {user_id} authenticated successfully via email")

        return ConversationHandler.END

    except Exception as e:
        logger.error(f"Error during password verification: {e}", exc_info=True)
        await edit_login_message(
            context=context,
            chat_id=update.effective_chat.id,
            text="Eroare la autentificare.\n\nIncearcă din nou cu /login"
        )
        return ConversationHandler.END


# ============================================================================
# CANCEL HANDLER
# ============================================================================

async def cancel_login(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Cancel conversation"""

    # EDIT login message to show cancellation (don't delete)
    if update.callback_query:
        # Called from button
        await edit_login_message(
            context=context,
            chat_id=update.effective_chat.id,
            text="Login anulat",
            reply_markup=None
        )
    elif update.message:
        # Called from /cancel command
        await edit_login_message(
            context=context,
            chat_id=update.effective_chat.id,
            text="Login anulat",
            reply_markup=None
        )

    # Clear context
    context.user_data.clear()

    return ConversationHandler.END


async def conversation_timeout(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Handler for conversation timeout"""

    # EDIT login message to show timeout
    await edit_login_message(
        context=context,
        chat_id=update.effective_chat.id,
        text="Sesiune expirată\n\nConversația a expirat după 5 minute.\n\nIncearcă din nou cu /login"
    )

    # Clear context
    context.user_data.clear()

    return ConversationHandler.END


# ============================================================================
# CONVERSATION HANDLER SETUP
# ============================================================================

email_login_handler = ConversationHandler(
    entry_points=[
        CommandHandler('login', login_command),
        CallbackQueryHandler(action_login_callback, pattern='^action:login$'),
        CallbackQueryHandler(email_login_callback, pattern='^email_login$')
    ],
    states={
        AWAITING_EMAIL: [
            MessageHandler(filters.TEXT & ~filters.COMMAND, receive_email)
        ],
        AWAITING_CODE: [
            MessageHandler(filters.TEXT & ~filters.COMMAND, receive_code),
            CallbackQueryHandler(resend_code_callback, pattern='^resend:')
        ],
        AWAITING_PASSWORD: [
            MessageHandler(filters.TEXT & ~filters.COMMAND, receive_password)
        ],
    },
    fallbacks=[
        CommandHandler('cancel', cancel_login),
        CallbackQueryHandler(cancel_login, pattern='^cancel$'),
        CallbackQueryHandler(web_login_info_callback, pattern='^web_login_info$')
    ],
    per_message=False,  # Track conversation per user, not per message
    allow_reentry=True,  # Allow starting new conversation even if previous one is active
    name="email_login_conversation"
)