# Simple Google Search with BrowserAgent
# =======================================
# This example demonstrates basic BrowserAgent capabilities for automated Google searches.
# 
# Features Demonstrated:
# 1. Basic Google search automation
# 2. Simple popup handling
# 3. Content extraction from search results
# 4. Vision-based element detection
#
# Workflow: Navigate → Search → Extract Results → Report

import asyncio
import logging
import os
import sys
from pathlib import Path

# Add src to Python path
sys.path.insert(0, str(Path(__file__).parent / "src"))

from agents.browser_agent import BrowserAgent
from agents.utils import init_agent_logging
from models.models import ModelConfig

# Configure logging
init_agent_logging(level=logging.INFO, clear_existing_handlers=True)
logger = logging.getLogger("SimpleGoogleSearch")


async def main():
    """Simple Google search demonstration."""
    
    # Load environment variables from .env file if it exists
    env_file_path = Path(__file__).parent / ".env"
    if env_file_path.exists():
        try:
            from dotenv import load_dotenv
            load_dotenv(env_file_path)
            logger.info(f"Loaded environment variables from {env_file_path}")
        except ImportError:
            logger.warning("python-dotenv not installed. Install with: pip install python-dotenv")
            logger.info("Proceeding with system environment variables only")
        except Exception as e:
            logger.warning(f"Failed to load .env file: {e}")
    else:
        logger.info("No .env file found, using system environment variables")
    
    # Get API key
    OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
    if not OPENAI_API_KEY:
        logger.error("OPENAI_API_KEY environment variable not set")
        logger.info("Please set OPENAI_API_KEY in your environment or create a .env file with:")
        logger.info("OPENAI_API_KEY=your_api_key_here")
        raise ValueError("OPENAI_API_KEY environment variable is required")
    
    # Model configurations
    browser_model_config = ModelConfig(
        type="api",
        provider="openai",
        name="gpt-4.1-mini",
        temperature=0.1,
        max_tokens=2000,
        api_key=OPENAI_API_KEY,
    )
    
    vision_model_config = ModelConfig(
        type="api",
        provider="openai", 
        name="gpt-4.1",
        temperature=0.2,
        max_tokens=1000,
        api_key=OPENAI_API_KEY,
    )
    
    # Create browser agent
    logger.info("Creating BrowserAgent...")
    browser_agent = await BrowserAgent.create_safe(
        model_config=browser_model_config,
        vision_model_config=vision_model_config,
        description="You are a web automation specialist that can perform Google searches and extract content.",
        agent_name="SimpleSearchAgent",
        headless=True,  # Change to False to see the browser
        browser_channel="chrome",
        auto_screenshot=True,  # Enable automatic screenshots for vision analysis
    )
    
    try:
        # Define search query
        search_query = "Python web automation 2024"
        logger.info(f"Performing search for: '{search_query}'")
        
        # Step 1: Navigate to Google
        logger.info("Step 1: Navigating to Google...")
        await browser_agent.navigate("https://google.com")
        
        # Take screenshot to see current state
        screenshot_path = "current_google_page.png"
        await browser_agent.screenshot(screenshot_path)
        logger.info(f"Screenshot saved: {screenshot_path}")
        
        # Step 2: Handle any popups using vision if needed
        logger.info("Step 2: Checking for popups...")
        try:
            # Use vision to detect and handle cookie consent or other popups
            popup_elements = await browser_agent.predict_interactive_elements(
                screenshot_path, 
                "Find any cookie consent, privacy policy, or popup buttons that need to be clicked to proceed"
            )
            
            if popup_elements and popup_elements.get('elements'):
                logger.info("Found popup elements, attempting to handle...")
                # Try to click accept/dismiss buttons
                for element in popup_elements['elements'][:2]:  # Try first 2 elements
                    try:
                        x, y = element['center']
                        await browser_agent.click(x, y)
                        await asyncio.sleep(1)
                        logger.info(f"Clicked popup element at ({x}, {y})")
                        break
                    except Exception as e:
                        logger.debug(f"Failed to click popup element: {e}")
            else:
                logger.info("No popups detected")
        except Exception as e:
            logger.warning(f"Error handling popups: {e}")
        
        # Step 3: Find search box and perform search
        logger.info("Step 3: Performing search...")
        try:
            # Try to find search input by common selectors
            search_selectors = [
                'input[name="q"]',
                'input[title="Search"]', 
                'textarea[name="q"]',
                'input[type="text"]'
            ]
            
            search_input_found = False
            for selector in search_selectors:
                try:
                    await browser_agent.fill(selector, search_query)
                    await browser_agent.press(selector, "Enter")
                    search_input_found = True
                    logger.info(f"Search executed using selector: {selector}")
                    break
                except Exception:
                    continue
            
            if not search_input_found:
                # Use vision to find search box
                logger.info("Using vision to find search box...")
                elements = await browser_agent.predict_interactive_elements(
                    screenshot_path,
                    "Find the main Google search input box where I can type a search query"
                )
                
                if elements and elements.get('elements'):
                    element = elements['elements'][0]
                    x, y = element['center']
                    await browser_agent.click(x, y)
                    await browser_agent.type_text(search_query)
                    await browser_agent.press_key("Enter")
                    logger.info("Search executed using vision detection")
                else:
                    raise Exception("Could not find search input box")
            
            # Wait for results to load
            await asyncio.sleep(3)
            
        except Exception as e:
            logger.error(f"Error performing search: {e}")
            return
        
        # Step 4: Extract search results
        logger.info("Step 4: Extracting search results...")
        try:
            # Take screenshot of results
            results_screenshot = "search_results.png"
            await browser_agent.take_screenshot(results_screenshot)
            logger.info(f"Results screenshot saved: {results_screenshot}")
            
            # Get page content
            page_content = await browser_agent.get_text("body")
            
            # Extract search result links (try common selectors)
            result_links = []
            link_selectors = [
                'a[href*="/url?q="]',  # Google result links
                'h3 a',                # Alternative structure
                'a[data-ved]'          # Another Google pattern
            ]
            
            for selector in link_selectors:
                try:
                    links = await browser_agent.get_attribute_all(selector, "href")
                    if links:
                        result_links.extend(links[:5])  # Get first 5 links
                        break
                except Exception:
                    continue
            
            # Clean up Google redirect URLs
            clean_links = []
            for link in result_links:
                if "/url?q=" in link:
                    # Extract actual URL from Google redirect
                    try:
                        actual_url = link.split("/url?q=")[1].split("&")[0]
                        from urllib.parse import unquote
                        clean_links.append(unquote(actual_url))
                    except Exception:
                        continue
                else:
                    clean_links.append(link)
            
            # Step 5: Compile and display results
            logger.info("Step 5: Compiling results...")
            
            print("\n" + "="*60)
            print("GOOGLE SEARCH RESULTS")
            print("="*60)
            print(f"Search Query: {search_query}")
            print(f"Results Found: {len(clean_links)}")
            print("\nTop Results:")
            
            for i, url in enumerate(clean_links[:5], 1):
                print(f"{i}. {url}")
            
            print(f"\nPage Content Preview (first 500 chars):")
            print("-" * 40)
            print(page_content[:500] + "..." if len(page_content) > 500 else page_content)
            print("="*60)
            
            logger.info("Search completed successfully!")
            
        except Exception as e:
            logger.error(f"Error extracting results: {e}")
        
    except Exception as e:
        logger.error(f"Error during search workflow: {e}", exc_info=True)
    
    finally:
        # Cleanup
        logger.info("Cleaning up...")
        await browser_agent.close()
        logger.info("Browser closed")


if __name__ == "__main__":
    print("Simple Google Search with BrowserAgent")
    print("=" * 40)
    print("This example demonstrates basic automated Google search")
    print("")
    print("Setup:")
    print("1. Set OPENAI_API_KEY environment variable, or")
    print("2. Create a .env file with: OPENAI_API_KEY=your_api_key_here")
    print("3. Optionally install: pip install python-dotenv")
    print("")
    
    # Run the example
    asyncio.run(main()) 