""" Bot de Telegram: quitar fondo + mejorar imagen, con 10 usos gratis por día por persona, y 10 Stars por cada imagen extra una vez agotados. Corre como una tarea en segundo plano dentro del mismo servidor FastAPI (remove-bg-api), así no hace falta un Space nuevo. """ import os import io import sqlite3 import uuid import asyncio from datetime import date import httpx from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup, LabeledPrice from telegram.request import HTTPXRequest from telegram.ext import ( Application, CommandHandler, MessageHandler, CallbackQueryHandler, PreCheckoutQueryHandler, ContextTypes, filters, ) BOT_TOKEN = os.environ.get("TELEGRAM_BOT_TOKEN", "") REMOVE_BG_URL = "https://zgrafic-remove-bg-api.hf.space/remove-bg" ENHANCE_URL = "https://zgrafic-enhance-photo-api.hf.space/enhance" FREE_USES_PER_DAY = 10 STARS_PER_IMAGE = 10 DB_PATH = "usage.db" # ---------- Base de datos: cuántos usos gratis lleva cada persona hoy ---------- # Aviso: en el plan gratuito de Hugging Face, este archivo puede borrarse # si el servidor se reinicia por completo (no hay almacenamiento # garantizado entre reinicios). Para un sistema de cobro 100% confiable, # lo ideal a futuro es una base de datos externa (ej. Supabase). def get_db(): conn = sqlite3.connect(DB_PATH) conn.execute(""" CREATE TABLE IF NOT EXISTS usage ( user_id INTEGER PRIMARY KEY, day TEXT, used INTEGER ) """) return conn def get_free_uses_left(user_id: int) -> int: today = date.today().isoformat() conn = get_db() row = conn.execute("SELECT day, used FROM usage WHERE user_id=?", (user_id,)).fetchone() if row is None or row[0] != today: conn.execute( "INSERT OR REPLACE INTO usage (user_id, day, used) VALUES (?, ?, 0)", (user_id, today), ) conn.commit() conn.close() return FREE_USES_PER_DAY conn.close() return max(0, FREE_USES_PER_DAY - row[1]) def register_free_use(user_id: int): today = date.today().isoformat() conn = get_db() conn.execute( "INSERT INTO usage (user_id, day, used) VALUES (?, ?, 1) " "ON CONFLICT(user_id) DO UPDATE SET day=excluded.day, " "used = CASE WHEN usage.day=excluded.day THEN usage.used + 1 ELSE 1 END", (user_id, today), ) conn.commit() conn.close() # ---------- Pedidos pendientes de pago (en memoria, viven poco tiempo) ---------- pending_jobs = {} # payload -> {"file_id":..., "mode":..., "chat_id":...} # ---------- Handlers ---------- async def start(update: Update, context: ContextTypes.DEFAULT_TYPE): await update.message.reply_text( "¡Hola! 👋 Mandame una foto y elegí qué querés hacer:\n\n" "✂️ Quitar el fondo\n" "✨ Mejorar la calidad\n\n" f"Tenés {FREE_USES_PER_DAY} usos gratis por día. " f"Después de eso, cada imagen nueva cuesta {STARS_PER_IMAGE} ⭐." ) async def on_photo(update: Update, context: ContextTypes.DEFAULT_TYPE): photo = update.message.photo[-1] keyboard = InlineKeyboardMarkup([ [ InlineKeyboardButton("✂️ Quitar fondo", callback_data=f"mode:remove:{photo.file_id}"), InlineKeyboardButton("✨ Mejorar imagen", callback_data=f"mode:enhance:{photo.file_id}"), ] ]) await update.message.reply_text("¿Qué querés hacer con esta foto?", reply_markup=keyboard) async def on_mode_chosen(update: Update, context: ContextTypes.DEFAULT_TYPE): query = update.callback_query await query.answer() _, mode, file_id = query.data.split(":", 2) user_id = query.from_user.id chat_id = query.message.chat_id free_left = get_free_uses_left(user_id) if free_left > 0: await query.edit_message_text(f"Procesando… (te quedan {free_left - 1} usos gratis hoy)") register_free_use(user_id) await process_and_send(context, chat_id, file_id, mode) else: payload = str(uuid.uuid4()) pending_jobs[payload] = {"file_id": file_id, "mode": mode, "chat_id": chat_id} title = "Quitar fondo" if mode == "remove" else "Mejorar imagen" await context.bot.send_invoice( chat_id=chat_id, title=title, description="Ya usaste tus 10 imágenes gratis de hoy. Esta imagen cuesta " f"{STARS_PER_IMAGE} Stars.", payload=payload, provider_token="", # vacío: pago con Telegram Stars currency="XTR", prices=[LabeledPrice(title, STARS_PER_IMAGE)], ) await query.edit_message_text("Te mandé una factura para pagar con Telegram Stars ⭐") async def on_pre_checkout(update: Update, context: ContextTypes.DEFAULT_TYPE): query = update.pre_checkout_query if query.invoice_payload in pending_jobs: await query.answer(ok=True) else: await query.answer(ok=False, error_message="Este pedido ya no es válido, mandá la foto de nuevo.") async def on_successful_payment(update: Update, context: ContextTypes.DEFAULT_TYPE): payload = update.message.successful_payment.invoice_payload job = pending_jobs.pop(payload, None) if not job: return await update.message.reply_text("¡Pago recibido! Procesando tu imagen…") await process_and_send(context, job["chat_id"], job["file_id"], job["mode"]) async def process_and_send(context: ContextTypes.DEFAULT_TYPE, chat_id: int, file_id: str, mode: str): try: tg_file = await context.bot.get_file(file_id) image_bytes = bytes(await tg_file.download_as_bytearray()) url = REMOVE_BG_URL if mode == "remove" else ENHANCE_URL files = {"file": ("imagen.jpg", image_bytes, "image/jpeg")} data = {} if mode == "remove" else {"type": "foto", "scale": "2"} async with httpx.AsyncClient(timeout=420.0) as client: resp = await client.post(url, files=files, data=data) if resp.status_code != 200: await context.bot.send_message(chat_id, "No se pudo procesar la imagen. Probá de nuevo en un rato.") return await context.bot.send_document( chat_id, document=io.BytesIO(resp.content), filename="resultado.png", ) except Exception as e: await context.bot.send_message(chat_id, f"Ocurrió un error procesando la imagen: {e}") def build_bot_app(): # Timeouts más largos: en el arranque del contenedor la red puede # tardar unos segundos en estar del todo lista. request = HTTPXRequest( connect_timeout=30.0, read_timeout=30.0, write_timeout=30.0, pool_timeout=30.0, ) application = Application.builder().token(BOT_TOKEN).request(request).build() application.add_handler(CommandHandler("start", start)) application.add_handler(MessageHandler(filters.PHOTO, on_photo)) application.add_handler(CallbackQueryHandler(on_mode_chosen, pattern=r"^mode:")) application.add_handler(PreCheckoutQueryHandler(on_pre_checkout)) application.add_handler(MessageHandler(filters.SUCCESSFUL_PAYMENT, on_successful_payment)) return application async def run_bot(): if not BOT_TOKEN: print("TELEGRAM_BOT_TOKEN no está configurado, el bot no se inicia.") return # Reintenta el arranque unas cuantas veces, por si la red del # contenedor todavía no está lista en el primer intento. for attempt in range(1, 6): try: application = build_bot_app() await application.initialize() await application.start() await application.updater.start_polling() print("Bot de Telegram conectado correctamente.") return except Exception as e: print(f"Intento {attempt}/5 de conectar el bot falló: {e}") await asyncio.sleep(10) print("No se pudo conectar el bot de Telegram después de varios intentos.")