import os
import re
import asyncio
import sys
import importlib.util
from pyrogram import Client, filters
from pyrogram.enums import ParseMode
import google.generativeai as genai
import subprocess

# check line 93 .. accordingly
# === CONFIG ===
API_ID = os.environ.get("API_ID")                 # Replace with your API ID
API_HASH = os.environ.get("API_HASH")   # Replace with your API HASH
BOT_TOKEN = os.environ.get("BOT_TOKEN")
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
if not BOT_TOKEN:
    raise RuntimeError("BOT_TOKEN environment variable not set!")
if not GEMINI_API_KEY:
    raise RuntimeError("GEMINI_API_KEY environment variable not set!")

# === SETUP ===
genai.configure(api_key=GEMINI_API_KEY)
model = genai.GenerativeModel("gemini-2.5-flash")
bot = Client("JarvisBot", api_id=API_ID, api_hash=API_HASH, bot_token=BOT_TOKEN)
os.makedirs("modules", exist_ok=True)

# === MODULE LOADER ===
def load_modules(folder="modules"):
    for file in os.listdir(folder):
        if file.endswith(".py"):
            path = os.path.join(folder, file)
            name = file[:-3]
            spec = importlib.util.spec_from_file_location(name, path)
            mod = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(mod)

load_modules()

# === HELPERS ===
def extract_module_name(code: str) -> str:
    match = re.search(r"def\\s+([a-zA-Z_][a-zA-Z0-9_]*)", code)
    return match.group(1).lower() if match else f"mod_{os.urandom(2).hex()}"

def extract_commands(code: str) -> list:
    return re.findall(r'filters\.command\(["\'](\w+)["\']', code)

def determine_intent(text: str) -> str:
    lowered = text.lower()
    if any(x in lowered for x in ["create", "make", "build"]):
        return "CREATE"
    elif any(x in lowered for x in ["edit", "modify"]):
        return "EDIT"
    elif "recode" in lowered:
        return "RECODE"
    else:
        return "UNSURE"

def clean_code_blocks(code: str) -> str:
    if code.startswith("```python"):
        code = code[9:]
    if code.startswith("```"):
        code = code[3:]
    if code.endswith("```"):
        code = code[:-3]
    return code.strip()

def extract_module_name_from_description(description: str) -> str:
    # Improved module name extraction - focus on key nouns and actions
    import re
    
    # Common words to remove
    stopwords = {
        'jarvis', 'make', 'create', 'build', 'generate', 'a', 'an', 'the', 'please', 'module', 
        'for', 'me', 'to', 'with', 'add', 'new', 'of', 'system', 'bot', 'command', 'that', 'and', 
        'in', 'on', 'by', 'as', 'is', 'it', 'my', 'this', 'do', 'can', 'you', 'i', 'want', 'need', 
        'from', 'like', 'using', 'feature', 'function', 'implement', 'write', 'code', 'file', 
        'python', 'telegram', 'pyrogram', 'handler', 'example', 'mod', 'mod_', 'particular', 'user'
    }
    
    # Clean the description
    desc = re.sub(r'[^a-zA-Z0-9_\s]', ' ', description.lower())
    words = desc.split()
    
    # Filter out stopwords and short words
    filtered_words = []
    for word in words:
        if word not in stopwords and len(word) > 2:
            filtered_words.append(word)
    
    # If we have good words, use them
    if filtered_words:
        # Take up to 3 most relevant words
        name_words = filtered_words[:3]
        name = '_'.join(name_words)
    else:
        # Fallback: extract any meaningful word
        meaningful_words = [w for w in words if len(w) > 3 and w not in stopwords]
        if meaningful_words:
            name = meaningful_words[0]
        else:
            # Last resort: use first few characters
            clean_desc = re.sub(r'[^a-zA-Z0-9]', '', description.lower())
            name = clean_desc[:10] if clean_desc else f"mod_{os.urandom(2).hex()}"
    
    # Clean up the name
    name = re.sub(r'_+', '_', name).strip('_')
    name = name.lower()
    
    # Ensure it's not empty
    if not name:
        name = f"mod_{os.urandom(2).hex()}"
    
    return name

async def restart_bot(chat_id):
    await bot.send_message(chat_id, "♻️ Restarting to apply new module...")
    await bot.stop()
    os.execl(sys.executable, sys.executable, *sys.argv)
    os._exit(0)  # Ensure process exits if execl fails

# === WORD TRIGGER ===
@bot.on_message(filters.private & filters.text)
async def jarvis_trigger(client, message):
    # Owner-only access control
    if message.from_user.id != 7361622601:
        return await message.reply("❌ Access denied. Only the owner can use this bot.")

    text = message.text.strip()
    if not text.lower().startswith("jarvis"):
        return

    description = text[6:].strip()
    if not description:
        return await message.reply("🤖 Hello, what would you like me to do?")

    intent = determine_intent(description)
    progress_msg = await message.reply(f"🤖 Acknowledged.\n🧠 Determining intent...\n📘 Intent: {intent}\n📝 Task: `{description}`", parse_mode=ParseMode.MARKDOWN)

    # --- NEW: Handle EDIT intent for existing files ---
    if intent == "EDIT":
        # Try to extract filename and edit instruction
        match = re.match(r"edit\s+([\w/\\.]+)\s+(.*)", description, re.IGNORECASE)
        if match:
            file_path = match.group(1)
            edit_instruction = match.group(2)
            if not os.path.exists(file_path):
                await progress_msg.edit(f"❌ File `{file_path}` not found.")
                return
            with open(file_path, "r", encoding="utf-8") as f:
                current_content = f.read()
            edit_prompt = (
                f"You are an expert Python developer. Here is the current content of `{file_path}`:\n"
                f"""\n{current_content}\n"""
                f"Please edit this file to: {edit_instruction}.\n"
                f"Output ONLY the new Python code, no explanations or markdown."
            )
            try:
                response = model.generate_content(edit_prompt)
                new_code = clean_code_blocks(response.text.strip())
                if not new_code or "def " not in new_code:
                    await progress_msg.edit(f"❌ Edit failed: No valid code returned.")
                    return
                with open(file_path, "w", encoding="utf-8") as f:
                    f.write(new_code)
                # Try to reload the module if it's in modules/
                if file_path.startswith("modules/") and file_path.endswith(".py"):
                    mod_name = os.path.basename(file_path)[:-3]
                    try:
                        spec = importlib.util.spec_from_file_location(mod_name, file_path)
                        mod = importlib.util.module_from_spec(spec)
                        mod.bot = bot
                        spec.loader.exec_module(mod)
                    except Exception as test_error:
                        await progress_msg.edit(f"⚠️ Edit applied, but module reload failed: {test_error}")
                        return
                await progress_msg.edit(f"✅ Edit applied to `{file_path}`. Restarting bot...")
                await asyncio.sleep(2)
                asyncio.create_task(restart_bot(message.chat.id))
                return
            except Exception as e:
                await progress_msg.edit(f"❌ Edit failed: `{str(e)[:100]}...`")
                return
        else:
            await progress_msg.edit("❗ Could not parse edit command. Use: 'edit <file> <instruction>'")
            return
    # --- END NEW ---

    success = False
    code = ""
    last_error = None
    for attempt in range(1, 6):
        await progress_msg.edit(f"`Attempt {attempt}/5:` Thinking and generating code...\n- Executing prerequisite shell commands...")

        try:
            prompt = (
                f"Write a full Pyrogram Telegram bot module that implements:\n"
                f"{description}.\n\n"
                f"IMPORTANT RULES:\n"
                f"1. Use 'bot' variable (not 'Client') - it will be injected\n"
                f"2. Include commands using @bot.on_message(filters.command(...))\n"
                f"3. Import only what you need from pyrogram\n and use 'from ..t1 import bot'\n 'try: from ..t1 import bot except (ImportError, ValueError): \n from __main__ import bot' \n"
                f"4. Don't create a new Client instance\n"
                f"5. Make functions async when using bot methods\n\n"
                f"Output ONLY Python code, no explanations or markdown."
            )
            if attempt > 1 and last_error is not None:
                prompt += f"\n\nNote: Previous attempt failed with error: {str(last_error)}. Fix that issue this time."

            response = model.generate_content(prompt)
            code = clean_code_blocks(response.text.strip())

            # === AUTO-FIX COMMON LLM MISTAKES ===
            code = code.replace('asyncio.PIPE', 'subprocess.PIPE')
            # Add more auto-fixes as needed

            # === LINTER/STATIC CHECKER ===
            with open('tmp_lint_module.py', 'w', encoding='utf-8') as tmpf:
                tmpf.write(code)
            lint_proc = subprocess.run(['pyflakes', 'tmp_lint_module.py'], capture_output=True, text=True)
            if lint_proc.returncode != 0 or lint_proc.stdout.strip() or lint_proc.stderr.strip():
                lint_output = (lint_proc.stdout + '\n' + lint_proc.stderr).strip()
                await progress_msg.edit(f"❌ Linter/static check failed!\n\n```\n{lint_output}\n```")
                os.remove('tmp_lint_module.py')
                continue
            os.remove('tmp_lint_module.py')

            if "def " not in code or "@bot" not in code:
                await progress_msg.edit(f"❌ Attempt {attempt}: Invalid function or handler.")
                continue

            # --- NEW: Use file path from description if present ---
            import re
            file_path_match = re.search(r"(modules/[\w\-]+\.py)", description)
            if file_path_match:
                mod_path = file_path_match.group(1)
                mod_name = os.path.basename(mod_path)[:-3]
            else:
                mod_name = extract_module_name_from_description(description)
                mod_path = f"modules/{mod_name}.py"
            # If file exists, use edit flow instead of creating a new file
            if os.path.exists(mod_path):
                with open(mod_path, "r", encoding="utf-8") as f:
                    current_content = f.read()
                edit_prompt = (
                    f"You are an expert Python developer. Here is the current content of `{mod_path}`:\n"
                    f"""\n{current_content}\n"""
                    f"Please update this file to: {description}.\n"
                    f"Output ONLY the new Python code, no explanations or markdown."
                )
                response = model.generate_content(edit_prompt)
                code = clean_code_blocks(response.text.strip())
            # --- END NEW ---

            with open(mod_path, "w", encoding="utf-8") as f:
                f.write(code)

            await progress_msg.edit(f"`Attempt {attempt}/5: Thinking and generating code...\n- Executing prerequisite shell commands...\n- Testing generated code...")

            try:
                spec = importlib.util.spec_from_file_location(mod_name, mod_path)
                mod = importlib.util.module_from_spec(spec)
                mod.bot = bot
                spec.loader.exec_module(mod)
            except Exception as test_error:
                raise RuntimeError(f"Testing error: {test_error}")

            await progress_msg.edit(f"`Attempt {attempt}/5:` Thinking and generating code...\n- Executing prerequisite shell commands...\n- Testing generated code...\n- Tests passed. Writing files...\n- Files written. Initiating restart...")

            # === NEW: Show module details before restart ===
            try:
                commands = extract_commands(code)
                commands_str = ', '.join(f'/{cmd}' for cmd in commands) if commands else 'No commands found.'
            except Exception as cmd_error:
                commands_str = f'Error extracting commands: {cmd_error}'
            
            try:
                explain_prompt = (
                    f"Explain in 2-3 sentences what this Pyrogram Telegram bot module does, given the following code and the user's request.\n"
                    f"User request: {description}\n"
                    f"Module code:\n{code}\n"
                    f"Be concise and clear."
                )
                explain_response = model.generate_content(explain_prompt)
                explanation = explain_response.text.strip()
            except Exception as exp_explain:
                explanation = f"Could not generate explanation: {exp_explain}"
            
            try:
                details_msg = (
                    f"✅ Module created successfully!\n\n"
                    f"**Commands:** {commands_str}\n\n"
                    f"**What it does:**\n{explanation}\n\n"
                    f"♻️ Restarting bot to load the new module..."
                )
                await progress_msg.edit(details_msg, parse_mode=ParseMode.MARKDOWN)
            except Exception as msg_error:
                # Fallback message if formatting fails
                await progress_msg.edit(f"✅ Module created successfully! Commands: {commands_str}\n\nWhat it does: {explanation}\n\n♻️ Restarting bot...")
            # === END NEW ===
            await asyncio.sleep(2)
            asyncio.create_task(restart_bot(message.chat.id))
            success = True
            break

        except Exception as e:
            last_error = e
            error_msg = f"❌ Attempt {attempt} Error: {str(e)}"
            print(f"DEBUG: Module generation error on attempt {attempt}: {e}")
            print(f"DEBUG: Error type: {type(e).__name__}")
            try:
                await progress_msg.edit(error_msg[:100] + "..." if len(error_msg) > 100 else error_msg)
            except Exception as edit_error:
                print(f"DEBUG: Failed to edit progress message: {edit_error}")
                # Try to send a new message if editing fails
                try:
                    await message.reply(f"❌ Attempt {attempt} failed. Check logs for details.")
                except:
                    pass

    if not success:
        await progress_msg.edit("❌ All 5 attempts failed. Please try again with a simpler instruction.")

# === /modules ===
@bot.on_message(filters.command("modules") & filters.private)
async def list_modules(client, message):
    # Owner-only access control
    if message.from_user.id != 7361622601:
        return await message.reply("❌ Access denied. Only the owner can use this bot.")

    modules = [f for f in os.listdir("modules") if f.endswith(".py")]
    if modules:
        module_list = "\n".join([f"• {m[:-3]}" for m in modules])
        await message.reply(f" Loaded Modules: \n{module_list}", parse_mode=ParseMode.MARKDOWN)
    else:
        await message.reply(" No modules found.", parse_mode=ParseMode.MARKDOWN)

# === /delete ===
@bot.on_message(filters.command("delete") & filters.private)
async def delete_module(client, message):
    # Owner-only access control
    if message.from_user.id != 7361622601:
        return await message.reply("❌ Access denied. Only the owner can use this bot.")

    if len(message.command) < 2:
        return await message.reply("❗ Specify module name.\n\nExample: `/delete calculator`", parse_mode=ParseMode.MARKDOWN)
    
    mod_name = message.command[1].lower()
    mod_path = f"modules/{mod_name}.py"
    
    if os.path.exists(mod_path):
        os.remove(mod_path)
        await message.reply(f"🗑 Deleted `{mod_name}.py`\n♻️ Restart required to take effect.", parse_mode=ParseMode.MARKDOWN)
    else:
        await message.reply(f"❌ Module `{mod_name}.py` not found.", parse_mode=ParseMode.MARKDOWN)

# === /what to do ===
@bot.on_message(filters.regex(r"(?i)what( can)? i do\??") & filters.private)
async def what_to_do(_, message):
    # Owner-only access control
    if message.from_user.id != 7361622601:
        return await message.reply("❌ Access denied. Only the owner can use this bot.")

    await message.reply(
        "**🧠 Jarvis Assistant**\n\n"
        "I can generate custom bot modules for you!\n\n"
        "**Commands:**\n"
        "`jarvis make a calculator` - Natural language trigger\n"
        "`/modules` - List all modules\n"
        "`/delete <name>` - Delete a module\n\n"
        "**Examples:**\n"
        "`jarvis build a reminder system`\n"
        "`jarvis create a dice game`\n"
        "`jarvis weather checker`\n",
        parse_mode=ParseMode.MARKDOWN
    )

# === START ===
print("🚀 Starting Jarvis Bot...")
load_modules()
print("✅ Bot is ready!")
bot.run()