screen-leads-app/core/controller.py

"""Capture controller: a background state machine driven by the dashboard.

States: IDLE -> RUNNING <-> PAUSED -> STOPPED. The loop detects the front tab,
stops (or waits) on a non-target page per config, and for each new profile URL
runs scroll -> screenshot -> extract -> normalise -> save.
"""
from __future__ import annotations

import threading
import time
import traceback
from enum import Enum
from typing import Any, Optional

import config
from agent import detector, screenshot, scroller
from ai import extractor, recipes
from core.models import Lead
from db import database


class State(str, Enum):
    IDLE = "IDLE"
    RUNNING = "RUNNING"
    PAUSED = "PAUSED"
    STOPPED = "STOPPED"


def _recipe_by_id(site_id: str):
    for r in recipes.REGISTRY:
        if r.site_id == site_id:
            return r
    return None


class CaptureController:
    def __init__(self) -> None:
        self._state = State.IDLE
        self._thread: Optional[threading.Thread] = None
        self._stop = threading.Event()
        self._pause = threading.Event()
        self._lock = threading.Lock()
        self.reason = ""
        self.last_event = ""
        self.captured_urls: set[str] = set()
        self.leads_captured = 0
        self._run_id: Optional[int] = None

    # --- introspection -------------------------------------------------------
    def status(self) -> dict[str, Any]:
        return {
            "state": self._state.value,
            "reason": self.reason,
            "last_event": self.last_event,
            "leads_captured": self.leads_captured,
            "captured_urls": len(self.captured_urls),
            "on_invalid_page": config.ON_INVALID_PAGE,
        }

    def _set(self, state: State, event: str = "") -> None:
        self._state = state
        if event:
            self.last_event = event

    # --- controls ------------------------------------------------------------
    def start(self) -> dict[str, Any]:
        with self._lock:
            if self._state in (State.RUNNING, State.PAUSED):
                return self.status()
            self._stop.clear()
            self._pause.clear()
            self.reason = ""
            self.leads_captured = 0
            self.captured_urls.clear()
            self._run_id = database.start_run()
            self._set(State.RUNNING, "started")
            self._thread = threading.Thread(target=self._loop, daemon=True)
            self._thread.start()
        return self.status()

    def pause(self) -> dict[str, Any]:
        if self._state == State.RUNNING:
            self._pause.set()
            self._set(State.PAUSED, "paused")
        return self.status()

    def resume(self) -> dict[str, Any]:
        if self._state == State.PAUSED:
            self._pause.clear()
            self._set(State.RUNNING, "resumed")
        return self.status()

    def stop(self) -> dict[str, Any]:
        if self._state in (State.RUNNING, State.PAUSED):
            self._stop.set()
            self._pause.clear()
            self.reason = self.reason or "stopped by user"
        return self.status()

    # --- main loop -----------------------------------------------------------
    def _finish(self, status: str, reason: str) -> None:
        self.reason = reason
        self._set(State.STOPPED, reason)
        if self._run_id is not None:
            database.finish_run(self._run_id, status, reason, self.leads_captured)
            self._run_id = None

    def _countdown(self) -> None:
        """Grace period so the user can focus the browser before capture."""
        remaining = int(round(config.START_DELAY))
        while remaining > 0 and not self._stop.is_set():
            self.last_event = f"focus your browser — starting in {remaining}s"
            time.sleep(1)
            remaining -= 1

    def _loop(self) -> None:
        try:
            self._countdown()
            while not self._stop.is_set():
                if self._pause.is_set():
                    time.sleep(0.3)
                    continue

                det = detector.detect()

                if not det.ok:
                    if config.ON_INVALID_PAGE == "wait":
                        self.last_event = f"waiting: {det.reason}"
                        time.sleep(config.WATCH_INTERVAL)
                        continue
                    self._finish("stopped", det.reason)
                    return

                url = det.url or ""
                if url and url in self.captured_urls:
                    self.last_event = "waiting for a new profile"
                    time.sleep(config.WATCH_INTERVAL)
                    continue

                recipe = _recipe_by_id(det.recipe_id or "")
                if recipe is None:
                    self._finish("error", f"no recipe for {det.recipe_id}")
                    return

                self.last_event = f"capturing {url or det.recipe_id}"
                lead = self._capture(recipe, url)
                if lead is not None:
                    self.leads_captured += 1
                    if url:
                        self.captured_urls.add(url)
                    self.last_event = f"saved: {lead.full_name or url}"

                time.sleep(config.WATCH_INTERVAL)

            self._finish("stopped", self.reason or "stopped by user")
        except Exception as e:  # keep the loop crash-visible on the dashboard
            self._finish("error", f"{e}\n{traceback.format_exc(limit=2)}")

    def _capture(self, recipe, url: str) -> Optional[Lead]:
        # scroll + screenshot the full page
        try:
            scroller.scroll_to_top()
        except Exception as e:
            self.last_event = f"scroll unavailable: {e}"
        shots: list[str] = []
        for i in range(config.SCROLL_STEPS):
            if self._stop.is_set():
                break
            shots.append(screenshot.capture(prefix=f"{recipe.site_id}"))
            try:
                scroller.scroll_down()
            except Exception:
                break
        if not shots:
            return None

        raw = extractor.extract_lead(recipe, shots)
        lead = recipe.to_canonical(raw)
        lead.source_url = url
        lead.screenshot_refs = shots
        database.upsert_lead(lead)
        return lead


# module-level singleton shared by the API
controller = CaptureController()