import os
import re
import time
import io
import asyncio
import aiohttp
from pyrogram import Client, filters
from pyrogram.types import InlineKeyboardMarkup, InlineKeyboardButton
from pyrogram.errors import MessageNotModified

from state import state
from logger import get_logger
import scraper

log = get_logger()

# --- CONFIG & AUTH ---
API_ID = int(os.environ.get("API_ID", 0))
API_HASH = os.environ.get("API_HASH", "")
BOT_TOKEN = os.environ.get("BOT_TOKEN", "")

# Keep the set for fast lookups if used elsewhere, but cast to list for Pyrogram's filter
ADMIN_IDS_SET = {int(x) for x in os.environ.get("ADMIN_IDS", "").split(",") if x}
ADMIN_IDS = list(ADMIN_IDS_SET)

app = Client("bot_session", api_id=API_ID, api_hash=API_HASH, bot_token=BOT_TOKEN)

# FIX: Pyrogram accepts list natively, bypassing the unhashable set error
is_admin = filters.user(ADMIN_IDS)

# --- DYNAMIC UI MENUS ---
def get_dynamic_menu():
    """Generates toggle switches based on the current in-memory state."""
    buttons = []
    
    if state.status == "STOPPED":
        buttons.append([InlineKeyboardButton("▶️ Start Scan", callback_data="toggle_start")])
    elif state.status == "RUNNING":
        buttons.append([
            InlineKeyboardButton("⏸ Pause", callback_data="toggle_pause"),
            InlineKeyboardButton("⏹ Stop", callback_data="toggle_stop")
        ])
    elif state.status == "PAUSED":
        buttons.append([
            InlineKeyboardButton("▶️ Resume", callback_data="toggle_resume"),
            InlineKeyboardButton("⏹ Stop", callback_data="toggle_stop")
        ])

    buttons.append([
        InlineKeyboardButton("📊 Live Status", callback_data="show_status"), 
        InlineKeyboardButton("📥 Export Data", callback_data="export_files")
    ])
    buttons.append([InlineKeyboardButton("⚙️ System Settings", callback_data="menu_settings")])
    
    return InlineKeyboardMarkup(buttons)

def get_settings_menu():
    return InlineKeyboardMarkup([
        [InlineKeyboardButton("💾 Load Words", callback_data="load_words"), InlineKeyboardButton("📏 Set Length", callback_data="set_min_length")],
        [InlineKeyboardButton("⚡ Set Speed", callback_data="set_speed"), InlineKeyboardButton("🔴 Clear Memory", callback_data="reset_confirm")],
        [InlineKeyboardButton("« Back to Main", callback_data="back_main")]
    ])

def generate_status_msg():
    qlen = state.queue.qsize() if state.queue else 0
    pct = (state.processed / state.total_words * 100) if state.total_words > 0 else 0
    
    run_str = f"🟢 {state.status}" if state.status == "RUNNING" else f"🔴 {state.status}"
    if state.status == "PAUSED": run_str = "⏸ PAUSED"
    
    bar_len = 20
    filled = int(bar_len * (state.processed / state.total_words)) if state.total_words > 0 else 0
    bar = "=" * filled + "-" * (bar_len - filled)
    
    return (
        f"**SYSTEM DASHBOARD**\n"
        f"━━━━━━━━━━━━━━━━━━━━━\n"
        f"**State:** `{run_str}`\n"
        f"**Queue:** `{qlen:,}` waiting\n"
        f"**Speed:** `{state.concurrency}` workers\n"
        f"**Proxies:** `{len(state.proxies)}` Active\n"
        f"**Progress:** `{pct:.2f}%`\n"
        f"`[{bar}]`\n"
        f"`{state.processed:,} / {state.total_words:,}`\n\n"
        f"**METRICS**\n"
        f"━━━━━━━━━━━━━━━━━━━━━\n"
        f"🔴 Taken      : `{state.counts.get('taken', 0):,}`\n"
        f"🚫 Unavail    : `{state.counts.get('unavailable', 0):,}`\n"
        f"💰 For Sale   : `{state.counts.get('forsale', 0):,}`\n"
        f"🔨 Auction    : `{state.counts.get('auction', 0):,}`\n"
        f"✅ Available  : `{state.counts.get('available', 0):,}`\n"
        f"🛒 Sold       : `{state.counts.get('sold', 0):,}`\n"
        f"⚠️ Errors     : `{state.counts.get('error', 0):,}`\n"
    )

# --- PROXY TESTER (Re-engineered Shared Session Management) ---
async def test_proxy(session, proxy_url):
    start = time.time()
    headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 Chrome/124.0 Safari/537.36"}
    try:
        async with session.get("https://fragment.com/", headers=headers, proxy=proxy_url, allow_redirects=False) as resp:
            if resp.status == 200:
                return proxy_url, round((time.time() - start) * 1000)
    except Exception:
        pass
    return proxy_url, None

# --- COMMANDS ---
@app.on_message(filters.command("start") & is_admin)
async def start_cmd(client, message):
    await message.reply_text("**Fragment Control System**", reply_markup=get_dynamic_menu())

@app.on_message(filters.command("proxies") & is_admin)
async def proxy_cmd(client, message):
    if not message.reply_to_message or not message.reply_to_message.document:
        return await message.reply_text("⚠️ Reply to a `proxies.txt` file with `/proxies`.")
        
    msg = await message.reply_text("⏳ Downloading proxies...")
    file_path = await message.reply_to_message.download(file_name="temp_proxies.txt")
    
    def read_proxies():
        with open(file_path, "r", encoding="utf-8") as f:
            return [line.strip() for line in f if line.strip()]
            
    raw_proxies = await asyncio.to_thread(read_proxies)
    if os.path.exists(file_path):
        os.remove(file_path)
        
    if not raw_proxies: 
        return await msg.edit_text("⚠️ No proxies found in file.")
        
    await msg.edit_text(f"🔍 Testing **{len(raw_proxies)}** proxies concurrently via global TCP pools...")
    
    timeout = aiohttp.ClientTimeout(total=10)
    connector = aiohttp.TCPConnector(limit=100, ssl=False)
    async with aiohttp.ClientSession(timeout=timeout, connector=connector) as shared_session:
        tasks = [test_proxy(shared_session, p) for p in raw_proxies]
        results = await asyncio.gather(*tasks)
    
    working = [p for p, lat in results if lat is not None]
    state.proxies = working 
    
    latency_board = sorted([(p, lat) for p, lat in results if lat is not None], key=lambda x: x[1])
    top_str = "\n".join([f"⏱ `{lat}ms` | {url.split('@')[-1] if '@' in url else url}" for url, lat in latency_board[:10]])
    
    report = (
        f"✅ **Proxy Test Complete**\n"
        f"━━━━━━━━━━━━━━━━\n"
        f"**Total Tested:** `{len(raw_proxies)}`\n"
        f"**Working & Saved:** `{len(working)}`\n\n"
        f"🏆 **Top Fastest Proxies:**\n{top_str if top_str else 'None'}"
    )
    await msg.edit_text(report, reply_markup=get_dynamic_menu())

@app.on_message(filters.command("upload") & is_admin)
async def upload_cmd(client, message):
    target_message = message if message.document else message.reply_to_message
    
    if not target_message or not target_message.document:
        return await message.reply_text("⚠️ Send a `.txt` file with the caption `/upload`, or reply to a file with `/upload`.")
        
    msg = await message.reply_text("⏳ Downloading `words.txt`...")
    
    try:
        file_path = await target_message.download()
        
        def safe_replace():
            if os.path.exists("words.txt"):
                os.remove("words.txt")
            os.replace(file_path, "words.txt")
            
        await asyncio.to_thread(safe_replace)
        await msg.edit_text("✅ **Uploaded!**\nGo to System Settings -> Load Words to add them to the queue.", reply_markup=get_dynamic_menu())
    except Exception as e:
        await msg.edit_text(f"❌ **Upload Error:** `{str(e)}`")
        log.error(f"Upload Error: {str(e)}")

@app.on_message(filters.command("check") & is_admin)
async def check_cmd(client, message):
    if len(message.command) < 2: return await message.reply_text("Usage: `/check <word>`")
    target = message.command[1].lower()
    msg = await message.reply_text(f"📸 Taking live screenshot of `@{target}` on Fragment...")
    
    url = f"https://fragment.com/username/{target}"
    img_filename = None
    
    try:
        def take_screenshot():
            from html2image import Html2Image
            hti = Html2Image(
                browser_executable='/usr/bin/chromium',
                custom_flags=['--no-sandbox', '--disable-gpu', '--hide-scrollbars', '--disable-dev-shm-usage']
            ) 
            hti.size = (1000, 750) 
            name = f"{target}_fragment.png"
            hti.screenshot(url=url, save_as=name)
            return name

        img_filename = await asyncio.to_thread(take_screenshot)
        await message.reply_document(document=img_filename, caption=f"**Live Web Snapshot:** `@{target}`")
        await msg.delete()
            
    except Exception as e:
        try: await msg.edit_text(f"❌ **Screenshot Error:**\n`{str(e)}`")
        except: pass
    finally:
        if img_filename and os.path.exists(img_filename):
            os.remove(img_filename)

@app.on_message(filters.command("html") & is_admin)
async def html_cmd(client, message):
    if len(message.command) < 2: return await message.reply_text("Usage: `/html <word>`")
    target = message.command[1].lower()
    msg = await message.reply_text(f"🔍 Fetching raw HTML for `@{target}`...")
    
    url = f"https://fragment.com/username/{target}"
    headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 Chrome/124.0 Safari/537.36"}
    
    try:
        # Failsafe IndexError checking fix
        proxy = state.proxies[0] if state.proxies else None
        
        timeout = aiohttp.ClientTimeout(total=15)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.get(url, headers=headers, proxy=proxy, allow_redirects=True) as resp:
                html_text = await resp.text()
                status_code = resp.status
                final_url = str(resp.url)
                
        f = io.BytesIO(html_text.encode('utf-8'))
        f.name = f"{target}_fragment.html"
        f.seek(0) 
        
        caption = (f"📄 **Raw Response for** `@{target}`\n"
                   f"**Requested URL:** `{url}`\n"
                   f"**Final URL:** `{final_url}`\n"
                   f"**Status Code:** `{status_code}`")
        
        await message.reply_document(document=f, caption=caption)
        await msg.delete() 
    except Exception as e:
        try: await msg.edit_text(f"❌ **Error fetching HTML:** `{str(e)}`")
        except: pass

# --- STRING SESSION STATE MACHINE ---
input_waiters = {}

@app.on_message(filters.text & is_admin, group=-1)
async def catch_input(client, message):
    chat_id = message.chat.id
    if chat_id in input_waiters and not input_waiters[chat_id].done():
        input_waiters[chat_id].set_result(message.text)
        message.stop_propagation()

async def get_response(chat_id, timeout=300):
    loop = asyncio.get_running_loop()
    future = loop.create_future()
    input_waiters[chat_id] = future
    try: return await asyncio.wait_for(future, timeout)
    finally: del input_waiters[chat_id]

@app.on_message(filters.command("getstring") & is_admin)
async def pyrogram_string_cmd(client, message):
    chat_id = message.chat.id
    try:
        await message.reply_text("🛠 **Pyrogram Session Generator**\nEnter your `API_ID`:")
        api_id = int((await get_response(chat_id)).strip())

        await message.reply_text("Enter your `API_HASH`:")
        api_hash = (await get_response(chat_id)).strip()

        await message.reply_text("Enter the **Phone Number** (e.g., +1234567890):")
        phone = (await get_response(chat_id)).replace(" ", "").strip()

        await message.reply_text("⏳ Requesting OTP from Telegram...")
        
        from pyrogram import Client as TmpClient
        from pyrogram.errors import SessionPasswordNeeded

        tmp_app = TmpClient("temp_session", api_id=api_id, api_hash=api_hash, in_memory=True)
        await tmp_app.connect()
        
        sent_code = await tmp_app.send_code(phone)
        
        await message.reply_text("📥 **Code Sent!**\n⚠️ Enter code with spaces (e.g., `1 2 3 4 5`):")
        otp_code = (await get_response(chat_id)).replace(" ", "").strip()
        
        try:
            await tmp_app.sign_in(phone, sent_code.phone_code_hash, otp_code)
        except SessionPasswordNeeded:
            await message.reply_text("🔐 **2FA Password Required:**")
            pwd = (await get_response(chat_id)).strip()
            await tmp_app.check_password(pwd)
        
        string_session = await tmp_app.export_session_string()
        await message.reply_text(f"✅ **Session Generated!**\n\n`{string_session}`\n\n⚠️ Keep this secret.")
    except asyncio.TimeoutError:
        await message.reply_text("❌ **Timeout.** Run `/getstring` again.")
    except Exception as e:
        await message.reply_text(f"❌ **Error:** `{str(e)}`")
    finally:
        if 'tmp_app' in locals():
            try: await tmp_app.disconnect()
            except: pass

# --- UI CALLBACKS ---
@app.on_callback_query(is_admin)
async def button_handler(client, query):
    data = query.data
    try:
        if data == "back_main": 
            await query.edit_message_text("**Fragment Control System**", reply_markup=get_dynamic_menu())
            
        elif data == "menu_settings": 
            await query.edit_message_text("**System Settings**", reply_markup=get_settings_menu())
            
        elif data == "toggle_start":
            if not state.queue or state.queue.empty():
                return await query.answer("Queue is empty! Load words first.", show_alert=True)
            state.status = "RUNNING"
            log.info("Scanner Started")
            await query.edit_message_text("▶️ Scanner Running", reply_markup=get_dynamic_menu())
            
        elif data == "toggle_pause":
            state.status = "PAUSED"
            log.info("Scanner Paused")
            await query.edit_message_text("⏸ Scanner Paused", reply_markup=get_dynamic_menu())
            
        elif data == "toggle_resume":
            state.status = "RUNNING"
            log.info("Scanner Resumed")
            await query.edit_message_text("▶️ Scanner Resumed", reply_markup=get_dynamic_menu())
            
        elif data == "toggle_stop":
            state.status = "STOPPED"
            log.info("Scanner Stopped")
            await query.edit_message_text("⏹ Scanner Stopped", reply_markup=get_dynamic_menu())
            
        elif data == "set_min_length":
            kb = InlineKeyboardMarkup([
                [InlineKeyboardButton("4 Letters", callback_data="min_len_4"), InlineKeyboardButton("5 Letters", callback_data="min_len_5")],
                [InlineKeyboardButton("6 Letters", callback_data="min_len_6"), InlineKeyboardButton("7 Letters", callback_data="min_len_7")],
                [InlineKeyboardButton("« Back", callback_data="menu_settings")]
            ])
            await query.edit_message_text("📏 **Select Minimum Word Length:**", reply_markup=kb)
            
        elif data.startswith("min_len_"):
            new_len = int(data.split("_")[2])
            state.min_length = new_len
            await query.edit_message_text(f"✅ Filter updated to `{new_len}` letters.", reply_markup=get_settings_menu())
            
        elif data == "load_words": 
            await query.edit_message_text("⏳ Processing and loading chunk into memory...")
            
            def process_words_file():
                try:
                    with open("words.txt", "r", encoding="utf-8") as f: 
                        all_lines = [line.strip().lower() for line in f if line.strip()]
                    return all_lines
                except Exception:
                    return None
                    
            all_lines = await asyncio.to_thread(process_words_file)
            if all_lines is None:
                return await query.edit_message_text("❌ `words.txt` not found. Upload it first.")
                
            CHUNK = 200000
            current = all_lines[:CHUNK]
            remaining = all_lines[CHUNK:]
            
            # Non-destructive memory leak sorting fix
            pattern = re.compile(rf'^[a-z0-9_]{{{state.min_length},32}}$')
            valid = [w for w in current if pattern.match(w)]
            
            def rewrite_words_file():
                with open("words.txt", "w", encoding="utf-8") as f: 
                    f.write("\n".join(remaining))
                    
            await asyncio.to_thread(rewrite_words_file)
            await state.add_to_queue(valid)
            await query.edit_message_text(f"✅ Loaded {len(valid):,} filtered items into queue.", reply_markup=get_dynamic_menu())
            
        elif data in ("show_status", "refresh_status"):
            msg = generate_status_msg()
            kb = InlineKeyboardMarkup([[InlineKeyboardButton("🔄 Refresh", callback_data="refresh_status")], [InlineKeyboardButton("« Back", callback_data="back_main")]])
            try: await query.edit_message_text(msg, reply_markup=kb)
            except MessageNotModified: pass
            
        elif data == "set_speed":
            kb = InlineKeyboardMarkup([
                [InlineKeyboardButton("10", callback_data="spd_10"), InlineKeyboardButton("20", callback_data="spd_20")], 
                [InlineKeyboardButton("30", callback_data="spd_30"), InlineKeyboardButton("50", callback_data="spd_50")], 
                [InlineKeyboardButton("« Back", callback_data="menu_settings")]
            ])
            await query.edit_message_text("⚡ **Select Concurrent Workers:**", reply_markup=kb)
            
        elif data.startswith("spd_"):
            state.concurrency = int(data.split("_")[1])
            await query.edit_message_text("✅ Speed Updated", reply_markup=get_settings_menu())

        elif data == "export_files":
            await query.edit_message_text("⏳ Preparing files from disk...")
            exported = 0
            for s in ["taken", "unavailable", "sold", "forsale", "auction", "available", "error"]:
                if os.path.exists(f"results/{s}.txt"):
                    await client.send_document(query.message.chat.id, document=f"results/{s}.txt")
                    exported += 1
                    
            if exported > 0:
                await query.edit_message_text(f"✅ Exported {exported} result files.", reply_markup=get_dynamic_menu())
            else:
                await query.edit_message_text("⚠️ No data files found.", reply_markup=get_dynamic_menu())
                
        elif data == "reset_confirm":
            kb = InlineKeyboardMarkup([[InlineKeyboardButton("CONFIRM WIPE", callback_data="reset_do")], [InlineKeyboardButton("Cancel", callback_data="menu_settings")]])
            await query.edit_message_text("⚠️ WIPE ALL LOCAL SCANNER MEMORY AND PROGRESS OUT_FILES?", reply_markup=kb)
            
        elif data == "reset_do":
            state.processed = 0
            state.total_words = 0
            state.counts = {k: 0 for k in state.counts}
            state.queue = asyncio.Queue()
            state.status = "STOPPED"
            
            def clear_results():
                if os.path.exists("results"):
                    for file in os.listdir("results"):
                        try: os.remove(f"results/{file}")
                        except: pass
            await asyncio.to_thread(clear_results)
            await query.edit_message_text("💣 Memory and Results purged.", reply_markup=get_settings_menu())
            
    except Exception as e:
        if not isinstance(e, MessageNotModified):
            log.exception("UI Callback Error:")

if __name__ == "__main__":
    from worker import manage_workers
    
    async def main_loop():
        await app.start()
        log.info("🤖 Pyrogram Bot Started!")
        asyncio.create_task(manage_workers())
        from pyrogram import idle
        await idle()
        await app.stop()
        
    app.run(main_loop())