""" Screenshot functionality for the auto clicker application. """ import base64 import os import time import requests import recorder import logging from typing import Dict from pathlib import Path import json as _json import boto3 from .constants import ( OPENAI_API_KEY, AWS_REGION, LOCAL_DEV, SCREENSHOT_S3_BUCKET, SCREENSHOT_S3_PREFIX, ) from .openai_analyzer import OpenAIAnalyzer from .manage_file_processing import add_or_update_file_record, add_or_append_screenshot_record, add_screenshot_record class ScreenshotManager: """Manages screenshot capture and analysis.""" def __init__(self, app): self.app = app self.openai_analyzer = OpenAIAnalyzer(app) def capture_screenshot(self): """Capture a screenshot of the current screen.""" try: import pyautogui from datetime import datetime # Create screenshots directory if it doesn't exist screenshots_dir = "screenshots" os.makedirs(screenshots_dir, exist_ok=True) # Generate filename with timestamp timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") filename = f"screenshot_{timestamp}.png" filepath = os.path.join(screenshots_dir, filename) # Take screenshot screenshot = pyautogui.screenshot() screenshot.save(filepath) logging.info(f"Screenshot captured: {filepath}") return filepath except Exception as e: logging.error(f"Error capturing screenshot: {e}") return None def capture_manual_screenshot(self): """Capture a manual screenshot and analyze it.""" pl_name = self.app.selected_playlist.get() if not pl_name or pl_name == 'Select playlist': return try: from mysql.mysql_client import get_playlist_id_by_name # Fetch playlist ID from MySQL self.app.show_loader('Fetching playlist from MySQL...') playlist_id = get_playlist_id_by_name(pl_name) if not playlist_id: self.app.hide_loader() return try: # Capture screenshot (prefer full-page Chrome via CDP) self.app.show_loader('Capturing screenshot...') screenshot_path = self._capture_fullpage_chrome_to_playlist(playlist_id, pl_name) if not screenshot_path: # Fallback to legacy full-screen grab screenshot_path = recorder.capture_screenshot( playlist_id, is_manual=True, playlist_name=pl_name ) self.app.hide_loader() # Upload to S3 (if configured) and record manage file processing row try: s3_url = self._upload_screenshot_to_s3(screenshot_path, pl_name) logging.info(f"[ManualScreenshot] Upload result: {s3_url}") except Exception: s3_url = None logging.exception("[ManualScreenshot] S3 upload failed") try: # Use user_id context if available from DI user_id = getattr(self.app, 'current_user_id', None) if user_id: # Prefer to append into a single per-step screenshots record with s3Path as {Bucket, keys: []} bucket = 'big-pond-openai' # Derive key if we uploaded ourselves; else keep local path as fallback if s3_url and s3_url.startswith('s3://'): # s3://bucket/key key = s3_url.split('://', 1)[1].split('/', 1)[1] else: key = os.path.basename(screenshot_path) # Step number from DI context if present step_details = getattr(self.app, 'current_di_step_details', None) or {} step_number = int(step_details.get('stepNumber') or 1) # Do not modify DI fields; we only pass DI step_details.fields (handled in helper) extra_fields = None logging.info(f"[ManualScreenshot] Upserting manage row: user={user_id} step={step_number} key={key}") # Insert one new row per screenshot db_res = add_screenshot_record( user_id=user_id, step_number=step_number, bucket=bucket, key=key, description=f"Manual screenshot for {pl_name}", di_step_details=step_details, flow_type_fallback='analyze', ) logging.info(f"[ManualScreenshot] Manage upsert result: {db_res}") # If uploaded to S3 and DB write succeeded, remove local file (unless debugging) try: from .constants import KEEP_SCREENSHOTS_FOR_DEBUG except Exception: KEEP_SCREENSHOTS_FOR_DEBUG = False if (not KEEP_SCREENSHOTS_FOR_DEBUG) and s3_url and isinstance(db_res, dict) and (db_res.get('updated') or db_res.get('inserted')): try: os.remove(screenshot_path) logging.info(f"[ManualScreenshot] Deleted local screenshot: {screenshot_path}") except Exception: logging.exception("[ManualScreenshot] Failed to delete local screenshot") except Exception: logging.exception("[ManualScreenshot] Manage upsert failed") # Analyze screenshot with OpenAI self.app.show_loader('Analyzing screenshot with OpenAI...') result = self.analyze_screenshot_with_openai(screenshot_path) self.app.hide_loader() # Save analysis to MySQL (legacy/local) self._save_analysis_to_mysql(pl_name, screenshot_path, result) except Exception as e: self.app.hide_loader() except Exception as e: self.app.hide_loader() def analyze_screenshot_with_openai(self, screenshot_path, user_prompt="Describe this screenshot."): """Analyze a screenshot using OpenAI Vision API.""" if not OPENAI_API_KEY: return 'OpenAI API key not set.' try: with open(screenshot_path, 'rb') as f: img_bytes = f.read() img_b64 = base64.b64encode(img_bytes).decode('utf-8') image_url = f"data:image/png;base64,{img_b64}" headers = { 'Authorization': f'Bearer {OPENAI_API_KEY}', 'Content-Type': 'application/json' } data = { "model": "gpt-4o", "messages": [ { "role": "user", "content": [ {"type": "text", "text": user_prompt}, {"type": "image_url", "image_url": {"url": image_url}} ] } ], "max_tokens": 1000 } response = requests.post( "https://api.openai.com/v1/chat/completions", headers=headers, json=data, timeout=60 ) response.raise_for_status() result = response.json() return result['choices'][0]['message']['content'] except Exception as e: return f"OpenAI API error: {e}" def _upload_screenshot_to_s3(self, screenshot_path: str, playlist_name: str | None) -> str | None: """Upload screenshot to S3 and return s3:// URL. Returns None in local dev. Destination per requirement: - Bucket: big-pond-openai - Key prefix: openai/{AWS_REGION}:{}/ - Filename: _ """ try: if LOCAL_DEV: return None s3 = boto3.client('s3', region_name=AWS_REGION) filename = os.path.basename(screenshot_path) ts = int(time.time() * 1000) # Build required destination try: user_uuid = getattr(self.app, 'current_user_id', None) or '' except Exception: user_uuid = '' bucket = 'big-pond-openai' key = f"openai/{AWS_REGION}:{user_uuid}/{ts}_{filename}" s3.upload_file(screenshot_path, bucket, key) return f"s3://{bucket}/{key}" except Exception: return None # --- New centralized upload helpers --- def upload_generic_screenshot_to_openai_s3(self, screenshot_path: str, *, step_details: Dict | None = None) -> Dict[str, str | bool | None]: """Upload a screenshot to the OpenAI S3 location used by normal screenshots. Key format: - With step details: openai/{AWS_REGION}:{userId}/Step{N}-{desc}/{timestamp}_{filename} - Without step details: openai/{AWS_REGION}:{userId}/{timestamp}_{filename} """ try: if LOCAL_DEV: return {"success": False, "skipped": True} s3 = boto3.client('s3', region_name=AWS_REGION) filename = os.path.basename(screenshot_path) ts = int(time.time() * 1000) try: user_uuid = getattr(self.app, 'current_user_id', None) or '' except Exception: user_uuid = '' bucket = 'big-pond-openai' # Optional step folder key_prefix = f"openai/{AWS_REGION}:{user_uuid}" try: if isinstance(step_details, dict): import re as _re step_number = int(step_details.get('stepNumber') or 1) raw_desc = step_details.get('description') or 'step' safe_desc = _re.sub(r'[^A-Za-z0-9 _.-]', '_', str(raw_desc)).strip() or 'step' folder = f"Step{step_number}-{safe_desc}" key_prefix = f"{key_prefix}/{folder}" except Exception: pass key = f"{key_prefix}/{ts}_{filename}" s3.upload_file(screenshot_path, bucket, key) return {"success": True, "s3_bucket": bucket, "s3_key": key, "s3_url": f"s3://{bucket}/{key}"} except Exception as e: try: logging.exception("[Screenshot] Generic S3 upload failed") except Exception: pass return {"success": False, "error": str(e)} def upload_supporting_doc_screenshot_to_openai_s3( self, screenshot_path: str, *, step_details: Dict | None = None, record_manage_table: bool = True, ) -> Dict[str, str | bool | None]: """Upload a screenshot under the supporting_documents structure (kept for supporting_docs flows). Also optionally writes a manage-file-processing row via add_screenshot_record. """ try: if LOCAL_DEV: return {"success": False, "skipped": True} # Derive context try: user_uuid = getattr(self.app, 'current_user_id', None) or '' except Exception: user_uuid = '' bucket = 'big-pond-openai' import re as _re # Step folder step_folder = None try: sd = step_details or {} step_number = int(sd.get('stepNumber') or 1) raw_desc = sd.get('description') or 'step' safe_desc = _re.sub(r'[^A-Za-z0-9 _.-]', '_', str(raw_desc)).strip() or 'step' step_folder = f"Step{step_number}-{safe_desc}" except Exception: step_folder = None base_prefix = f"openai/{AWS_REGION}:{user_uuid}" if step_folder: base_prefix = f"{base_prefix}/{step_folder}/supporting_documents" else: base_prefix = f"{base_prefix}/supporting_documents" # Optional invoice folder (best-effort from app.current_data_row) try: data = getattr(self.app, 'variable_input_data', None) current_row = getattr(self.app, 'current_data_row', 0) invoice_folder = None if isinstance(data, list) and 0 <= int(current_row) < len(data): row = data[int(current_row)] folder_candidate = None if isinstance(row, dict) and row: for _v in row.values(): if _v is not None and str(_v).strip(): folder_candidate = str(_v).strip(); break if not folder_candidate: for _k in ['invoice', 'invoiceNumber', 'invoice_no', 'invoice_num', 'ref', 'reference']: if _k in row and row[_k] is not None and str(row[_k]).strip(): folder_candidate = str(row[_k]).strip(); break if folder_candidate: _s = _re.sub(r'[^A-Za-z0-9 _.-]', '_', folder_candidate).strip().replace('/', '_') if _s: invoice_folder = _s[:64] if invoice_folder: base_prefix = f"{base_prefix}/{invoice_folder}" except Exception: pass # Upload s3 = boto3.client('s3', region_name=AWS_REGION) ts = int(time.time() * 1000) key = f"{base_prefix}/{ts}_{os.path.basename(screenshot_path)}" s3.upload_file(screenshot_path, bucket, key) # Optional DB record if record_manage_table: try: user_id = getattr(self.app, 'current_user_id', None) if user_id: sd2 = step_details or {} try: step_num2 = int(sd2.get('stepNumber') or 1) except Exception: step_num2 = 1 add_screenshot_record( user_id=user_id, step_number=step_num2, bucket=bucket, key=key, description=sd2.get('description') or 'Supporting doc screenshot', di_step_details=sd2, flow_type_fallback='analyze', ) except Exception: pass return {"success": True, "s3_bucket": bucket, "s3_key": key, "s3_url": f"s3://{bucket}/{key}"} except Exception as e: try: logging.exception("[Screenshot] Supporting-docs S3 upload failed") except Exception: pass return {"success": False, "error": str(e)} def capture_and_analyze_supporting_documents(self): """ Capture screenshot, analyze with OpenAI for supporting document links, and perform automated clicking. Works without requiring playlist selection. """ try: self.app.show_loader("Analyzing screenshot for supporting documents...") wait_time = float(os.getenv("SCREENSHOT_PRE_DELAY_SEC", "5.0")) time.sleep(wait_time) # Heartbeat before heavy capture/analysis begins try: hb_interval = float(os.getenv("HB_INTERVAL_SEC", "60")) self.app.send_heartbeat_throttled(min_interval_sec=hb_interval) except Exception: pass # Use the OpenAI analyzer's full resolution screenshot method screenshot_path = self.openai_analyzer.capture_full_resolution_screenshot() if not screenshot_path: self.app.hide_loader() self.app.status_var.set("Failed to capture screenshot") return # Analyze with OpenAI self.app.status_var.set("Analyzing with OpenAI...") # Heartbeat before OpenAI analysis try: hb_interval = float(os.getenv("HB_INTERVAL_SEC", "60")) self.app.send_heartbeat_throttled(min_interval_sec=hb_interval) except Exception: pass clickable_elements = self.openai_analyzer._analyze_screenshot_for_links(screenshot_path) if not clickable_elements: self.app.hide_loader() self.app.status_var.set("No supporting document links found") return # Create minimalist view of found links button_count = sum(1 for link in clickable_elements if link.get('button', False)) text_count = len(clickable_elements) - button_count status_parts = [f"Found {len(clickable_elements)} links"] if button_count > 0: status_parts.append(f"{button_count} SD/button") if text_count > 0: status_parts.append(f"{text_count} text") status_message = f"{'. '.join(status_parts)}. Starting automated clicking..." self.app.status_var.set(status_message) try: # Prefer the analyzer's Rekognition-enabled clicker # Heartbeat before long clicking loop try: hb_interval = float(os.getenv("HB_INTERVAL_SEC", "60")) self.app.send_heartbeat_throttled(min_interval_sec=hb_interval) except Exception: pass self.openai_analyzer._perform_automated_link_clicking( original_screenshot_path=screenshot_path, clickable_links=clickable_elements, playlist_name=None, ) except Exception: # Fallback to legacy clicker if needed self.openai_analyzer.perform_automated_clicking(clickable_elements, screenshot_path) self.app.hide_loader() # Create minimalist completion view button_count = sum(1 for link in clickable_elements if link.get('button', False)) text_count = len(clickable_elements) - button_count completion_parts = [f"Analysis complete! Processed {len(clickable_elements)} links"] if button_count > 0: completion_parts.append(f"{button_count} SD/button") if text_count > 0: completion_parts.append(f"{text_count} text") completion_message = f"{'. '.join(completion_parts)}" self.app.status_var.set(completion_message) except Exception as e: self.app.hide_loader() self.app.status_var.set(f"Error: {str(e)}") logging.error(f"Error in capture_and_analyze_supporting_documents: {e}") def capture_and_analyze_sd_icons(self): """ Capture screenshot, analyze with OpenAI for SD (Supporting Document) icons, and perform automated clicking. Specialized for detecting document icons like the example provided. """ try: self.app.show_loader("Analyzing screenshot for SD icons...") wait_time = float(os.getenv("SCREENSHOT_PRE_DELAY_SEC", "5.0")) time.sleep(wait_time) # Heartbeat before heavy capture/analysis begins try: hb_interval = float(os.getenv("HB_INTERVAL_SEC", "60")) self.app.send_heartbeat_throttled(min_interval_sec=hb_interval) except Exception: pass # Use the OpenAI analyzer's full resolution screenshot method screenshot_path = self.openai_analyzer.capture_full_resolution_screenshot() if not screenshot_path: self.app.hide_loader() self.app.status_var.set("Failed to capture screenshot") return # Analyze with OpenAI using SD icons prompt self.app.status_var.set("Analyzing with OpenAI for SD icons...") # Heartbeat before OpenAI analysis try: hb_interval = float(os.getenv("HB_INTERVAL_SEC", "60")) self.app.send_heartbeat_throttled(min_interval_sec=hb_interval) except Exception: pass clickable_icons = self.openai_analyzer._analyze_screenshot_for_sd_icons(screenshot_path) if not clickable_icons: self.app.hide_loader() self.app.status_var.set("No SD icons found") return # Create minimalist view of found icons icon_count = len(clickable_icons) status_message = f"Found {icon_count} SD icons. Starting automated clicking..." self.app.status_var.set(status_message) try: # Prefer the analyzer's Rekognition-enabled clicker # Heartbeat before long clicking loop try: hb_interval = float(os.getenv("HB_INTERVAL_SEC", "60")) self.app.send_heartbeat_throttled(min_interval_sec=hb_interval) except Exception: pass self.openai_analyzer._perform_automated_link_clicking( original_screenshot_path=screenshot_path, clickable_links=clickable_icons, playlist_name=None, ) except Exception: # Fallback to legacy clicker if needed self.openai_analyzer.perform_automated_clicking(clickable_icons, screenshot_path) self.app.hide_loader() # Create minimalist completion view completion_message = f"Analysis complete! Processed {icon_count} SD icons" self.app.status_var.set(completion_message) except Exception as e: self.app.hide_loader() self.app.status_var.set(f"Error: {str(e)}") logging.error(f"Error in capture_and_analyze_sd_icons: {e}") # --- Full-page Chrome capture helpers ------------------------------------------------- def _capture_fullpage_chrome_to_playlist(self, playlist_id: int, playlist_name: str) -> str | None: """Attempt full-page capture of the active Chrome tab via CDP. Saves under screenshots//fullpage_.png and records in recorder.screenshot_buffer. Returns the file path, or None if CDP capture is unavailable/failed. """ try: out_dir = Path('screenshots') / str(playlist_name) out_dir.mkdir(parents=True, exist_ok=True) ts = int(time.time() * 1000) out_path = out_dir / f"fullpage_{ts}.png" # Try CDP capture png_b64 = self._cdp_capture_fullpage_png() if not png_b64: return None # Write PNG import base64 as _b64 img_bytes = _b64.b64decode(png_b64) with open(out_path, 'wb') as f: f.write(img_bytes) # Record for legacy/local DB like recorder.capture_screenshot rel_path = os.path.relpath(str(out_path)) try: recorder.screenshot_buffer.append({ 'playlist_id': playlist_id, 'path': rel_path, 'is_manual': 1 }) except Exception: pass return str(out_path) except Exception: return None def _cdp_capture_fullpage_png(self) -> str | None: """Use Chrome DevTools Protocol to capture a full-page PNG of the active tab. Requires Chrome launched with --remote-debugging-port (default 9222). Env overrides: - CHROME_REMOTE_DEBUGGING_PORT (e.g., 9222) - CHROME_REMOTE_DEBUGGING_HOST (default: localhost) Returns base64 PNG, or None if not available/failed. """ try: host = os.getenv('CHROME_REMOTE_DEBUGGING_HOST', '127.0.0.1') port = int(os.getenv('CHROME_REMOTE_DEBUGGING_PORT', '9222')) base_url = f"http://{host}:{port}" # Enumerate targets targets = None try: r = requests.get(base_url + "/json/list", timeout=1.5) if r.ok: targets = r.json() except Exception: pass if not targets: try: r = requests.get(base_url + "/json", timeout=1.5) if r.ok: targets = r.json() except Exception: pass if not targets or not isinstance(targets, list): return None # Pick a page target with a websocket URL page = None for t in targets: try: if (t.get('type') == 'page') and t.get('webSocketDebuggerUrl'): page = t break except Exception: continue if not page: return None ws_url = page.get('webSocketDebuggerUrl') if not ws_url: return None # Lazy import to avoid mandatory dependency during non-use try: import websocket # type: ignore except Exception: return None # Connect and issue commands ws = None try: ws = websocket.create_connection(ws_url, timeout=3) msg_id = 0 def _send(method: str, params: dict | None = None) -> dict | None: nonlocal msg_id msg_id += 1 payload = {"id": msg_id, "method": method} if params: payload["params"] = params ws.send(_json.dumps(payload)) # Wait for matching id while True: raw = ws.recv() if not raw: return None data = _json.loads(raw) if isinstance(data, dict) and data.get('id') == msg_id: return data.get('result') or {} # Enable page domain _send('Page.enable') # Get content size (fallback to layout viewport if needed) lm = _send('Page.getLayoutMetrics') or {} content = lm.get('contentSize') or {} width = int(float(content.get('width') or 0)) height = int(float(content.get('height') or 0)) if width <= 0 or height <= 0: # Fallback to viewport metrics v = lm.get('layoutViewport') or {} width = int(float(v.get('clientWidth') or 0)) or 1280 height = int(float(v.get('clientHeight') or 0)) or 800 # Apply device metrics override to encompass full content _send('Emulation.setDeviceMetricsOverride', { 'width': width, 'height': height, 'deviceScaleFactor': 1, 'mobile': False, 'scale': 1 }) # Capture PNG beyond viewport with explicit clip cap = _send('Page.captureScreenshot', { 'format': 'png', 'fromSurface': True, 'captureBeyondViewport': True, 'clip': { 'x': 0, 'y': 0, 'width': float(width), 'height': float(height), 'scale': 1.0 } }) or {} # Clear override (best-effort) try: _send('Emulation.clearDeviceMetricsOverride') except Exception: pass data = cap.get('data') return data if isinstance(data, str) and data else None finally: try: if ws is not None: ws.close() except Exception: pass except Exception: return None