sos-bot/core/engine.py

"""
Engine — the main loop that ties everything together:
  screenshot → detect state → lookup playbook → run actions.
"""
import logging
import os
import time
import yaml

from .adb import AdbClient
from .vision import Vision
from .action_loader import load_actions

logger = logging.getLogger("sos.engine")


class Engine:
    """
    Runs the bot loop.  Each iteration:
      1. Capture screenshot (in-memory)
      2. Evaluate each playbook step (top-to-bottom)
      3. Execute the first matching step's action sequence
      4. Sleep the configured delay
    """

    def __init__(self, playbook_path: str, actions_dir: str = None):
        self.playbook_path = playbook_path
        self.actions_dir = actions_dir
        self.adb = AdbClient()
        self.vision = Vision(self.adb)
        self.registry = None
        self.running = False
        self._step_count = 0

    # ── lifecycle ────────────────────────────────────────────

    def start(self):
        """Load config, connect ADB, load actions, enter main loop."""
        import os
        # Resolve paths
        pb_path = os.path.abspath(self.playbook_path)
        project_root = os.path.dirname(os.path.dirname(pb_path))
        self._templates_dir = os.path.join(project_root, "templates")
        self._project_root = project_root

        # 1. Load playbook
        with open(pb_path) as f:
            config = yaml.safe_load(f)
        self.config = config
        logger.info("Loaded playbook: %s", self.playbook_path)

        # 2. Connect ADB
        host = config.get("adb_host", "127.0.0.1")
        port = config.get("adb_port", 5555)
        self.adb = AdbClient(host=host, port=port)
        self.adb.connect()
        self.vision = Vision(self.adb)

        if not self.adb.is_connected():
            raise RuntimeError(f"ADB not connected to {host}:{port}")

        logger.info("ADB connected: %s", self.adb.serial)

        # 3. Load actions
        actions_dir = self.actions_dir or os.path.join(project_root, "actions")
        self.registry = load_actions(actions_dir)

        # 4. Main loop
        self.running = True
        self._loop()

    def stop(self):
        """Signal the main loop to stop."""
        self.running = False
        logger.info("Stop signalled — finishing current iteration")

    # ── main loop ────────────────────────────────────────────

    def _loop(self):
        delay = self.config.get("loop_delay", 2.0)
        steps = self.config.get("steps", [])

        logger.info("Bot loop started (%d steps, delay=%.1fs)", len(steps), delay)

        while self.running:
            self._step_count += 1
            try:
                # 1. Capture
                self.vision.capture()

                # 2. Evaluate steps in order
                matched = False
                for step in steps:
                    if self._evaluate_step(step):
                        logger.info(
                            "Step #%d — %s — matched → executing",
                            self._step_count,
                            step.get("name", "?"),
                        )
                        self._execute_actions(step.get("then", []))
                        matched = True
                        break  # only first match per iteration

                if not matched:
                    logger.debug("Step #%d — no step matched", self._step_count)

            except KeyboardInterrupt:
                self.running = False
                break
            except Exception as e:
                logger.error("Loop error on step #%d: %s", self._step_count, e)
                import traceback
                traceback.print_exc()

            if self.running:
                time.sleep(delay)

        logger.info("Bot loop ended (%d iterations)", self._step_count)

    # ── step evaluation ──────────────────────────────────────

    def _evaluate_step(self, step: dict) -> bool:
        """
        Check if a playbook step's condition matches the current screen.
        Conditions (checked in order):
          - if_template:  template matching
          - if_pixel:     pixel color check
          - if_multiple:  all templates must match (AND)
        Returns True if the step's trigger condition is met.
        """
        # Template match
        if "if_template" in step:
            tpl = step["if_template"]
            conf = step.get("confidence", 0.8)
            region = step.get("region")  # optional [x1,y1,x2,y2]
            tpl_path = os.path.join(self._templates_dir, tpl)
            return self.vision.find_template(tpl_path, conf, region) is not None

        # Multiple templates (AND)
        if "if_multiple" in step:
            for tpl in step["if_multiple"]:
                tpl_path = os.path.join(self._templates_dir, tpl)
                conf = step.get("confidence", 0.8)
                if self.vision.find_template(tpl_path, conf) is None:
                    return False
            return True

        # Pixel color match
        if "if_pixel" in step:
            px = step["if_pixel"]
            tol = px.get("tolerance", 10)
            color = tuple(px["color"])
            return self.vision.color_match(px["x"], px["y"], color, tol)

        # Counter-based trigger (every N iterations)
        if "if_counter" in step:
            n = step["if_counter"]
            return self._step_count > 0 and self._step_count % n == 0

        # Default: always match (fallback / idle step)
        return step.get("always", False)

    # ── action execution ─────────────────────────────────────

    def _execute_actions(self, actions: list):
        """Run a sequence of actions."""
        ctx = {
            "adb": self.adb,
            "vision": self.vision,
            "engine": self,
            "screenshot": self.vision.current,
            "log": logger,
            "_templates_dir": self._templates_dir,
        }
        for item in actions:
            name = item["action"]
            params = item.get("params", {})
            try:
                self.registry.run(name, ctx, params)
            except KeyError as e:
                logger.error("Action error: %s", e)
                break
            except Exception as e:
                logger.error("Action '%s' failed: %s", name, e)
                raise