marzban_tg_bot/handlers/payment.py

from aiogram import Router, F
from aiogram.types import Message, CallbackQuery, LabeledPrice, InlineKeyboardMarkup, InlineKeyboardButton
from aiogram.fsm.context import FSMContext
from datetime import datetime
import logging

from config import CONFIG, PLANS
from database import db
from marzban import marzban
from keyboards import main_keyboard

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

async def grant_subscription(user_id: int, plan_id: str, promo_code: str, amount: int, bonus_days: int = 0):
    total_days = 0
    data_limit = 0
    plan_name = "VIP Sub"
    marzban_days = 0
    
    if plan_id:
        plan = PLANS[plan_id]
        total_days = plan['days'] + bonus_days
        marzban_days = total_days
        data_limit = plan['data_limit']
        plan_name = plan['name']
    else:
        # VIP case without plan
        total_days = 365 * 99 # For DB (99 years)
        marzban_days = 0 # For Marzban (Unlimited/None)
        data_limit = 0 # Unlimited
        plan_name = "VIP"

    user = await db.get_user(user_id)
    
    tg_username = user['username']
    note = f"@{tg_username}" if tg_username else ""
    
    sub_until = user['subscription_until']
    if isinstance(sub_until, str):
        try:
            sub_until = datetime.strptime(sub_until, '%Y-%m-%d %H:%M:%S')
        except ValueError:
            try:
                 sub_until = datetime.strptime(sub_until, '%Y-%m-%d %H:%M:%S.%f')
            except: pass

    # Marzban
    try:
        marzban_username = user['marzban_username']
        resp = None
        
        # Если есть подписка и она активна (и не бесконечна, хотя тут не важно)
        is_sub_active = sub_until and sub_until > datetime.now()
        
        if is_sub_active:
            logger.info(f"Attempting to modify existing user: {marzban_username}")
            resp = await marzban.modify_user(
                marzban_username,
                data_limit,
                marzban_days,
                note
            )
            
            if isinstance(resp, dict) and resp.get('detail') == 'User not found':
                logger.info(f"User {marzban_username} missing in Marzban, re-creating...")
                resp = await marzban.create_user(
                    marzban_username,
                    data_limit,
                    marzban_days,
                    note
                )
        else:
            logger.info(f"Creating/Reactivating user: {marzban_username}")
            resp = await marzban.create_user(
                marzban_username,
                data_limit,
                marzban_days,
                note
            )
    except Exception as e:
        logger.error(f"Marzban error in grant_subscription: {e}")
    
    # DB
    await db.update_subscription(user_id, total_days, data_limit)
    await db.add_payment(
        user_id,
        plan_id or "vip", # Store 'vip' as plan name/id
        amount,
        promo_code
    )
    
    if promo_code:
        await db.decrement_promo_usage(promo_code)
        
    # Return dummy plan dict for display
    return {'name': plan_name, 'days': total_days, 'data_limit': data_limit}, total_days

@router.callback_query(F.data == "pay_now")
async def process_payment(callback: CallbackQuery, state: FSMContext):
    data = await state.get_data()
    plan_id = data.get('selected_plan')
    is_vip = data.get('is_unlimited_promo')
    
    if not plan_id and not is_vip:
        await callback.answer("❌ Сессия истекла. Пожалуйста, выберите тариф снова.", show_alert=True)
        return
        
    # Если VIP без плана, ставим дефолтные значения
    plan_name = "VIP"
    plan_days = 3650
    plan_limit = 0
    base_price = 0
    
    if plan_id:
        plan = PLANS[plan_id]
        plan_name = plan['name']
        plan_days = plan['days']
        plan_limit = plan['data_limit']
        base_price = plan['price']

    final_price = int(data.get('final_price', base_price))
    promo_code = data.get('promo_code')
    bonus_days = data.get('bonus_days', 0)
    
    if final_price <= 0:
        sticky_msg = ""
        if promo_code:
             p_data = await db.get_promo_code(promo_code)
             # Use explicit key access with fallback logic if needed, but keys exist
             if p_data and p_data['is_sticky']:
                  user = await db.get_user(callback.from_user.id)
                  u_disc = user['personal_discount'] if user and user['personal_discount'] else 0
                  
                  if p_data['discount'] > u_disc:
                       await db.set_user_discount(callback.from_user.id, p_data['discount'])
                       sticky_msg = "\n🔐 Скидка закреплена навсегда!"

        plan, date_days = await grant_subscription(callback.from_user.id, plan_id, promo_code, 0, bonus_days)
        await callback.message.edit_text(
            f"✅ Подписка активирована бесплатно!\n\n"
            f"План: {plan['name']}\n"
            f"Срок: {date_days} дней\n"
            f"Трафик: {plan['data_limit'] if plan['data_limit'] > 0 else '∞'} ГБ\n"
            f"{sticky_msg}\n"
            f"Настройте подключение в меню: 📊 Моя подписка",
            reply_markup=main_keyboard(callback.from_user.id in CONFIG["ADMIN_IDS"])
        )
        await state.clear()
    else:
        # Создаем инвойс для Telegram Stars
        await callback.message.answer_invoice(
            title=f"Подписка VPN - {plan['name']}",
            description=f"Трафик: {plan['data_limit']} ГБ на {plan['days']} дней",
            payload=f"{plan_id}:{data.get('promo_code', '')}",
            currency="XTR",  # Telegram Stars
            prices=[LabeledPrice(label=plan['name'], amount=final_price)],
            reply_markup=InlineKeyboardMarkup(inline_keyboard=[
                [InlineKeyboardButton(text="💳 Оплатить", pay=True)]
            ])
        )

@router.pre_checkout_query()
async def checkout_process(pre_checkout_query):
    await pre_checkout_query.answer(ok=True)

@router.message(F.successful_payment)
async def successful_payment(message: Message):
    payment = message.successful_payment
    plan_id, promo_code = payment.invoice_payload.split(":")
    
    # We can reuse grant_subscription helper
    promo_code = promo_code if promo_code else None
    
    bonus_days = 0
    sticky_text = ""
    
    if promo_code:
        # Fetch actual promo details
        promo_data = await db.get_promo_code(promo_code)
        if promo_data:
            bonus_days = promo_data['bonus_days']
            
            # STICKY LOGIC
            # Access by key (sqlite3.Row has no .get method)
            is_sticky = False
            try:
                is_sticky = promo_data['is_sticky']
            except IndexError:
                pass # Column missing?
                
            if is_sticky:
                user = await db.get_user(message.from_user.id)
                current_discount = user['personal_discount'] if user and user['personal_discount'] else 0
                new_discount = promo_data['discount']
                
                if new_discount > current_discount:
                    await db.set_user_discount(message.from_user.id, new_discount)
                    sticky_text = f"\n🔐 <b>Скидка {new_discount}% закреплена за вами НАВСЕГДА!</b>"

    plan, date_days = await grant_subscription(
        message.from_user.id,
        plan_id,
        promo_code,
        payment.total_amount,
        bonus_days
    )
    
    await message.answer(
        f"✅ Оплата успешна!\n\n"
        f"Ваша подписка активирована на {date_days} дней.\n"
        f"Трафик: {plan['data_limit']} ГБ\n"
        f"{sticky_text}\n"
        f"Получите конфигурацию через меню: 📊 Моя подписка",
        reply_markup=main_keyboard(message.from_user.id in CONFIG["ADMIN_IDS"]),
        parse_mode="HTML"
    )