Aetheel/main.py

#!/usr/bin/env python3
"""
Aetheel — Main Entry Point
============================
Starts the AI assistant with multi-channel adapters, memory, skills,
scheduled tasks, and subagent support.

Usage:
    python main.py              Start with Slack + AI handler
    python main.py --telegram   Also enable Telegram adapter
    python main.py --claude     Use Claude Code runtime
    python main.py --test       Echo handler for testing
    python main.py --model anthropic/claude-sonnet-4-20250514
    python main.py --log DEBUG  Debug logging
"""

import argparse
import asyncio
import logging
import os
import re
import sys
import threading
from datetime import datetime

from dotenv import load_dotenv

# Load .env file
load_dotenv()

from adapters.base import BaseAdapter, IncomingMessage
from adapters.slack_adapter import SlackAdapter
from agent.claude_runtime import ClaudeCodeConfig, ClaudeCodeRuntime
from agent.opencode_runtime import (
    AgentResponse,
    OpenCodeConfig,
    OpenCodeRuntime,
    RuntimeMode,
    build_aetheel_system_prompt,
)
from agent.subagent import SubagentManager
from memory import MemoryManager
from memory.types import MemoryConfig
from scheduler import Scheduler
from scheduler.store import ScheduledJob
from skills import SkillsManager

logger = logging.getLogger("aetheel")

# Type alias for either runtime
AnyRuntime = OpenCodeRuntime | ClaudeCodeRuntime

# Global instances (initialized in main)
_runtime: AnyRuntime | None = None
_memory: MemoryManager | None = None
_skills: SkillsManager | None = None
_scheduler: Scheduler | None = None
_subagent_mgr: SubagentManager | None = None
_adapters: dict[str, BaseAdapter] = {}  # source_name -> adapter

# Runtime config (stored for subagent factory)
_use_claude: bool = False
_cli_args: argparse.Namespace | None = None

# Regex for parsing action tags from AI responses
_ACTION_RE = re.compile(r"\[ACTION:remind\|(\d+)\|(.+?)\]", re.DOTALL)
_CRON_RE = re.compile(r"\[ACTION:cron\|([\d\*/,\- ]+)\|(.+?)\]", re.DOTALL)
_SPAWN_RE = re.compile(r"\[ACTION:spawn\|(.+?)\]", re.DOTALL)


# ---------------------------------------------------------------------------
# Message Handlers
# ---------------------------------------------------------------------------


def echo_handler(msg: IncomingMessage) -> str:
    """Simple echo handler for testing."""
    response_lines = [
        f"👋 *Aetheel received your message!*",
        "",
        f"📝 *Text:* {msg.text}",
        f"👤 *From:* {msg.user_name} (`{msg.user_id}`)",
        f"📍 *Channel:* {msg.channel_name} (`{msg.channel_id}`)",
        f"💬 *Source:* {msg.source}",
        f"🧵 *ConvID:* `{msg.conversation_id[:15]}...`",
        f"🕐 *Time:* {msg.timestamp.strftime('%Y-%m-%d %H:%M:%S UTC')}",
        "",
        f"_This is an echo response from the Aetheel test handler._",
    ]
    return "\n".join(response_lines)


def _build_context(msg: IncomingMessage) -> str:
    """
    Build full context to inject into the system prompt.

    Combines:
      - Identity files (SOUL.md, USER.md, MEMORY.md)
      - Relevant memory search results
      - Relevant skills for this message
      - Available skills summary
    """
    global _memory, _skills

    sections: list[str] = []

    # ── Identity: SOUL.md ──
    if _memory:
        soul = _memory.read_soul()
        if soul:
            sections.append(f"# Your Identity (SOUL.md)\n\n{soul}")

        # ── User profile: USER.md ──
        user = _memory.read_user()
        if user:
            sections.append(f"# About the User (USER.md)\n\n{user}")

        # ── Long-term memory: MEMORY.md ──
        ltm = _memory.read_long_term_memory()
        if ltm:
            sections.append(f"# Long-Term Memory (MEMORY.md)\n\n{ltm}")

        # ── Relevant memory search results ──
        try:
            results = asyncio.run(_memory.search(msg.text, max_results=3, min_score=0.2))
            if results:
                snippets = []
                for r in results:
                    if r.path in ("SOUL.md", "USER.md", "MEMORY.md"):
                        continue
                    snippets.append(
                        f"**{r.path}** (lines {r.start_line}-{r.end_line}, "
                        f"relevance {r.score:.0%}):\n{r.snippet[:500]}"
                    )
                if snippets:
                    sections.append(
                        "# Relevant Memory Context\n\n"
                        + "\n\n---\n\n".join(snippets)
                    )
        except Exception as e:
            logger.debug(f"Memory search failed: {e}")

    # ── Skills context ──
    if _skills:
        # Inject matching skill instructions
        skill_context = _skills.get_context(msg.text)
        if skill_context:
            sections.append(skill_context)

        # Always show available skills summary
        skills_summary = _skills.get_all_context()
        if skills_summary:
            sections.append(skills_summary)

    return "\n\n---\n\n".join(sections)


def ai_handler(msg: IncomingMessage) -> str:
    """
    AI-powered handler — the heart of Aetheel.

    Flow:
      Message → context (memory + skills) → system prompt → runtime.chat()
      → action tags → session log → response
    """
    global _runtime, _memory, _scheduler

    if _runtime is None:
        return "⚠️ AI runtime not initialized. Please restart the service."

    text_lower = msg.text.lower().strip()

    # Built-in commands (bypass AI)
    if text_lower in ("status", "/status", "ping"):
        return _format_status()

    if text_lower in ("help", "/help"):
        return _format_help()

    if text_lower == "time":
        return f"🕐 Server time: *{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*"

    if text_lower in ("sessions", "/sessions"):
        return _format_sessions()

    # Cron management commands
    if text_lower.startswith("/cron"):
        return _handle_cron_command(text_lower)

    # Build context from memory + skills
    context = _build_context(msg)

    # Route to AI via runtime
    system_prompt = build_aetheel_system_prompt(
        user_name=msg.user_name,
        channel_name=msg.channel_name,
        is_dm=msg.is_dm,
        extra_context=context,
    )

    response = _runtime.chat(
        message=msg.text,
        conversation_id=msg.conversation_id,
        system_prompt=system_prompt,
    )

    if not response.ok:
        error_msg = response.error or "Unknown error"
        logger.error(f"AI error: {error_msg}")

        if "not found" in error_msg.lower() or "not installed" in error_msg.lower():
            return (
                "⚠️ AI CLI is not available.\n"
                "Check the runtime installation docs."
            )
        if "timeout" in error_msg.lower():
            return (
                "⏳ The AI took too long to respond. "
                "Try a shorter or simpler question."
            )
        return f"⚠️ AI error: {error_msg[:200]}"

    logger.info(
        f"🤖 AI response: {len(response.text)} chars, "
        f"{response.duration_ms}ms"
    )

    # Parse and execute action tags (reminders, cron, spawn)
    reply_text = _process_action_tags(response.text, msg)

    # Log conversation to memory session log
    if _memory:
        try:
            channel = "dm" if msg.is_dm else msg.channel_name or msg.source
            _memory.log_session(
                f"**User ({msg.user_name}):** {msg.text}\n\n"
                f"**Aetheel:** {reply_text}",
                channel=channel,
            )
        except Exception as e:
            logger.debug(f"Session logging failed: {e}")

    return reply_text


# ---------------------------------------------------------------------------
# Action Tag Processing
# ---------------------------------------------------------------------------


def _process_action_tags(text: str, msg: IncomingMessage) -> str:
    """
    Parse and execute action tags from the AI response.

    Supports:
      [ACTION:remind|<minutes>|<message>]     → one-shot reminder
      [ACTION:cron|<cron_expr>|<prompt>]       → recurring cron job
      [ACTION:spawn|<task description>]        → background subagent
    """
    cleaned = text

    # ── Remind tags (one-shot) ──
    for match in _ACTION_RE.finditer(text):
        minutes_str, reminder_msg = match.group(1), match.group(2)
        try:
            minutes = int(minutes_str)
            if _scheduler:
                _scheduler.add_once(
                    delay_minutes=minutes,
                    prompt=reminder_msg.strip(),
                    channel_id=msg.channel_id,
                    channel_type=msg.source,
                    thread_id=msg.raw_event.get("thread_id"),
                    user_name=msg.user_name,
                )
            logger.info(
                f"⏰ Reminder scheduled: '{reminder_msg.strip()[:50]}' "
                f"in {minutes} min for {msg.source}/{msg.channel_name}"
            )
        except Exception as e:
            logger.warning(f"Failed to schedule reminder: {e}")
        cleaned = cleaned.replace(match.group(0), "").strip()

    # ── Cron tags (recurring) ──
    for match in _CRON_RE.finditer(text):
        cron_expr, cron_prompt = match.group(1).strip(), match.group(2).strip()
        try:
            if _scheduler:
                job_id = _scheduler.add_cron(
                    cron_expr=cron_expr,
                    prompt=cron_prompt,
                    channel_id=msg.channel_id,
                    channel_type=msg.source,
                    thread_id=msg.raw_event.get("thread_id"),
                    user_name=msg.user_name,
                )
                logger.info(
                    f"🔄 Cron scheduled: '{cron_prompt[:50]}' ({cron_expr}) "
                    f"job_id={job_id}"
                )
        except Exception as e:
            logger.warning(f"Failed to schedule cron job: {e}")
        cleaned = cleaned.replace(match.group(0), "").strip()

    # ── Spawn tags (subagent) ──
    for match in _SPAWN_RE.finditer(text):
        spawn_task = match.group(1).strip()
        try:
            if _subagent_mgr:
                task_id = _subagent_mgr.spawn(
                    task=spawn_task,
                    channel_id=msg.channel_id,
                    channel_type=msg.source,
                    thread_id=msg.raw_event.get("thread_id"),
                    user_name=msg.user_name,
                )
                logger.info(
                    f"🚀 Subagent spawned: '{spawn_task[:50]}' "
                    f"task_id={task_id}"
                )
        except Exception as e:
            logger.warning(f"Failed to spawn subagent: {e}")
        cleaned = cleaned.replace(match.group(0), "").strip()

    return cleaned


# ---------------------------------------------------------------------------
# Cron Management Commands
# ---------------------------------------------------------------------------


def _handle_cron_command(text: str) -> str:
    """Handle /cron subcommands."""
    global _scheduler

    if not _scheduler:
        return "⚠️ Scheduler not initialized."

    parts = text.strip().split(maxsplit=2)

    if len(parts) < 2 or parts[1] == "list":
        jobs = _scheduler.list_jobs()
        if not jobs:
            return "📋 No scheduled jobs."
        lines = ["📋 *Scheduled Jobs:*\n"]
        for job in jobs:
            kind = f"🔄 `{job.cron_expr}`" if job.is_recurring else "⏰ one-shot"
            lines.append(
                f"• `{job.id}` — {kind} — {job.prompt[:60]}"
            )
        return "\n".join(lines)

    if parts[1] == "remove" and len(parts) >= 3:
        job_id = parts[2].strip()
        if _scheduler.remove(job_id):
            return f"✅ Job `{job_id}` removed."
        return f"⚠️ Job `{job_id}` not found."

    return (
        "Usage: `/cron list` or `/cron remove <id>`"
    )


# ---------------------------------------------------------------------------
# Scheduler Callback
# ---------------------------------------------------------------------------


def _on_scheduled_job(job: ScheduledJob) -> None:
    """
    Called by the scheduler when a job fires.

    Creates a synthetic IncomingMessage and routes it through ai_handler,
    then sends the response to the right channel.
    """
    logger.info(f"🔔 Scheduled job firing: {job.id} — '{job.prompt[:50]}'")

    # Build a synthetic message
    msg = IncomingMessage(
        text=job.prompt,
        user_id="system",
        user_name=job.user_name or "Scheduler",
        channel_id=job.channel_id,
        channel_name=f"scheduled-{job.id}",
        conversation_id=f"cron-{job.id}",
        source=job.channel_type,
        is_dm=True,
        raw_event={"thread_id": job.thread_id},
    )

    # Route through the AI handler
    try:
        response = ai_handler(msg)
        if response:
            _send_to_channel(
                channel_id=job.channel_id,
                text=response,
                thread_id=job.thread_id,
                channel_type=job.channel_type,
            )
    except Exception as e:
        logger.error(f"Scheduled job {job.id} handler failed: {e}", exc_info=True)


# ---------------------------------------------------------------------------
# Multi-Channel Send
# ---------------------------------------------------------------------------


def _send_to_channel(
    channel_id: str,
    text: str,
    thread_id: str | None,
    channel_type: str,
) -> None:
    """
    Send a message to a specific channel via the right adapter.

    Used by the scheduler and subagent manager to route responses
    back to the correct platform.
    """
    adapter = _adapters.get(channel_type)
    if adapter:
        adapter.send_message(
            channel_id=channel_id,
            text=text,
            thread_id=thread_id,
        )
    else:
        # Fallback: try the first available adapter
        for a in _adapters.values():
            a.send_message(channel_id=channel_id, text=text, thread_id=thread_id)
            break
        else:
            logger.warning(
                f"No adapter for '{channel_type}' — cannot send message"
            )


# ---------------------------------------------------------------------------
# Runtime Factory (for subagents)
# ---------------------------------------------------------------------------


def _make_runtime() -> AnyRuntime:
    """Create a fresh runtime instance (used by subagent manager)."""
    global _use_claude, _cli_args

    if _use_claude:
        config = ClaudeCodeConfig.from_env()
        if _cli_args and _cli_args.model:
            config.model = _cli_args.model
        return ClaudeCodeRuntime(config)
    else:
        config = OpenCodeConfig.from_env()
        if _cli_args and _cli_args.model:
            config.model = _cli_args.model
        return OpenCodeRuntime(config)


# ---------------------------------------------------------------------------
# Formatting Helpers
# ---------------------------------------------------------------------------


def _format_status() -> str:
    """Format the /status response with runtime info."""
    global _runtime, _scheduler, _skills, _subagent_mgr

    lines = [
        "🟢 *Aetheel is online*",
        "",
    ]

    if _runtime:
        status = _runtime.get_status()
        lines.extend([
            f"• *Mode:* {status['mode']}",
            f"• *Model:* {status['model']}",
            f"• *Provider:* {status['provider']}",
            f"• *Active Sessions:* {status['active_sessions']}",
        ])

    # Adapter status
    if _adapters:
        adapter_names = ", ".join(_adapters.keys())
        lines.append(f"• *Channels:* {adapter_names}")

    # Skills status
    if _skills:
        skill_count = len(_skills.skills)
        lines.append(f"• *Skills Loaded:* {skill_count}")

    # Scheduler status
    if _scheduler:
        jobs = _scheduler.list_jobs()
        lines.append(f"• *Scheduled Jobs:* {len(jobs)}")

    # Subagents status
    if _subagent_mgr:
        active = _subagent_mgr.list_active()
        lines.append(f"• *Active Subagents:* {len(active)}")

    lines.extend([
        "",
        f"• *Time:* {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
    ])

    return "\n".join(lines)


def _format_help() -> str:
    """Format the /help response."""
    return (
        "🦾 *Aetheel — AI-Powered Assistant*\n"
        "\n"
        "*Built-in Commands:*\n"
        "• `status` — Check bot and AI runtime status\n"
        "• `help` — Show this help message\n"
        "• `time` — Current server time\n"
        "• `sessions` — Active session count\n"
        "• `/cron list` — List scheduled jobs\n"
        "• `/cron remove <id>` — Remove a scheduled job\n"
        "\n"
        "*AI Chat:*\n"
        "• Send any message and the AI will respond\n"
        "• Each thread maintains its own conversation\n"
        "• DMs work too — just message me directly\n"
        "\n"
        "*AI Actions:*\n"
        "• The AI can schedule reminders\n"
        "• The AI can set up recurring cron jobs\n"
        "• The AI can spawn background subagents for long tasks\n"
    )


def _format_sessions() -> str:
    """Format session info."""
    global _runtime
    if _runtime:
        count = _runtime.get_status()["active_sessions"]
        cleaned = _runtime.cleanup_sessions()
        return (
            f"🧵 *Active Sessions:* {count}\n"
            f"🧹 *Cleaned up:* {cleaned} stale sessions"
        )
    return "⚠️ Runtime not initialized."


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main():
    global _runtime, _memory, _skills, _scheduler, _subagent_mgr
    global _adapters, _use_claude, _cli_args

    parser = argparse.ArgumentParser(
        description="Aetheel — AI-Powered Personal Assistant",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""\
Examples:
    python main.py              Start with Slack + AI handler
    python main.py --telegram   Also enable Telegram adapter
    python main.py --claude     Use Claude Code runtime
    python main.py --test       Echo-only handler
    python main.py --model anthropic/claude-sonnet-4-20250514
    python main.py --log DEBUG  Debug logging
        """,
    )
    parser.add_argument("--test", action="store_true", help="Use echo handler for testing")
    parser.add_argument("--claude", action="store_true", help="Use Claude Code runtime")
    parser.add_argument("--cli", action="store_true", help="Force CLI mode (OpenCode)")
    parser.add_argument("--sdk", action="store_true", help="Force SDK mode (OpenCode)")
    parser.add_argument("--telegram", action="store_true", help="Enable Telegram adapter")
    parser.add_argument("--model", default=None, help="Model to use")
    parser.add_argument(
        "--log",
        default=os.environ.get("LOG_LEVEL", "INFO"),
        help="Log level (DEBUG, INFO, WARNING, ERROR)",
    )
    args = parser.parse_args()
    _cli_args = args
    _use_claude = args.claude

    # Configure logging
    logging.basicConfig(
        level=getattr(logging, args.log.upper(), logging.INFO),
        format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )

    # -------------------------------------------------------------------
    # 1. Initialize Memory System
    # -------------------------------------------------------------------

    workspace_dir = os.environ.get(
        "AETHEEL_WORKSPACE", os.path.expanduser("~/.aetheel/workspace")
    )
    db_path = os.environ.get(
        "AETHEEL_MEMORY_DB", os.path.expanduser("~/.aetheel/memory.db")
    )

    try:
        mem_config = MemoryConfig(
            workspace_dir=workspace_dir,
            db_path=db_path,
        )
        _memory = MemoryManager(mem_config)
        logger.info(f"Memory system initialized: workspace={workspace_dir}")

        stats = asyncio.run(_memory.sync())
        logger.info(
            f"Memory sync: {stats.get('files_indexed', 0)} files indexed, "
            f"{stats.get('chunks_created', 0)} chunks"
        )
    except Exception as e:
        logger.warning(f"Memory system init failed (continuing without): {e}")
        _memory = None

    # -------------------------------------------------------------------
    # 2. Initialize Skills System
    # -------------------------------------------------------------------

    try:
        _skills = SkillsManager(workspace_dir)
        loaded = _skills.load_all()
        logger.info(f"Skills system initialized: {len(loaded)} skill(s)")
    except Exception as e:
        logger.warning(f"Skills system init failed (continuing without): {e}")
        _skills = None

    # -------------------------------------------------------------------
    # 3. Initialize AI Runtime
    # -------------------------------------------------------------------

    runtime_label = "echo (test mode)"
    if not args.test:
        if args.claude:
            claude_config = ClaudeCodeConfig.from_env()
            if args.model:
                claude_config.model = args.model
            _runtime = ClaudeCodeRuntime(claude_config)
            runtime_label = f"claude-code, model={claude_config.model or 'default'}"
        else:
            config = OpenCodeConfig.from_env()
            if args.cli:
                config.mode = RuntimeMode.CLI
            elif args.sdk:
                config.mode = RuntimeMode.SDK
            if args.model:
                config.model = args.model
            _runtime = OpenCodeRuntime(config)
            runtime_label = (
                f"opencode/{config.mode.value}, "
                f"model={config.model or 'default'}"
            )

    # -------------------------------------------------------------------
    # 4. Initialize Scheduler
    # -------------------------------------------------------------------

    try:
        _scheduler = Scheduler(callback=_on_scheduled_job)
        _scheduler.start()
        logger.info("Scheduler initialized")
    except Exception as e:
        logger.warning(f"Scheduler init failed (continuing without): {e}")
        _scheduler = None

    # -------------------------------------------------------------------
    # 5. Initialize Subagent Manager
    # -------------------------------------------------------------------

    if _runtime:
        try:
            _subagent_mgr = SubagentManager(
                runtime_factory=_make_runtime,
                send_fn=_send_to_channel,
            )
            logger.info("Subagent manager initialized")
        except Exception as e:
            logger.warning(f"Subagent manager init failed: {e}")
            _subagent_mgr = None

    # -------------------------------------------------------------------
    # 6. Initialize Channel Adapters
    # -------------------------------------------------------------------

    # Choose the message handler
    handler = echo_handler if args.test else ai_handler

    # Slack adapter (always enabled if tokens are present)
    slack_token = os.environ.get("SLACK_BOT_TOKEN")
    slack_app_token = os.environ.get("SLACK_APP_TOKEN")

    if slack_token and slack_app_token:
        try:
            slack = SlackAdapter(log_level=args.log)
            slack.on_message(handler)
            _adapters["slack"] = slack
            logger.info("Slack adapter registered")
        except Exception as e:
            logger.error(f"Slack adapter failed to initialize: {e}")
    else:
        logger.warning("Slack tokens not set — Slack adapter disabled")

    # Telegram adapter (enabled with --telegram flag)
    if args.telegram:
        telegram_token = os.environ.get("TELEGRAM_BOT_TOKEN")
        if telegram_token:
            try:
                from adapters.telegram_adapter import TelegramAdapter

                telegram = TelegramAdapter()
                telegram.on_message(handler)
                _adapters["telegram"] = telegram
                logger.info("Telegram adapter registered")
            except Exception as e:
                logger.error(f"Telegram adapter failed to initialize: {e}")
        else:
            logger.error(
                "TELEGRAM_BOT_TOKEN not set — cannot enable Telegram. "
                "Get a token from @BotFather on Telegram."
            )

    if not _adapters:
        print("❌ No channel adapters initialized!")
        print("   Set SLACK_BOT_TOKEN + SLACK_APP_TOKEN or use --telegram")
        sys.exit(1)

    # Start file watching for automatic memory re-indexing
    if _memory:
        _memory.start_watching()

    # -------------------------------------------------------------------
    # 7. Start Adapters
    # -------------------------------------------------------------------

    logger.info("=" * 60)
    logger.info("  Aetheel Starting")
    logger.info("=" * 60)
    logger.info(f"  Runtime:    {runtime_label}")
    logger.info(f"  Channels:   {', '.join(_adapters.keys())}")
    logger.info(f"  Skills:     {len(_skills.skills) if _skills else 0}")
    logger.info(f"  Scheduler:  {'✅' if _scheduler else '❌'}")
    logger.info(f"  Subagents:  {'✅' if _subagent_mgr else '❌'}")
    logger.info("=" * 60)

    try:
        if len(_adapters) == 1:
            # Single adapter — start it blocking
            adapter = next(iter(_adapters.values()))
            adapter.start()
        else:
            # Multiple adapters — start all but last async, last blocking
            adapter_list = list(_adapters.values())
            for adapter in adapter_list[:-1]:
                adapter.start_async()
            adapter_list[-1].start()  # blocking

    except KeyboardInterrupt:
        logger.info("Shutting down...")
    finally:
        # Cleanup
        for adapter in _adapters.values():
            try:
                adapter.stop()
            except Exception:
                pass
        if _scheduler:
            _scheduler.stop()
        if _memory:
            _memory.close()
        logger.info("Aetheel stopped. Goodbye! 👋")


if __name__ == "__main__":
    main()