ajarbot/bot_runner.py

"""
Multi-platform bot runner for ajarbot.

Usage:
    python bot_runner.py                       # Run with config from adapters.yaml
    python bot_runner.py --config custom.yaml  # Use custom config file
    python bot_runner.py --init                # Generate config template

Environment variables:
    AJARBOT_SLACK_BOT_TOKEN     # Slack bot token (xoxb-...)
    AJARBOT_SLACK_APP_TOKEN     # Slack app token (xapp-...)
    AJARBOT_TELEGRAM_BOT_TOKEN  # Telegram bot token
"""

import argparse
import asyncio
import signal
import traceback
from dotenv import load_dotenv
from telegram.error import NetworkError as TelegramNetworkError

# Load environment variables from .env file
load_dotenv()

from adapters.base import AdapterConfig
from adapters.runtime import AdapterRuntime
from adapters.slack.adapter import SlackAdapter
from adapters.telegram.adapter import TelegramAdapter
from agent import Agent
from agent_registry import register_agent
from config.config_loader import ConfigLoader
from google_tools.oauth_manager import GoogleOAuthManager
from scheduled_tasks import TaskScheduler

# Adapter class registry mapping platform names to their classes
_ADAPTER_CLASSES = {
    "slack": SlackAdapter,
    "telegram": TelegramAdapter,
}


class BotRunner:
    """Main bot runner that manages all adapters."""

    def __init__(self, config_file: str = "adapters.yaml") -> None:
        self.config_loader = ConfigLoader()
        self.config = self.config_loader.load(config_file)
        self.runtime: AdapterRuntime = None
        self.agent: Agent = None
        self.scheduler: TaskScheduler = None
        self.shutdown_event = asyncio.Event()

    def _load_adapter(self, platform: str) -> bool:
        """Load and register a single adapter. Returns True if loaded."""
        if not self.config_loader.is_adapter_enabled(platform):
            return False

        print(f"\n[Setup] Loading {platform.title()} adapter...")
        adapter_cls = _ADAPTER_CLASSES[platform]
        platform_config = self.config_loader.get_adapter_config(platform)
        adapter = adapter_cls(AdapterConfig(
            platform=platform,
            enabled=True,
            credentials=platform_config.get("credentials", {}),
            settings=platform_config.get("settings", {}),
        ))
        self.runtime.add_adapter(adapter)
        print(f"[Setup] {platform.title()} adapter loaded")
        return True

    def setup(self) -> bool:
        """Set up agent and adapters."""
        print("=" * 60)
        print("Ajarbot Multi-Platform Runner")
        print("=" * 60)

        print("\n[Setup] Initializing agent...")
        self.agent = Agent(
            provider="claude",
            workspace_dir="./memory_workspace",
        )
        print("[Setup] Agent initialized")

        # Register agent in global registry for MCP tool access (delegate_task)
        register_agent(self.agent)

        self.runtime = AdapterRuntime(self.agent)

        # Wire child safety filter (no-op if child_safety block is absent from config)
        try:
            from pathlib import Path as _Path
            from child_safety import ChildSafetyConfig, ChildSafetyFilter, ChildAuditLogger
            _cs_config = ChildSafetyConfig.from_yaml(_Path("config/adapters.local.yaml"))
            if _cs_config:
                _cs_audit = ChildAuditLogger(_Path("./memory_workspace"))
                _cs_audit.cleanup_old_logs(_cs_config.audit_retention_days)
                _cs_filter = ChildSafetyFilter(_cs_config, _cs_audit)
                self.runtime.add_preprocessor(_cs_filter.preprocess_adapter)
                self.runtime.add_postprocessor(_cs_filter.postprocess_adapter)
                print(f"[Setup] Child safety filter active for: {_cs_config.restricted_users}")
        except Exception as _e:
            print(f"[Setup] Child safety filter not loaded: {_e}")

        enabled_count = sum(
            self._load_adapter(platform)
            for platform in _ADAPTER_CLASSES
        )

        # Load user mappings
        # Config format: "platform_userid" — runtime lookup format: "platform:userid"
        user_mapping = self.config_loader.get_user_mapping()
        for platform_user_id, username in user_mapping.items():
            if "_" in platform_user_id:
                platform, uid = platform_user_id.split("_", 1)
                self.runtime.map_user(f"{platform}:{uid}", username)
            else:
                self.runtime.map_user(platform_user_id, username)
            print(f"[Setup] User mapping: {platform_user_id} -> {username}")

        if enabled_count == 0:
            print("\nWARNING: No adapters enabled!")
            print("Edit config/adapters.local.yaml and set enabled: true")
            print("Or run: python bot_runner.py --init")
            return False

        print(f"\n[Setup] {enabled_count} adapter(s) configured")

        # Initialize scheduler
        print("\n[Setup] Initializing task scheduler...")
        self.scheduler = TaskScheduler(
            self.agent,
            config_file="config/scheduled_tasks.yaml"
        )

        # Register adapters with scheduler
        for platform, adapter in self.runtime.registry._adapters.items():
            self.scheduler.add_adapter(platform, adapter)

        # List scheduled tasks
        tasks = self.scheduler.list_tasks()
        enabled_tasks = [t for t in tasks if t.get("enabled")]
        if enabled_tasks:
            print(f"[Setup] {len(enabled_tasks)} scheduled task(s) enabled:")
            for task_info in enabled_tasks:
                print(f"  - {task_info['name']}: {task_info['schedule']}")
                if task_info.get("send_to"):
                    print(f"    → {task_info['send_to']}")

        return True

    async def run(self) -> None:
        """Start all adapters and run until interrupted."""
        if not self.setup():
            return

        # Set up signal handlers
        loop = asyncio.get_running_loop()

        def signal_handler(signum, frame):
            print(f"\n\n[Shutdown] Received signal {signum}...")
            loop.call_soon_threadsafe(self.shutdown_event.set)

        # Register signal handlers (works on both Windows and Unix)
        signal.signal(signal.SIGINT, signal_handler)
        signal.signal(signal.SIGTERM, signal_handler)

        print("\n" + "=" * 60)
        print("Starting bot...")
        print("=" * 60 + "\n")

        _startup_delays = [10, 30, 60, 120, 300]  # backoff between full restart attempts

        for attempt, delay in enumerate([0] + _startup_delays, 0):
            if self.shutdown_event.is_set():
                break
            if attempt > 0:
                print(f"\n[Reconnect] Waiting {delay}s before restart attempt {attempt}...")
                try:
                    await asyncio.wait_for(self.shutdown_event.wait(), timeout=delay)
                    break  # Shutdown was requested during the wait
                except asyncio.TimeoutError:
                    pass
                print(f"[Reconnect] Restarting adapters (attempt {attempt})...")

            try:
                await self.runtime.start()

                # Start scheduler if configured
                if self.scheduler:
                    self.scheduler.start()
                    print("[Scheduler] Task scheduler started\n")

                print("=" * 60)
                print("Bot is running! Press Ctrl+C to stop.")
                print("=" * 60 + "\n")

                # Wait for shutdown signal
                await self.shutdown_event.wait()
                break  # Clean shutdown — don't retry

            except TelegramNetworkError as e:
                print(f"\n[Reconnect] Telegram network error during startup: {e}")
                if attempt >= len(_startup_delays):
                    print("[Reconnect] Max retries reached. Giving up.")
                    traceback.print_exc()
                    break
                # Will retry in next loop iteration
            except Exception as e:
                print(f"\n[Error] {e}")
                traceback.print_exc()
                break  # Non-network errors are not retried
            finally:
                if self.scheduler:
                    self.scheduler.stop()
                    print("[Scheduler] Task scheduler stopped")
                await self.runtime.stop()

        print("\n[Shutdown] Bot stopped cleanly")

    async def health_check(self) -> None:
        """Check health of all adapters."""
        if not self.runtime:
            print("Runtime not initialized")
            return

        status = await self.runtime.health_check()
        print("\n" + "=" * 60)
        print("Health Check")
        print("=" * 60)
        print(f"\nRuntime running: {status['runtime_running']}")
        print("\nAdapters:")
        for platform, adapter_status in status["adapters"].items():
            print(f"\n  {platform.upper()}:")
            for key, value in adapter_status.items():
                print(f"    {key}: {value}")


def main() -> None:
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Ajarbot Multi-Platform Runner"
    )
    parser.add_argument(
        "--config",
        default="adapters.yaml",
        help="Config file to use (default: adapters.yaml)",
    )
    parser.add_argument(
        "--init",
        action="store_true",
        help="Generate config template",
    )
    parser.add_argument(
        "--health",
        action="store_true",
        help="Run health check",
    )
    parser.add_argument(
        "--setup-google",
        action="store_true",
        help="Set up Google OAuth for Gmail/Calendar integration",
    )
    parser.add_argument(
        "--manual",
        action="store_true",
        help="Use manual OAuth code entry (for headless servers)",
    )

    args = parser.parse_args()

    if args.init:
        print("Generating configuration template...")
        loader = ConfigLoader()
        path = loader.save_template()
        print(f"\nConfiguration template created at: {path}")
        print("\nNext steps:")
        print("1. Edit the file with your credentials")
        print("2. Set enabled: true for adapters you want to use")
        print("3. Run: python bot_runner.py")
        return

    if args.setup_google:
        print("=" * 60)
        print("Google OAuth Setup")
        print("=" * 60)
        print()
        oauth_manager = GoogleOAuthManager()

        if oauth_manager.is_authorized():
            print("✓ Already authorized!")
            print(f"✓ Tokens found at {oauth_manager.token_file}")
            print("\nTo re-authorize, delete the token file and run this command again.")
            return

        success = oauth_manager.run_oauth_flow(manual=args.manual)
        if success:
            print("You can now use Gmail and Calendar tools!")
            print("\nTest it:")
            print("  Via Telegram: \"What's on my calendar?\"")
            print("  Via Telegram: \"Send an email to john@example.com\"")
        else:
            print("\nSetup failed. Please check:")
            print("1. config/google_credentials.yaml exists with valid client_id/client_secret")
            print("2. You authorized the app in your browser")
            print("3. No firewall blocking localhost:8080")
        return

    runner = BotRunner(config_file=args.config)

    if args.health:
        asyncio.run(runner.health_check())
        return

    try:
        asyncio.run(runner.run())
    except KeyboardInterrupt:
        print("\nExiting...")


if __name__ == "__main__":
    main()