import gradio as gr
import subprocess
import asyncio
import threading
from queue import Queue, Empty
import os
import time
from telegram import Update
from telegram.ext import ApplicationBuilder, CommandHandler, ContextTypes

# --- Config ---
BOT_TOKEN = os.environ.get("BOT_TOKEN", "YOUR_BOT_TOKEN_HERE")

# --- Globals ---
bot_app = None
bot_running = False
log_queue = Queue()
telegram_cmd_logs = []

# --- Telegram Bot Command Handler ---
async def bash_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
    cmd = ' '.join(context.args)
    if not cmd:
        await update.message.reply_text("Usage: /bash <command>")
        return

    log_queue.put(f"[Telegram] Received command: {cmd}")
    try:
        proc = await asyncio.create_subprocess_shell(
            cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
        )
        stdout, _ = await proc.communicate()
        output = stdout.decode().strip()
        if len(output) > 4000:
            output = output[:4000] + "\n...[truncated]"
    except Exception as e:
        output = f"Error: {e}"

    await update.message.reply_text(f"$ {cmd}\n{output}")
    telegram_cmd_logs.append(f"$ {cmd}\n{output}")
    if len(telegram_cmd_logs) > 50:
        telegram_cmd_logs.pop(0)

# --- Terminal Command Function ---
def live_terminal(cmd):
    if not cmd.strip():
        yield "$ "
        return
    yield f"$ {cmd}\n"
    process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    for line in iter(process.stdout.readline, ''):
        yield line
    process.stdout.close()
    process.wait()
    yield "$ "

# --- Bot Start/Stop Functions ---
def start_bot_sync():
    global bot_app, bot_running
    if bot_running:
        log_queue.put("[Bot] Already running.")
        return

    async def main():
        global bot_app, bot_running
        bot_app = ApplicationBuilder().token(BOT_TOKEN).build()
        bot_app.add_handler(CommandHandler("bash", bash_command))
        bot_running = True
        log_queue.put("[Bot] Started.")
        await bot_app.initialize()
        await bot_app.start()
        await bot_app.updater.start_polling()

    asyncio.run(main())

def stop_bot_sync():
    global bot_app, bot_running
    if not bot_running:
        log_queue.put("[Bot] Not running.")
        return

    async def shutdown():
        await bot_app.updater.stop()
        await bot_app.stop()
        await bot_app.shutdown()
        log_queue.put("[Bot] Stopped.")

    asyncio.run(shutdown())
    bot_running = False

# --- Log Collector ---
def update_log_box():
    logs = []
    try:
        while True:
            logs.append(log_queue.get_nowait())
    except Empty:
        pass
    logs.extend(telegram_cmd_logs[-20:])
    return "\n".join(logs[-20:])

# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.Markdown("## 🖥️ Live Terminal + Telegram Bot Interface")

    with gr.Row():
        terminal_output = gr.Textbox(label="Live Terminal", lines=20, interactive=False)
        cmd_input = gr.Textbox(placeholder="Type shell command", label="Command Input")

    with gr.Row():
        run_btn = gr.Button("▶️ Run")
        start_btn = gr.Button("🤖 Start Telegram Bot")
        stop_btn = gr.Button("🛑 Stop Telegram Bot")

    log_output = gr.Textbox(label="📜 Bot & Telegram Logs", lines=15, interactive=False)

    # Streaming terminal output
    run_btn.click(fn=live_terminal, inputs=cmd_input, outputs=terminal_output)

    # Start/Stop bot
    start_btn.click(lambda: update_log_box() or start_bot_sync(), outputs=log_output)
    stop_btn.click(lambda: update_log_box() or stop_bot_sync(), outputs=log_output)

    # Auto log updater every 3 seconds (Gradio v5 style)
    def background_log_updater():
        while True:
            time.sleep(3)
            yield update_log_box()

    demo.load(background_log_updater, outputs=log_output)

demo.launch()