core/bot.py

import os
import asyncio
import logging
from telegram import Update
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes

# Using safe relative package imports
from .rate_limiter import RateLimitManager
from .planner import TaskPlanner
from .executor import ExecutionEngine
from models.router import ModelRouter

log = logging.getLogger("bot")

class AICoderBot:
    def __init__(self):
        self.token = os.getenv("TELEGRAM_BOT_TOKEN")
        self.owner_id = int(os.getenv("TELEGRAM_OWNER_ID", "0"))
        self.app = Application.builder().token(self.token).build()
        
        # Initialize modules
        self.rl = RateLimitManager()
        self.router = ModelRouter(self.rl)
        self.planner = TaskPlanner(self.router)
        
        # Async hooks for conversation state locking
        self._user_reply_event = asyncio.Event()
        self._last_user_message = ""

    async def start(self):
        log.info("Registering bot command handlers...")
        self.app.add_handler(CommandHandler("task", self.handle_task))
        self.app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, self.handle_incoming_text))
        
        await self.app.initialize()
        await self.app.start()
        await self.app.updater.start_polling()
        log.info("Telegram engine connected and polling updates.")
        
        # Keeping loop alive
        while True:
            await asyncio.sleep(3600)

    async def handle_incoming_text(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        if update.effective_user.id != self.owner_id:
            return
        self._last_user_message = update.message.text
        self._user_reply_event.set()

    async def ask_owner_hook(self, question: str) -> str:
        """Suspends code execution dynamically until you reply on Telegram."""
        msg = f"⚠️ *QUESTION FROM AGENT*:\n{question}\n\n_Type your reply directly to respond._"
        await self.app.bot.send_message(chat_id=self.owner_id, text=msg, parse_mode="Markdown")
        
        self._user_reply_event.clear()
        await self._user_reply_event.wait()
        return self._last_user_message

    async def send_progress(self, message: str):
        """Sends instant execution updates to your Telegram chat."""
        try:
            await self.app.bot.send_message(chat_id=self.owner_id, text=f"ℹ️ {message}")
        except Exception as e:
            log.error(f"Failed sending update to Telegram: {e}")

    async def handle_task(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        if update.effective_user.id != self.owner_id:
            await update.message.reply_text("Access denied.")
            return

        task_description = " ".join(context.args)
        if not task_description:
            await update.message.reply_text("Please specify a task! Example: `/task Build a Flask webapp`", parse_mode="Markdown")
            return

        await update.message.reply_text(f"🚀 *Task Received:* \"{task_description}\"\nStructuring implementation plan...")

        try:
            project_workspace = "./workspace/project"
            os.makedirs(project_workspace, exist_ok=True)
            plan = await self.planner.create_plan(task_description, project_workspace)
        except Exception as e:
            await update.message.reply_text(f"❌ Failed to parse task steps: {e}")
            return

        # Start execution in a safe background task
        engine = ExecutionEngine(
            model_router=self.router,
            progress_cb=self.send_progress,
            ask_owner_cb=self.ask_owner_hook
        )
        asyncio.create_task(self._run_engine_safely(engine, plan))

    async def _run_engine_safely(self, engine: ExecutionEngine, plan):
        await self.send_progress(f"Starting execution of *{len(plan.steps)} planned steps*.")
        success = await engine.execute_plan(plan)
        if success:
            await self.send_progress("🎉 *Task Complete!* All files written and local compile checks passed.")
        else:
            await self.send_progress("🛑 *Task Terminated*: Execution broke down or was aborted.")