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

# 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 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")

        self.runtime = AdapterRuntime(self.agent)

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

        # Load user mappings
        user_mapping = self.config_loader.get_user_mapping()
        for platform_user_id, username in user_mapping.items():
            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")

        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()

        except Exception as e:
            print(f"\n[Error] {e}")
            traceback.print_exc()
        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()