import os
import sys
import traceback
import io
import asyncio
from pyrogram import Client, filters
from config import ADMIN

@Client.on_message(filters.user(ADMIN) & filters.command("sh"))
async def shell_runner(client, message):
    if len(message.command) < 2:
        return await message.reply_text("<b>Usage:</b> <code>/sh [command]</code>")
    
    code = message.text.split(None, 1)[1]
    msg = await message.reply_text("<code>Executing...</code>")
    
    process = await asyncio.create_subprocess_shell(
        code,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    stdout, stderr = await process.communicate()
    
    result = str(stdout.decode().strip()) + str(stderr.decode().strip())
    
    if len(result) > 4096:
        with io.BytesIO(str.encode(result)) as out_file:
            out_file.name = "shell.txt"
            await message.reply_document(document=out_file, caption=f"<code>{code}</code>")
            await msg.delete()
    else:
        await msg.edit(f"<b>Command:</b>\n<code>{code}</code>\n\n<b>Result:</b>\n<code>{result}</code>")

@Client.on_message(filters.user(ADMIN) & filters.command("eval"))
async def eval_runner(client, message):
    if len(message.command) < 2:
        return await message.reply_text("<b>Usage:</b> <code>/eval [code]</code>")
        
    status_message = await message.reply_text("<code>Processing...</code>")
    cmd = message.text.split(None, 1)[1]

    old_stderr = sys.stderr
    old_stdout = sys.stdout
    redirected_output = io.StringIO()
    redirected_error = io.StringIO()
    sys.stdout = redirected_output
    sys.stderr = redirected_error
    stdout, stderr, exc = None, None, None

    try:
        await aexec(cmd, client, message)
    except Exception:
        exc = traceback.format_exc()

    stdout = redirected_output.getvalue()
    stderr = redirected_error.getvalue()
    sys.stdout = old_stdout
    sys.stderr = old_stderr

    evaluation = ""
    if exc:
        evaluation = exc
    elif stderr:
        evaluation = stderr
    elif stdout:
        evaluation = stdout
    else:
        evaluation = "Success"

    final_output = f"<b>EVAL</b>: <code>{cmd}</code>\n\n<b>OUTPUT</b>:\n<code>{evaluation}</code>"

    if len(final_output) > 4096:
        with io.BytesIO(str.encode(evaluation)) as out_file:
            out_file.name = "eval.txt"
            await message.reply_document(
                document=out_file,
                caption=f"<code>{cmd}</code>",
                quote=True
            )
            await status_message.delete()
    else:
        await status_message.edit(final_output)

async def aexec(code, client, message):
    exec(
        f"async def __aexec(client, message): "
        + "".join(f"\n {l}" for l in code.split("\n"))
    )
    return await locals()["__aexec"](client, message)