import os
import html
from pyrogram import Client, filters, enums
from pyrogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton, CallbackQuery
from pyromod.exceptions import ListenerTimeout
from helper.database import codeflixbots as db

EDIT_STATE = {}

@Client.on_message(filters.command("edit") & filters.private)
async def edit_cmd(client: Client, message: Message):
    if not await db.is_user_authorized(message.from_user.id):
        return
        
    await client.stop_listening(chat_id=message.chat.id, user_id=message.from_user.id)
    
    if message.reply_to_message and message.reply_to_message.forward_from_chat:
        fwd_msg = message.reply_to_message
    else:
        try:
            fwd_msg = await client.ask(
                chat_id=message.chat.id,
                text="**FORWARD a message from your channel here to edit its buttons:**\n\n(Or send /cancel to abort)",
                filters=filters.forwarded | filters.text,
                user_id=message.from_user.id,
                timeout=300
            )
        except ListenerTimeout:
            return
        except Exception:
            return
            
        if not fwd_msg.forward_from_chat and getattr(fwd_msg, "text", "") == "/cancel":
            await message.reply_text("❌ Cancelled edit operation.")
            return

    if not fwd_msg.forward_from_chat or fwd_msg.forward_from_chat.type != enums.ChatType.CHANNEL:
        await message.reply_text("❌ Please **forward a valid message from a Channel**.")
        return
        
    chat_id = fwd_msg.forward_from_chat.id
    msg_id = fwd_msg.forward_from_message_id
    
    try:
        orig_msg = await client.get_messages(chat_id, msg_id)
        if not orig_msg:
            await message.reply_text("❌ Could not fetch original message from channel.")
            return
    except Exception as e:
        await message.reply_text(f"❌ Error fetching message: {e}\n(Make sure Bot is Admin in the channel)")
        return
        
    # Extract existing inline buttons
    buttons = []
    if orig_msg.reply_markup and orig_msg.reply_markup.inline_keyboard:
        for row in orig_msg.reply_markup.inline_keyboard:
            for btn in row:
                if btn.url:
                    buttons.append({"text": btn.text, "url": btn.url})
                    
    EDIT_STATE[message.from_user.id] = {
        "chat_id": chat_id,
        "msg_id": msg_id,
        "buttons": buttons,
        "layout": []
    }
    
    await render_edit_menu(client, message.from_user.id, message.chat.id)


async def render_edit_menu(client, user_id, chat_id, m_edit=None):
    state = EDIT_STATE.get(user_id)
    if not state:
        return
        
    text = f"🛠 **Message Editor Mode**\n\n"
    text += f"**Channel ID:** `{state['chat_id']}`\n"
    text += f"**Message ID:** `{state['msg_id']}`\n\n"
    text += "**Current Buttons:**\n"
    
    if not state["buttons"]:
        text += "_No buttons attached to this message._\n"
    else:
        for i, btn in enumerate(state["buttons"]):
            text += f"{i+1}. {btn['text']} — `{btn['url']}`\n"
            
    markup = []
    markup.append([InlineKeyboardButton("➕ Add New Button", callback_data="edbtn_add")])
    
    if state["buttons"]:
        markup.append([
            InlineKeyboardButton("➖ Remove Button", callback_data="edbtn_rem"),
            InlineKeyboardButton("📟 Change Layout", callback_data="edbtn_layout")
        ])
        
    markup.append([InlineKeyboardButton("💾 Save & Apply To Channel", callback_data="edbtn_save")])
    markup.append([InlineKeyboardButton("🗑 Cancel", callback_data="edbtn_cancel")])
    
    if m_edit:
        await m_edit.edit_text(text, reply_markup=InlineKeyboardMarkup(markup))
    else:
        await client.send_message(chat_id, text, reply_markup=InlineKeyboardMarkup(markup))


@Client.on_callback_query(filters.regex(r"^edbtn_(.*)"))
async def edit_callback_handler(client: Client, query: CallbackQuery):
    action = query.matches[0].group(1)
    user_id = query.from_user.id
    state = EDIT_STATE.get(user_id)
    
    if not state:
        await query.answer("Session expired. Send /edit again.", show_alert=True)
        return
        
    if action == "cancel":
        del EDIT_STATE[user_id]
        await query.message.edit_text("❌ Edit Session Cancelled.")
        return
        
    elif action == "add":
        await query.answer()
        try:
            btn_info = await client.ask(
                chat_id=query.message.chat.id,
                text=(
                    "**Send button details:**\n\n"
                    "Use one `Button Name - https://yourlink.com` per line. You can send several lines at once:\n\n"
                    "`WATCH HERE - https://t.me/...`\n"
                    "`SHARE CHANNEL - https://t.me/...`\n\n"
                    "(Send /cancel to abort)"
                ),
                user_id=user_id,
                timeout=120
            )
        except ListenerTimeout:
            return
        except Exception:
            return
            
        if (btn_info.text or "").strip() == "/cancel":
            await render_edit_menu(client, user_id, query.message.chat.id)
            return

        valid_buttons = []
        invalid_lines = []
        for line in (btn_info.text or "").splitlines():
            line = line.strip()
            if not line:
                continue
            # Keep dashes inside URLs intact; the separator is " - ".
            if " - " not in line:
                invalid_lines.append(line)
                continue
            button_text, button_url = (part.strip() for part in line.split(" - ", 1))
            if button_text and button_url.startswith(("http://", "https://")):
                valid_buttons.append((button_text, button_url))
            else:
                invalid_lines.append(line)

        # Telegram currently limits an inline keyboard to 100 buttons.
        remaining = max(0, 100 - len(state["buttons"]))
        if len(valid_buttons) > remaining:
            invalid_lines.extend(f"{label} - {url}" for label, url in valid_buttons[remaining:])
            valid_buttons = valid_buttons[:remaining]

        for button_text, button_url in valid_buttons:
            state["buttons"].append({"text": button_text, "url": button_url})

        if valid_buttons:
            reply = f"✅ Added {len(valid_buttons)} button(s)."
            if invalid_lines:
                reply += f"\n⚠️ Skipped {len(invalid_lines)} invalid line(s). Use `Button Name - https://link` for each button."
            await btn_info.reply_text(reply)
        else:
            await btn_info.reply_text("⚠️ No valid buttons found. Use one `Button Name - https://link` per line.")
                
        # Reset layout when a new button is added since count changes
        state["layout"] = []
        await render_edit_menu(client, user_id, query.message.chat.id)
        
    elif action == "layout":
        await query.answer()
        num_btns = len(state["buttons"])
        if num_btns == 0:
            return
            
        try:
            lay_info = await client.ask(
                chat_id=query.message.chat.id,
                text=f"**Send Button Layout Structure:**\n\nYou have **{num_btns}** buttons.\nSend spaces between numbers to set rows.\nExample: `2 1 1` (means 2 in first row, 1 in next, etc.)\n\n(Send /cancel to abort)",
                user_id=user_id,
                timeout=120
            )
        except Exception:
            return
            
        if lay_info.text == "/cancel":
            await render_edit_menu(client, user_id, query.message.chat.id)
            return
            
        try:
            parts = lay_info.text.strip().split()
            layout = [int(x) for x in parts]
            if sum(layout) != num_btns:
                await lay_info.reply_text(f"⚠️ **Error:** Sum of your layout ({sum(layout)}) does not match your total buttons ({num_btns}).")
            else:
                state["layout"] = layout
                await lay_info.reply_text("✅ **Custom Layout Applied!**")
        except Exception:
            await lay_info.reply_text("⚠️ **Invalid format.** Just send numbers separated by spaces.")
            
        await render_edit_menu(client, user_id, query.message.chat.id)
        
    elif action == "rem":
        if not state["buttons"]:
            await query.answer("No buttons to remove.", show_alert=True)
            return
            
        await query.answer()
        rem_buttons = []
        for i, b in enumerate(state["buttons"]):
            rem_buttons.append([InlineKeyboardButton(f"Remove: {b['text']}", callback_data=f"edrem_{i}")])
        rem_buttons.append([InlineKeyboardButton("🔙 Back to Edit Menu", callback_data="edbtn_back")])
        
        await query.message.edit_text("✂️ **Select a button to remove:**", reply_markup=InlineKeyboardMarkup(rem_buttons))
        
    elif action == "back":
        await render_edit_menu(client, user_id, query.message.chat.id, m_edit=query.message)
        
    elif action == "save":
        await query.message.edit_text("⏳ Saving changes to channel...")
        chat_id = state["chat_id"]
        msg_id = state["msg_id"]
        
        final_markup = None
        if state["buttons"]:
            kb = []
            layout = state.get("layout", [])
            btns = state["buttons"].copy()
            
            if layout and sum(layout) == len(btns):
                idx = 0
                for row_len in layout:
                    row = []
                    for _ in range(row_len):
                        b = btns[idx]
                        row.append(InlineKeyboardButton(b["text"], url=b["url"]))
                        idx += 1
                    kb.append(row)
            else:
                # Default 1 per row
                for b in btns:
                    kb.append([InlineKeyboardButton(b["text"], url=b["url"])])
            final_markup = InlineKeyboardMarkup(kb)
            
        try:
            await client.edit_message_reply_markup(chat_id, msg_id, reply_markup=final_markup)
            await query.message.edit_text("✅ **Successfully updated the message buttons in the channel!**")
            del EDIT_STATE[user_id]
        except Exception as e:
            await query.message.edit_text(f"❌ Failed to update channel message: {e}")


@Client.on_callback_query(filters.regex(r"^edrem_(\d+)"))
async def edit_remove_handler(client: Client, query: CallbackQuery):
    user_id = query.from_user.id
    state = EDIT_STATE.get(user_id)
    if not state:
        await query.answer("Session expired.", show_alert=True)
        return
        
    idx = int(query.matches[0].group(1))
    if 0 <= idx < len(state["buttons"]):
        removed = state["buttons"].pop(idx)
        await query.answer(f"Removed: {removed['text']}")
        
    await render_edit_menu(client, user_id, query.message.chat.id, m_edit=query.message)