"""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._advances = 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, "auto_next": config.AUTO_NEXT, "advances": self._advances, } 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._advances = 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}" # Auto-advance: click a Next control, scroll to top, capture again. if config.AUTO_NEXT and not self._stop.is_set() and self._advance_next(): continue 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 def _advance_next(self) -> bool: """Look for a 'Next' control on screen; if found, click it and scroll to top so the loop captures the next page. Returns True if it advanced.""" if self._advances >= config.MAX_AUTO_NEXT: self.last_event = f"auto-next cap ({config.MAX_AUTO_NEXT}) reached" return False try: shot = screenshot.capture(prefix="next") info = extractor.find_next_button(shot) if not info.get("found"): return False self.last_event = f"clicking '{info.get('label') or 'Next'}' → next page" scroller.click_at_fraction(float(info.get("x", 0)), float(info.get("y", 0))) time.sleep(config.NEXT_LOAD_PAUSE) scroller.scroll_to_top() self._advances += 1 return True except Exception as e: self.last_event = f"next-button step skipped: {e}" return False # module-level singleton shared by the API controller = CaptureController()