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
import traceback

# === CONFIG ===
API_ID = os.environ.get("API_ID")
API_HASH = os.environ.get("API_HASH")
BOT_TOKEN = os.environ.get("BOT_TOKEN")
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
OWNER_ID = int(os.environ.get("OWNER_ID", "7361622601"))

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"):
    loaded_count = 0
    for file in os.listdir(folder):
        if file.endswith(".py"):
            try:
                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)
                mod.bot = bot  # Inject bot instance
                spec.loader.exec_module(mod)
                loaded_count += 1
                print(f"✅ Loaded module: {name}")
            except Exception as e:
                print(f"❌ Failed to load {file}: {e}")
    print(f"📦 Loaded {loaded_count} 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", "generate", "add"]):
        return "CREATE"
    elif any(x in lowered for x in ["edit", "modify", "update", "change"]):
        return "EDIT"
    elif "recode" in lowered:
        return "RECODE"
    else:
        return "CREATE"  # Default to CREATE instead of UNSURE

def clean_code_blocks(code: str) -> str:
    # Remove markdown code blocks
    code = re.sub(r'^```python\n?', '', code, flags=re.MULTILINE)
    code = re.sub(r'^```\n?', '', code, flags=re.MULTILINE)
    code = re.sub(r'\n?```$', '', code, flags=re.MULTILINE)
    return code.strip()

def extract_module_name_from_description(description: str) -> str:
    """Enhanced module name extraction"""
    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', 'particular', 'user'
    }
    
    # Extract meaningful words
    desc = re.sub(r'[^a-zA-Z0-9_\s]', ' ', description.lower())
    words = [w for w in desc.split() if w not in stopwords and len(w) > 2]
    
    if words:
        name = '_'.join(words[:2])  # Take first 2 meaningful words
    else:
        # Fallback to extracting from original description
        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 and validate
    name = re.sub(r'_+', '_', name).strip('_').lower()
    return name if name else f"mod_{os.urandom(2).hex()}"

def get_improved_prompt(description: str, attempt: int = 1, last_error: str = None) -> str:
    """Generate improved prompts for better code generation"""
    
    base_prompt = f"""You are an expert Python developer creating a Pyrogram Telegram bot module.

TASK: Create a module that implements: {description}

CRITICAL REQUIREMENTS:
1. Use the existing 'bot' variable - DO NOT create a new Client instance
2. Import bot using this exact pattern:
   ```python
   try:
       from __main__ import bot
   except ImportError:
       from ..main import bot
   ```
3. Import only necessary Pyrogram components:
   ```python
   from pyrogram import filters
   from pyrogram.enums import ParseMode
   ```
4. Use proper async/await for all bot operations
5. Handle errors gracefully with try/except blocks
6. Include proper message validation and user feedback

HANDLER PATTERN:
```python
@bot.on_message(filters.command("commandname") & filters.private)
async def function_name(client, message):
    try:
        # Your code here
        await message.reply("Response")
    except Exception as e:
        await message.reply(f"❌ Error: {{str(e)}}")
```

SPECIFIC GUIDELINES:
- Only create commands explicitly mentioned in the request
- Use descriptive function names
- Include input validation
- Provide helpful error messages
- Use markdown formatting for responses when appropriate
- Handle edge cases (empty inputs, invalid data, etc.)

CODE STRUCTURE:
- Start with imports
- Define helper functions if needed
- Create command handlers
- No main execution code (no if __name__ == "__main__")

OUTPUT: Provide ONLY the Python code, no explanations or markdown formatting."""

    if attempt > 1 and last_error:
        base_prompt += f"\n\nPREVIOUS ERROR TO FIX: {last_error}\n"
        base_prompt += "Make sure to address this specific error in your implementation."
    
    return base_prompt

def get_edit_prompt(file_path: str, current_content: str, edit_instruction: str) -> str:
    """Generate prompt for editing existing modules"""
    return f"""You are an expert Python developer. Edit the following Pyrogram bot module.

CURRENT FILE: {file_path}
CURRENT CONTENT:
```python
{current_content}
```

EDIT INSTRUCTION: {edit_instruction}

REQUIREMENTS:
1. Maintain the existing import structure
2. Keep the bot variable usage consistent
3. Preserve working functionality while making requested changes
4. Follow the same coding patterns as the existing code
5. Add proper error handling for new features

OUTPUT: Provide ONLY the complete updated Python code, no explanations."""

async def restart_bot(chat_id):
    """Restart the bot to load new modules"""
    await bot.send_message(chat_id, "♻️ Restarting to apply changes...")
    await asyncio.sleep(1)
    await bot.stop()
    os.execl(sys.executable, sys.executable, *sys.argv)

def validate_generated_code(code: str) -> tuple[bool, str]:
    """Validate generated code before saving"""
    if not code.strip():
        return False, "Empty code generated"
    
    if "def " not in code:
        return False, "No function definitions found"
    
    if "@bot.on_message" not in code:
        return False, "No message handlers found"
    
    if "import" not in code:
        return False, "No imports found"
    
    # Check for common issues
    if "Client(" in code:
        return False, "Code creates new Client instance (should use existing bot)"
    
    if "__main__" in code and "if __name__" in code:
        return False, "Code contains main execution block (not allowed in modules)"
    
    return True, "Code validation passed"

# === MAIN HANDLER ===
@bot.on_message(filters.private & filters.text)
async def jarvis_trigger(client, message):
    # Owner-only access control
    if message.from_user.id != OWNER_ID:
        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 create?")

    intent = determine_intent(description)
    progress_msg = await message.reply(
        f"🤖 **Jarvis Assistant**\n"
        f"📝 Task: `{description}`\n"
        f"🧠 Intent: {intent}\n"
        f"⏳ Processing...", 
        parse_mode=ParseMode.MARKDOWN
    )

    # Handle EDIT intent
    if intent == "EDIT":
        edit_match = re.match(r"edit\s+([\w/\\.]+)\s+(.*)", description, re.IGNORECASE)
        if edit_match:
            file_path = edit_match.group(1)
            edit_instruction = edit_match.group(2)
            
            if not os.path.exists(file_path):
                return await progress_msg.edit(f"❌ File `{file_path}` not found.")
            
            try:
                with open(file_path, "r", encoding="utf-8") as f:
                    current_content = f.read()
                
                edit_prompt = get_edit_prompt(file_path, current_content, edit_instruction)
                
                await progress_msg.edit(f"🔄 Editing {file_path}...")
                response = await asyncio.to_thread(model.generate_content, edit_prompt)
                new_code = clean_code_blocks(response.text.strip())
                
                is_valid, validation_msg = validate_generated_code(new_code)
                if not is_valid:
                    return await progress_msg.edit(f"❌ Edit failed: {validation_msg}")
                
                with open(file_path, "w", encoding="utf-8") as f:
                    f.write(new_code)
                
                await progress_msg.edit(f"✅ Successfully edited `{file_path}`\n♻️ Restarting bot...")
                await asyncio.sleep(2)
                asyncio.create_task(restart_bot(message.chat.id))
                return
                
            except Exception as e:
                return await progress_msg.edit(f"❌ Edit failed: {str(e)}")
        else:
            return await progress_msg.edit("❗ Format: `jarvis edit <file> <instruction>`")

    # Handle CREATE intent
    success = False
    last_error = None
    
    for attempt in range(1, 4):  # Reduced to 3 attempts
        try:
            await progress_msg.edit(
                f"🤖 **Jarvis Assistant**\n"
                f"📝 Task: `{description}`\n"
                f"🧠 Intent: {intent}\n"
                f"🔄 Attempt {attempt}/3: Generating code...",
                parse_mode=ParseMode.MARKDOWN
            )
            
            # Generate code
            prompt = get_improved_prompt(description, attempt, last_error)
            response = await asyncio.to_thread(model.generate_content, prompt)
            code = clean_code_blocks(response.text.strip())
            
            # Validate generated code
            is_valid, validation_msg = validate_generated_code(code)
            if not is_valid:
                last_error = validation_msg
                continue
            
            # Auto-fix common issues
            code = code.replace('asyncio.PIPE', 'subprocess.PIPE')
            code = code.replace('asyncio.create_subprocess_shell', 'subprocess.run')
            
            # Determine module name and path
            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"
            
            # Save the module
            await progress_msg.edit(
                f"🤖 **Jarvis Assistant**\n"
                f"📝 Task: `{description}`\n"
                f"🧠 Intent: {intent}\n"
                f"💾 Saving module: {mod_name}...",
                parse_mode=ParseMode.MARKDOWN
            )
            
            with open(mod_path, "w", encoding="utf-8") as f:
                f.write(code)
            
            # Test the module
            await progress_msg.edit(
                f"🤖 **Jarvis Assistant**\n"
                f"📝 Task: `{description}`\n"
                f"🧠 Intent: {intent}\n"
                f"🧪 Testing module...",
                parse_mode=ParseMode.MARKDOWN
            )
            
            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:
                last_error = f"Module test failed: {str(test_error)}"
                continue
            
            # Success! Show details
            try:
                commands = extract_commands(code)
                commands_str = ', '.join(f'/{cmd}' for cmd in commands) if commands else 'No commands found'
                
                # Generate explanation
                explain_prompt = f"""Briefly explain what this Telegram bot module does based on the user's request and the code.

User request: {description}
Commands found: {commands_str}

Provide a concise 1-2 sentence explanation."""
                
                explain_response = await asyncio.to_thread(model.generate_content, explain_prompt)
                explanation = explain_response.text.strip()
                
                await progress_msg.edit(
                    f"✅ **Module Created Successfully!**\n\n"
                    f"📄 **File:** `{mod_path}`\n"
                    f"🔧 **Commands:** {commands_str}\n"
                    f"📋 **Description:** {explanation}\n\n"
                    f"♻️ Restarting bot to load the module...",
                    parse_mode=ParseMode.MARKDOWN
                )
                
            except Exception:
                await progress_msg.edit(
                    f"✅ **Module Created Successfully!**\n\n"
                    f"📄 **File:** `{mod_path}`\n"
                    f"🔧 **Commands:** {commands_str}\n\n"
                    f"♻️ Restarting bot to load the module...",
                    parse_mode=ParseMode.MARKDOWN
                )
            
            await asyncio.sleep(2)
            asyncio.create_task(restart_bot(message.chat.id))
            success = True
            break
            
        except Exception as e:
            last_error = str(e)
            print(f"❌ Attempt {attempt} failed: {e}")
            print(f"❌ Traceback: {traceback.format_exc()}")
    
    if not success:
        await progress_msg.edit(
            f"❌ **Failed to create module**\n\n"
            f"All attempts failed. Last error:\n`{last_error}`\n\n"
            f"Please try with a simpler request or check the logs.",
            parse_mode=ParseMode.MARKDOWN
        )

# === UTILITY COMMANDS ===
@bot.on_message(filters.command("modules") & filters.private)
async def list_modules(client, message):
    if message.from_user.id != OWNER_ID:
        return await message.reply("❌ Access denied.")

    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)

@bot.on_message(filters.command("delete") & filters.private)
async def delete_module(client, message):
    if message.from_user.id != OWNER_ID:
        return await message.reply("❌ Access denied.")

    if len(message.command) < 2:
        return await message.reply("❗ **Usage:** `/delete <module_name>`", 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.", parse_mode=ParseMode.MARKDOWN)
    else:
        await message.reply(f"❌ Module `{mod_name}.py` not found.", parse_mode=ParseMode.MARKDOWN)

@bot.on_message(filters.regex(r"(?i)(what|help|commands)") & filters.private)
async def help_command(_, message):
    if message.from_user.id != OWNER_ID:
        return await message.reply("❌ Access denied.")

    await message.reply(
        "🤖 **Jarvis Assistant - Help**\n\n"
        "**Natural Language:**\n"
        "• `jarvis create a calculator`\n"
        "• `jarvis make a weather bot`\n"
        "• `jarvis build a reminder system`\n\n"
        "**Commands:**\n"
        "• `/modules` - List all modules\n"
        "• `/delete <name>` - Delete a module\n\n"
        "**Examples:**\n"
        "• `jarvis create a /joke command`\n"
        "• `jarvis make a password generator`\n"
        "• `jarvis build a text to speech module`\n",
        parse_mode=ParseMode.MARKDOWN
    )

# === STARTUP ===
if __name__ == "__main__":
    print("🚀 Starting Jarvis Bot...")
    load_modules()
    print("✅ Bot is ready!")
    bot.run()