"""
Test script for Google Programmable Search Engine integration.
Run this script to verify that your Google PSE API key and search engine ID are working.
"""

import logging
import os
import random
import time

import requests
from dotenv import load_dotenv
from requests.exceptions import RequestException

from local_deep_research.web_search_engines.search_engine_factory import (
    create_search_engine,
)

# Set up logging
logging.basicConfig(
    level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# Load environment variables from .env file
load_dotenv()


def print_step(message):
    """Helper function to print and log a step with a timestamp"""
    print(f"[{time.strftime('%H:%M:%S')}] {message}")
    logger.info(message)


def check_api_quota(api_key, search_engine_id):
    """
    Make a direct minimal request to check API quota status
    Returns a tuple of (is_quota_ok, error_message)
    """
    url = "https://www.googleapis.com/customsearch/v1"
    params = {
        "key": api_key,
        "cx": search_engine_id,
        "q": "test",  # Minimal query
        "num": 1,  # Request only 1 result
    }

    try:
        response = requests.get(url, params=params, timeout=10)

        # Check for quota errors specifically
        if response.status_code == 429:
            return (
                False,
                "API quota exceeded. Google PSE has a limit of 100 requests per day on the free tier.",
            )
        elif response.status_code != 200:
            return False, f"API error: {response.status_code} - {response.text}"

        # If we get here, the API is working
        return True, None

    except Exception as e:
        return False, f"Error checking API: {str(e)}"


def test_google_pse_search(max_retries=3, retry_delay=2):
    """
    Test Google PSE search engine with retry logic and rate limiting

    Args:
        max_retries: Maximum number of retries on failure
        retry_delay: Base delay between retries in seconds
    """
    print_step("Starting Google Programmable Search Engine test...")

    # Check if API key and search engine ID are set
    api_key = os.getenv("GOOGLE_PSE_API_KEY")
    search_engine_id = os.getenv("GOOGLE_PSE_ENGINE_ID")

    # Mask the key for logging but show if it exists
    api_key_masked = (
        f"{api_key[:4]}...{api_key[-4:]}" if api_key and len(api_key) > 8 else None
    )
    print_step(f"API Key found: {api_key_masked is not None}")
    print_step(f"Search Engine ID: {search_engine_id}")

    if not api_key:
        print_step("❌ GOOGLE_PSE_API_KEY not set in environment variables")
        return False

    if not search_engine_id:
        print_step("❌ GOOGLE_PSE_ENGINE_ID not set in environment variables")
        return False

    print_step("✅ API key and search engine ID found")

    # Check API quota before proceeding
    print_step("Checking API quota status...")
    quota_ok, error_message = check_api_quota(api_key, search_engine_id)

    if not quota_ok:
        print_step(f"❌ {error_message}")
        print_step(
            "Please wait for quota to reset (typically after 24 hours) or use a different API key."
        )
        return False

    print_step("✅ API quota check passed")

    # Create search engine with reduced max_results to minimize API usage
    try:
        print_step("Creating search engine instance...")
        # Create search engine without LLM to avoid hanging or errors if Ollama is not running
        # Using None instead of get_llm() to skip LLM initialization
        engine = create_search_engine(
            "google_pse",
            llm=None,  # Skip LLM to avoid potential connection issues
            max_results=3,  # Reduced from 5 to minimize API usage
            region="us",
            safe_search=True,
            search_language="English",
        )

        if not engine:
            print_step("❌ Failed to create Google PSE search engine")
            return False

        print_step("✅ Search engine created successfully")

        # Run a simple test query with retry logic
        query = "artificial intelligence latest developments"
        print_step(f"Running test query: '{query}'")

        attempt = 0
        results = None
        rate_limited = False

        while attempt < max_retries:
            try:
                # Add a small random jitter to the delay to avoid thundering herd
                if attempt > 0:
                    jitter = random.uniform(0.5, 1.5)
                    sleep_time = retry_delay * (2 ** (attempt - 1)) * jitter
                    print_step(
                        f"Retry attempt {attempt}/{max_retries}. Waiting {sleep_time:.2f} seconds..."
                    )
                    time.sleep(sleep_time)

                # Execute the search
                print_step(
                    f"Executing search, attempt {attempt + 1} / {max_retries}..."
                )
                results = engine.run(query)

                # If successful, break out of the retry loop
                if results:
                    print_step("Search successful!")
                    break
                else:
                    print_step("Search returned no results")

            except RequestException as e:
                print_step(
                    f"Network error on attempt {attempt + 1} / {max_retries}:"
                    f" {str(e)}"
                )

                # Check for rate limiting
                if "429" in str(e):
                    rate_limited = True
                    print_step("⚠️ API quota has been exceeded")
                    break

            except Exception as e:
                print_step(
                    f"Error on attempt {attempt + 1} / {max_retries}:" f" {str(e)}"
                )

            attempt += 1

        # Special handling for rate limiting
        if rate_limited:
            print_step("❌ Google PSE API quota exceeded (HTTP 429)")
            print_step(
                "The free tier of Google PSE allows 100 requests per day. Please wait for quota to reset or use a different API key."
            )
            return False

        # Check results after all retry attempts
        if not results:
            print_step("❌ No results returned after all retry attempts")
            return False

        print_step(f"✅ {len(results)} results returned")

        # Print the first result
        print_step("\nFirst result:")
        first_result = results[0]
        print_step(f"Title: {first_result['title']}")
        print_step(f"URL: {first_result['url']}")
        print_step(f"Snippet: {first_result['snippet'][:100]}...")

        return True

    except Exception as e:
        print_step(f"❌ Error: {str(e)}")
        logger.exception("Exception in test_google_pse_search")
        return False


if __name__ == "__main__":
    try:
        print_step("Test script started")
        success = test_google_pse_search()
        if success:
            print_step(
                "\n✅ Google Programmable Search Engine test completed successfully"
            )
        else:
            print_step("\n❌ Google Programmable Search Engine test failed")
            print_step("\nPlease check that:")
            print_step(
                "1. You have set GOOGLE_PSE_API_KEY and GOOGLE_PSE_ENGINE_ID in .env file"
            )
            print_step("2. The API key is valid and has not reached its quota")
            print_step(
                "3. The search engine ID is correct and the search engine is properly configured"
            )
            print_step("4. Your network connection allows access to Google APIs")
    except Exception as e:
        logger.exception("Unhandled exception in main")
        print(f"Critical error: {str(e)}")
