bot/checker.py

import asyncio
import datetime
import random
from pyrogram import enums
from pyrogram.raw import functions
from pyrogram.errors import FloodWait, UsernameInvalid, UsernameOccupied, RPCError
from pyrogram.types import InlineKeyboardMarkup, InlineKeyboardButton

from bot.state import state
from bot.database import db
from utils.logger import log
from bot.alerts import alert_admin

ACCOUNT_LOCKS = {}

def get_acc_lock(acc_index):
    if acc_index not in ACCOUNT_LOCKS:
        ACCOUNT_LOCKS[acc_index] = asyncio.Lock()
    return ACCOUNT_LOCKS[acc_index]

async def check_username(username: str, acc) -> str:
    if not acc or not acc.client.is_connected:
        return "ERROR"

    try:
        now = datetime.datetime.utcnow()
        if not hasattr(acc, 'hour_start_time') or (now - acc.hour_start_time).total_seconds() > 3600:
            acc.hour_start_time = now
            acc.checks_this_hour = 0
            
        acc.checks_this_hour += 1
        result = await acc.client.invoke(functions.account.CheckUsername(username=username))
        return "AVAILABLE" if result else "TAKEN"

    except UsernameInvalid:
        return "INVALID"
    except UsernameOccupied:
        return "TAKEN"
    except FloodWait as e:
        log(f"⚠️ Account {acc.index} hit FloodWait for {e.value} seconds.")
        acc.status = "cooling"
        acc.cooldown_until = datetime.datetime.utcnow() + datetime.timedelta(seconds=e.value)
        return "ERROR"
    except RPCError as e:
        if "FROZEN" in str(e).upper():
            log(f"🧊 Account {acc.index} is FROZEN by Telegram.")
            acc.status = "frozen"
        else:
            log(f"⚠️ Account {acc.index} RPC Error on @{username}: {e}")
        return "ERROR"
    except Exception as e:
        if "FROZEN" in str(e).upper():
            log(f"🧊 Account {acc.index} is FROZEN by Telegram.")
            acc.status = "frozen"
        else:
            log(f"⚠️ Account {acc.index} Exception on @{username}: {e}")
        return "ERROR"

async def trigger_burst(username: str) -> str:
    state["burst_last_triggered"] = datetime.datetime.utcnow()
    state["cancel_burst"] = False 
    
    log(f"🚀 BURST MODE ACTIVATED FOR @{username}!")
    
    if not state["autoclaim_on"]:
        await alert_admin(f"🚨 **TARGET DROPPED!**\n`@{username}` is available, but Autoclaim is OFF!")
        return "FAILED"

    available_accounts = [a for a in state["account_pool"] if a.status in ["active", "paused"] and a.client.is_connected]
    
    if not available_accounts:
        log("❌ Burst failed: No connected accounts available.")
        return "FAILED"

    random.shuffle(available_accounts)

    pref_acc_idx = state.get("preferred_channel_account")
    if pref_acc_idx is not None:
        available_accounts.sort(key=lambda a: 0 if a.index == pref_acc_idx else 1)

    abort_markup = InlineKeyboardMarkup([[InlineKeyboardButton("🛑 ABORT GATLING-GUN", callback_data="abort_burst", style=enums.ButtonStyle.DANGER)]])
    await alert_admin(f"🚀 <b>GATLING-GUN ENGAGED!</b>\nFiring at <code>@{username}</code> with {len(available_accounts)} accounts...", markup=abort_markup)

    result_status = "FAILED"
    for acc in available_accounts:
        if state.get("cancel_burst"):
            log(f"🛑 GATLING-GUN ABORTED BY ADMIN FOR @{username}.")
            await alert_admin(f"🛑 <b>BURST ABORTED</b> for <code>@{username}</code>.")
            break

        success = await autoclaim_flow(username, acc)
        
        if success == True:
            log(f"✅ Burst sequence finished. @{username} secured.")
            result_status = "CLAIMED"
            break
        elif success == "OCCUPIED":
            log(f"⚠️ Telegram is delaying release of @{username}. Aborting burst to save API limits.")
            result_status = "OCCUPIED"
            break
        elif success == "INVALID":
            log(f"🚫 @{username} is permanently invalid/banned by Telegram. Aborting burst.")
            result_status = "INVALID"
            break
            
    return result_status

async def autoclaim_flow(username: str, acc):
    priority = state.get("channel_claim_priority", "account_first")
    lock = get_acc_lock(acc.index)
    
    async with lock:
        # ⚡ ENTERPRISE REFACTOR: Single definitive state check inside the lock.
        if username in state.get("unavailable_set", set()):
            return True
            
        try:
            me = await acc.client.get_users("me")
            current_username = me.username
            acc.profile_username = current_username
        except Exception as e:
            if "FROZEN" in str(e).upper():
                acc.status = "frozen"
            log(f"⚠️ Could not fetch live profile for Acc {acc.index}, falling back to cache. Error: {e}")
            current_username = getattr(acc, 'profile_username', None)

        has_profile_username = bool(current_username)
        display_name = current_username if current_username else "UNKNOWN"

        log(f"⚡ Autoclaim triggered for @{username} on Acc {acc.index}. Priority: {priority}")
        success = False

        if priority == "channel_first":
            success = await try_claim_channel(username, acc)
            if not success:
                if has_profile_username:
                    log(f"⚠️ Skipped account claim for @{username} to prevent overwriting @{display_name}.")
                elif success != "OCCUPIED" and success != "INVALID":
                    success = await try_claim_account(username, acc)
        else:
            if not has_profile_username:
                success = await try_claim_account(username, acc)
                if not success and success != "OCCUPIED" and success != "INVALID":
                    success = await try_claim_channel(username, acc)
            else:
                log(f"⚠️ Skipped account claim for @{username} to prevent overwriting @{display_name}.")
                success = await try_claim_channel(username, acc)

        if success == True:
            state["unavailable_set"].add(username)
            state["total_claimed"] += 1
            await db.add_unavailable(username)
            await db.save_setting("total_claimed", state["total_claimed"])
            return True
            
        return success

async def get_or_create_standby_channel(acc):
    if hasattr(acc, 'standby_channel_id') and acc.standby_channel_id:
        return acc.standby_channel_id 

    try:
        async for dialog in acc.client.get_dialogs():
            if dialog.chat.type in [enums.ChatType.CHANNEL, enums.ChatType.SUPERGROUP] and dialog.chat.is_creator:
                if dialog.chat.title == "[STANDBY_SLOT]":
                    acc.standby_channel_id = dialog.chat.id
                    return dialog.chat.id
    except Exception as e:
        if "FROZEN" in str(e).upper():
            acc.status = "frozen"
        log(f"⚠️ Error searching dialogs for Acc {acc.index}: {e}")

    if acc.status == "frozen":
        return None

    try:
        log(f"🔧 Pre-warming new Standby Slot for Acc {acc.index}...")
        channel_result = await acc.client.create_channel(title="[STANDBY_SLOT]", description="Ready for claim.")
        
        if hasattr(channel_result, "chats"):
            acc.standby_channel_id = channel_result.chats[0].id
            return channel_result.chats[0].id
        else:
            acc.standby_channel_id = channel_result.id
            return channel_result.id
    except Exception as e:
        if "FROZEN" in str(e).upper():
            acc.status = "frozen"
        log(f"❌ Failed to create Standby Slot for Acc {acc.index}: {e}")
        return None

async def try_claim_channel(username: str, acc):
    standby_id = await get_or_create_standby_channel(acc)
    if not standby_id:
        return False

    try:
        peer = await acc.client.resolve_peer(standby_id)
        await acc.client.invoke(functions.channels.UpdateUsername(channel=peer, username=username))
        
        desc = state.get("custom_claim_message", "Secured")
        await acc.client.invoke(functions.channels.EditTitle(channel=peer, title=username))
        await acc.client.set_chat_description(standby_id, desc)
        
        acc.standby_channel_id = None 
        
        custom_txt = state.get("custom_success_text", "")
        gif_setting = state.get("custom_claim_gif", "")

        # ⚡ ENTERPRISE FIX: Robust Media Upload Handling with Fallbacks
        try:
            if gif_setting:
                try:
                    # Attempt 1: Try sending as Animation (GIF) via URL or File ID
                    await acc.client.send_animation(chat_id=standby_id, animation=gif_setting, caption=custom_txt)
                    log(f"✅ Successfully posted Welcome GIF to @{username}")
                except Exception as anim_err:
                    log(f"⚠️ Failed to send as animation: {anim_err}. Trying as video...")
                    try:
                        # Attempt 2: Fallback to Video format
                        await acc.client.send_video(chat_id=standby_id, video=gif_setting, caption=custom_txt)
                        log(f"✅ Successfully posted Welcome Video to @{username}")
                    except Exception as vid_err:
                        log(f"⚠️ Failed to send as video: {vid_err}. Falling back to text.")
                        # Attempt 3: Ultimate Fallback to just text
                        if custom_txt:
                            await acc.client.send_message(chat_id=standby_id, text=custom_txt, disable_web_page_preview=True)
                            log(f"✅ Successfully posted Welcome Text to @{username}")
            elif custom_txt:
                await acc.client.send_message(chat_id=standby_id, text=custom_txt, disable_web_page_preview=True)
                log(f"✅ Successfully posted Welcome Text to @{username}")
        except Exception as post_err:
            log(f"🚨 CRITICAL: Secured @{username}, but totally failed to post welcome message: {post_err}")
        
        text = f"🏆 <b>SUCCESS!</b>\nAutomatically claimed <code>@{username}</code> via Channel on <b>Account {acc.index}</b>!"
        if custom_txt: text += f"\n\n<blockquote>{custom_txt}</blockquote>"
        
        buttons = [
            [InlineKeyboardButton(f"👀 View @{username}", url=f"https://t.me/{username}", style=enums.ButtonStyle.PRIMARY)],
            [InlineKeyboardButton("🗑️ Undo (Delete Channel)", callback_data=f"undo_ch|{standby_id}|{username}|{acc.index}", style=enums.ButtonStyle.DANGER)],
            [InlineKeyboardButton("❌ Dismiss", callback_data=f"skip|{username}", style=enums.ButtonStyle.DEFAULT)]
        ]
        
        await alert_admin(text, markup=InlineKeyboardMarkup(buttons))
        return True
        
    except Exception as e:
        err_str = str(e).upper()
        log(f"❌ Channel autoclaim failed for @{username} on Acc {acc.index}: {e}")

        if "FROZEN" in err_str: 
            acc.status = "frozen"

        if "CHANNEL_INVALID" in err_str or "PEER_ID_INVALID" in err_str:
            acc.standby_channel_id = None 

        if "USERNAME_OCCUPIED" in err_str: return "OCCUPIED"
        if "USERNAME_INVALID" in err_str:
            state["unavailable_set"].add(username)
            await db.add_unavailable(username)
            return "INVALID"

        return False

async def try_claim_account(username: str, acc):
    try:
        await acc.client.invoke(functions.account.UpdateUsername(username=username))
        
        custom_txt = state.get("custom_success_text", "")
        text = f"🏆 <b>SUCCESS!</b>\nAutomatically claimed <code>@{username}</code> on <b>Account {acc.index}</b> profile!"
        if custom_txt: text += f"\n\n<blockquote>{custom_txt}</blockquote>"
        
        buttons = [
            [InlineKeyboardButton(f"👀 View @{username}", url=f"https://t.me/{username}", style=enums.ButtonStyle.PRIMARY)],
            [InlineKeyboardButton("🗑️ Undo (Remove Username)", callback_data=f"undo_acc|{username}|{acc.index}", style=enums.ButtonStyle.DANGER)],
            [InlineKeyboardButton("❌ Dismiss", callback_data=f"skip|{username}", style=enums.ButtonStyle.DEFAULT)]
        ]
        
        await alert_admin(text, markup=InlineKeyboardMarkup(buttons))
        return True
        
    except Exception as e:
        err_str = str(e).upper()
        log(f"❌ Account autoclaim failed for @{username} on Acc {acc.index}: {e}")

        if "FROZEN" in err_str: 
            acc.status = "frozen"
        
        if "USERNAME_OCCUPIED" in err_str: return "OCCUPIED"
        if "USERNAME_INVALID" in err_str:
            state["unavailable_set"].add(username)
            await db.add_unavailable(username)
            return "INVALID"
            
        return False