bot.py

"""
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.")