telegram_bot.py

import os
import json
import logging
import asyncio
from datetime import date
from telegram import Update
from telegram.ext import ApplicationBuilder, CommandHandler, MessageHandler, filters, ContextTypes, ConversationHandler

# Import from app.py
from app import (
    fetch_unread_emails, summarize_email, draft_reply_email,
    parse_schedule_prompt, create_calendar_event,
    build, CHATGPT_MODEL, client, get_auth_flow
)

logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO)

# --- AUTH CONVERSATION STATES ---
EMAIL, AUTH_CODE = range(2)

# --- HELPERS ---
def get_creds(context):
    return context.user_data.get('creds')

# --- HANDLER FUNCTIONS ---
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    await update.message.reply_text("👋 Welcome! Please enter your **Gmail Address**:")
    return EMAIL

async def get_email(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_email = update.message.text
    context.user_data['email'] = user_email
    flow = get_auth_flow()
    auth_url, _ = flow.authorization_url(prompt='consent', access_type='offline')
    await update.message.reply_text(f"Login here: [Google Login]({auth_url})\nPaste the code/URL below.", parse_mode='Markdown')
    return AUTH_CODE

async def get_code(update: Update, context: ContextTypes.DEFAULT_TYPE):
    raw_text = update.message.text
    code = raw_text.split("code=")[1].split("&")[0] if "code=" in raw_text else raw_text
    try:
        flow = get_auth_flow()
        flow.fetch_token(code=code)
        context.user_data['creds'] = flow.credentials
        await update.message.reply_text("✅ Auth successful! Try /inbox.")
        return ConversationHandler.END
    except Exception as e:
        await update.message.reply_text(f"❌ Auth failed: {str(e)}")
        return ConversationHandler.END

async def inbox_handler(update: Update, context: ContextTypes.DEFAULT_TYPE):
    creds = get_creds(context)
    if not creds:
        await update.message.reply_text("❌ Please /start first.")
        return
    service = build("gmail", "v1", credentials=creds)
    emails = fetch_unread_emails(service)
    if not emails:
        await update.message.reply_text("📭 No unread emails.")
        return
    summary_text = "🎯 *Summaries:*\n\n"
    for email in emails:
        summary = summarize_email(email['body'], email['from'], email['subject'])
        summary_text += f"{summary}\n\n"
        context.user_data[email['from'].lower()] = email
    await update.message.reply_text(summary_text, parse_mode='Markdown')

async def reply_handler(update: Update, context: ContextTypes.DEFAULT_TYPE):
    if not context.args:
        await update.message.reply_text("Usage: `/reply Name`")
        return
    target = " ".join(context.args).lower()
    found = next((v for k, v in context.user_data.items() if target in str(k)), None)
    if found:
        reply = draft_reply_email(found['subject'], found['body'], "AI Assistant")
        await update.message.reply_text(f"📝 *Draft:* \n\n`{reply}`", parse_mode='Markdown')
    else:
        await update.message.reply_text("❌ Sender not found.")

async def schedule_handler(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_prompt = " ".join(context.args)
    creds = get_creds(context)
    if not user_prompt or not creds:
        await update.message.reply_text("Usage: `/schedule meeting at 2pm` (and ensure you are logged in)")
        return
    json_res = parse_schedule_prompt(user_prompt, date.today(), "UTC")
    try:
        event_data = json.loads(json_res)
        result = create_calendar_event(creds, event_data, "UTC")
        await update.message.reply_text(result, parse_mode='Markdown')
    except Exception as e:
        await update.message.reply_text(f"❌ Error: {str(e)}")

# --- MAIN RUNNER ---
# if __name__ == '__main__':
#     token = os.environ.get("TELEGRAM_BOT_TOKEN")
    
#     # 1. Define the ConversationHandler AFTER the functions are defined
#     auth_conv = ConversationHandler(
#         entry_points=[CommandHandler('start', start)],
#         states={
#             EMAIL: [MessageHandler(filters.TEXT & ~filters.COMMAND, get_email)],
#             AUTH_CODE: [MessageHandler(filters.TEXT & ~filters.COMMAND, get_code)],
#         },
#         fallbacks=[CommandHandler('start', start)],
#     )

#     # 2. Build application
#     application = ApplicationBuilder().token(token).build()
    
#     # 3. Add handlers
#     application.add_handler(auth_conv)
#     application.add_handler(CommandHandler('inbox', inbox_handler))
#     application.add_handler(CommandHandler('reply', reply_handler))
#     application.add_handler(CommandHandler('schedule', schedule_handler))

#     print("✅ Bot is live and polling...")
#     application.run_polling(drop_pending_updates=True)



# --- At the bottom of telegram_bot.py ---

if __name__ == '__main__':
    import time
    
    token = os.environ.get("TELEGRAM_BOT_TOKEN")
    if not token:
        print("❌ CRITICAL ERROR: TELEGRAM_BOT_TOKEN is missing from environment variables.")
        exit(1)

    # 1. Wait for the container network to actually wake up
    print("⏳ Waiting 10 seconds for network stability...")
    time.sleep(10)

    # 2. Define handlers
    auth_conv = ConversationHandler(
        entry_points=[CommandHandler('start', start)],
        states={
            EMAIL: [MessageHandler(filters.TEXT & ~filters.COMMAND, get_email)],
            AUTH_CODE: [MessageHandler(filters.TEXT & ~filters.COMMAND, get_code)],
        },
        fallbacks=[CommandHandler('start', start)],
    )

    # 3. Initialize Application
    application = ApplicationBuilder().token(token).build()
    application.add_handler(auth_conv)
    application.add_handler(CommandHandler('inbox', inbox_handler))
    application.add_handler(CommandHandler('reply', reply_handler))
    application.add_handler(CommandHandler('schedule', schedule_handler))

    # 4. Start with a Retry Loop
    print("🚀 Attempting to connect to Telegram...")
    
    max_retries = 5
    for attempt in range(max_retries):
        try:
            print(f"🔄 Connection attempt {attempt + 1}/{max_retries}...")
            application.run_polling(drop_pending_updates=True)
            break # If it succeeds, break the loop
        except Exception as e:
            print(f"⚠️ Connection failed: {e}")
            if attempt < max_retries - 1:
                print("😴 Waiting 15 seconds before retrying...")
                time.sleep(15)
            else:
                print("❌ All connection attempts failed. Please check Hugging Face network settings.")