citation.py

from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException
from anthropic import Anthropic
import time
import requests
from bs4 import BeautifulSoup
import base64
from io import BytesIO
import os
from dotenv import load_dotenv
import tempfile
from trafilatura import fetch_url, extract
from twocaptcha import TwoCaptcha

class ImageCitationFinder:
    def __init__(self, show_browser=False):  # Changed parameter name to be more descriptive
        load_dotenv()
        self.api_key = os.getenv('CLAUDE_API_KEY')
        if not self.api_key:
            raise ValueError("CLAUDE_API_KEY not found in environment variables")
        
        self.client = Anthropic(api_key=self.api_key)
        self.show_browser = show_browser  # Store parameter for text extraction

        self.captcha_api_key = os.getenv('CAPTCHA_API_KEY')
        if not self.captcha_api_key:
            print("Warning: CAPTCHA_API_KEY not found in environment variables. Captcha solving will be disabled.")
        self.solver = TwoCaptcha(self.captcha_api_key) if self.captcha_api_key else None
        
        # Setup Selenium for image search (always headless)
        options = webdriver.ChromeOptions()
        options.add_experimental_option('excludeSwitches', ['enable-logging'])
        options.add_argument("--log-level=1")
        options.add_argument('--headless')  # Always headless for image search
        self.driver = webdriver.Chrome(options=options)
        
    def search_image(self, image_bytes):
        """Perform reverse image search using Google Images."""
        # print("Starting image search...")
        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
            tmp_file.write(image_bytes)
            tmp_path = tmp_file.name
        
        try:
            # Go to Google Images
            self.driver.get('https://images.google.com')
            
            # Click on camera icon
            camera_button = WebDriverWait(self.driver, 10).until(
                EC.presence_of_element_located((By.CSS_SELECTOR, "[aria-label='Search by image']"))
            )
            camera_button.click()
            
            # Upload image
            file_input = WebDriverWait(self.driver, 10).until(
                EC.presence_of_element_located((By.CSS_SELECTOR, "input[type='file']"))
            )
            file_input.send_keys(tmp_path)
            
            # Wait for results and click on "Exact matches" tab
            try:
                # print("Looking for 'Exact matches' tab...")
                exact_matches_tab = WebDriverWait(self.driver, 10).until(
                    EC.presence_of_element_located((By.XPATH, "//div[text()='Exact matches']"))
                )
                # Add a small wait to let any overlays clear
                time.sleep(1)
                
                # Try regular click first
                try:
                    WebDriverWait(self.driver, 5).until(
                        EC.element_to_be_clickable((By.XPATH, "//div[text()='Exact matches']"))
                    ).click()
                except Exception:
                    # If regular click fails, try JavaScript click
                    self.driver.execute_script("arguments[0].click();", exact_matches_tab)
                
                # Wait a bit longer for results after successful click
                time.sleep(2)
                
                # Get all links and filter
                all_links = self.driver.find_elements(By.TAG_NAME, "a")
                urls = []
                excluded_domains = ['google.co', 'gstatic.com', 'javascript:void(0)']
                
                for link in all_links:
                    url = link.get_attribute('href')
                    if url and not any(domain in url.lower() for domain in excluded_domains):
                        urls.append(url)
                
                # Get first 3 unique URLs
                unique_urls = list(dict.fromkeys(urls))[:3]  # Remove duplicates and get first 3
                
                # print(f"Found {len(unique_urls)} unique external links: {unique_urls}")
                return unique_urls
                
            except TimeoutException as e:
                print("Error occurred:", e)
                return []
            
        finally:
            os.unlink(tmp_path)

    def solve_captcha(self, url, driver):
        """
        Solve captcha using 2captcha service.
        Returns True if solved successfully, False otherwise.
        """
        if not self.solver:
            print("Captcha solver not initialized - missing API key")
            return False

        try:
            # Wait for page to fully load
            time.sleep(3)
            
            # Handle reCAPTCHA v2
            recaptcha_element = driver.find_elements(By.CLASS_NAME, "g-recaptcha")
            if recaptcha_element:
                print("Found reCAPTCHA v2")
                site_key = recaptcha_element[0].get_attribute("data-sitekey")
                result = self.solver.recaptcha(
                    sitekey=site_key,
                    url=url
                )
                
                # Wait for the response
                time.sleep(2)
                
                # Execute JavaScript to set the response
                driver.execute_script(
                    'document.querySelector("[name=g-recaptcha-response]").innerHTML = "{}";'.format(result["code"])
                )
                
                # Trigger the grecaptcha callback if it exists
                driver.execute_script(
                    'if (typeof ___grecaptcha_cfg !== "undefined") {'
                    '  Object.entries(___grecaptcha_cfg.clients).forEach(([k, v]) => {'
                    '    if (v["K"]) {'
                    '      v["K"].callback("{}");'.format(result["code"]) +
                    '    }'
                    '  });'
                    '}'
                )
                
                # Find and click the form submit button
                try:
                    submit_button = WebDriverWait(driver, 5).until(
                        EC.element_to_be_clickable((By.CSS_SELECTOR, "button[type='submit'], input[type='submit']"))
                    )
                    submit_button.click()
                except:
                    print("No submit button found, continuing...")
                
                return True

            # Handle reCAPTCHA v3
            recaptcha_v3 = driver.find_elements(By.XPATH, "//script[contains(@src, 'recaptcha/api.js?render=')]")
            if recaptcha_v3:
                print("Found reCAPTCHA v3")
                script_src = recaptcha_v3[0].get_attribute("src")
                site_key = script_src.split("render=")[1].split("&")[0]  # Handle additional URL parameters
                result = self.solver.recaptcha(
                    sitekey=site_key,
                    url=url,
                    version='v3',
                    action='submit'
                )
                
                # Wait for the response
                time.sleep(2)
                
                # Set the response token
                driver.execute_script(
                    f'document.querySelector("[name=g-recaptcha-response]").innerHTML = "{result["code"]}";'
                )
                return True

            # Handle hCaptcha
            hcaptcha_element = driver.find_elements(By.CLASS_NAME, "h-captcha")
            if hcaptcha_element:
                print("Found hCaptcha")
                site_key = hcaptcha_element[0].get_attribute("data-sitekey")
                result = self.solver.hcaptcha(
                    sitekey=site_key,
                    url=url
                )
                
                # Wait for the response
                time.sleep(2)
                
                # Set the response token
                driver.execute_script(
                    f'document.querySelector("[name=h-captcha-response]").innerHTML = "{result["code"]}";'
                )
                
                try:
                    submit_button = WebDriverWait(driver, 5).until(
                        EC.element_to_be_clickable((By.CSS_SELECTOR, "button[type='submit'], input[type='submit']"))
                    )
                    submit_button.click()
                except:
                    print("No submit button found, continuing...")
                
                return True

            print(f"No supported captcha type found on {url}")
            return False

        except Exception as e:
            print(f"Error solving captcha: {e}")
            return False


    def extract_text_from_url(self, url):
        """Extract main text content from a URL. Falls back to Selenium if direct request fails."""
        # Try direct request first
        try:
            response = requests.get(url)
            # print(url)
            # print(response.status_code)
            if response.status_code == 200:
                downloaded = response.text
                if downloaded:
                    return extract(downloaded)
        except Exception as e:
            print(f"Direct request failed: {e}")

        # Fall back to Selenium if direct request fails
        try:
            # Create a new Chrome driver - headless unless show_browser is True
            options = webdriver.ChromeOptions()
            options.add_experimental_option('excludeSwitches', ['enable-logging'])
            if not self.show_browser:  # Only add headless if show_browser is False
                options.add_argument('--headless')
            options.add_argument("--log-level=1")
            text_extraction_driver = webdriver.Chrome(options=options)
            
            try:
                text_extraction_driver.get(url)
                
                safe_filename = url.replace('://', '_').replace('/', '_')[:100]
                screenshot_path = f"C:\\Users\\hp\\copyright_handling\\error_ss\\{safe_filename}_{response.status_code}.png"
                text_extraction_driver.save_screenshot(screenshot_path)
                print(f"Screenshot saved as {screenshot_path}")

                # Check for captchas and solve if needed
                captcha_selectors = [
                    "//iframe[contains(@src, 'recaptcha')]",
                    "//div[contains(@class, 'captcha')]",
                    "//form[contains(@action, 'captcha')]"
                ]
                
                for selector in captcha_selectors:
                    try:
                        captcha_element = WebDriverWait(text_extraction_driver, 3).until(
                            EC.presence_of_element_located((By.XPATH, selector))
                        )
                        if captcha_element:
                            solved = self.solve_captcha(url)
                            print("solving captcha")
                            if not solved:
                                print(f"Captcha solving failed for {url}")
                                return ""
                            time.sleep(2)
                        break
                    except TimeoutException:
                        continue
                
                # Wait for body content to load
                wait = WebDriverWait(text_extraction_driver, 10)
                content = wait.until(EC.presence_of_element_located((By.TAG_NAME, "body")))
                page_text = content.text
                if page_text:
                    return page_text
                    
            finally:
                text_extraction_driver.quit()
                
        except Exception as e:
            print(f"Selenium extraction failed: {e}")
        
        return ""
    
    
    def generate_citation(self, image_bytes, extracted_texts):
        """Generate citation recommendation using Claude."""
        prompt = """Based on the following extracted texts from potential source websites of an image, 
        generate an appropriate citation or attribution text for the image. You must strictly follow one of these formats:

        1. © [Organization]. "[Paper/Article Title]" by [Author]. All rights reserved. This content is excluded from our Creative Commons license.
        2. Images courtesy of [Organization] from "[Paper/Article Title]" by [Author]. This image is in the public domain.
        3. Image courtesy of [Organization] on [Platform], from "[Paper/Article Title]" by [Author]. Used under [License].
        4. Source: [Organization], "[Paper/Article Title]" by [Author]. This image is in the public domain.
        5. Courtesy of [Author] on [Platform], "[Paper/Article Title]". Used with permission.

        Choose the most appropriate format and fill in all available information. If certain information (like paper title or author) 
        is not available, you may omit those parts while keeping the basic structure intact.

        If the provided texts do not contain relevant information about the image source, respond with: "Unable to determine image source information."

        Extracted texts:
        {texts}

        Respond only with the formatted citation or the unable to determine message, nothing else."""
        
        response = self.client.messages.create(
            model="claude-3-5-sonnet-latest",
            max_tokens=300,
            messages=[{
                "role": "user",
                "content": prompt.format(texts="\n\n".join(extracted_texts))
            }]
        )
        
        return response.content[0].text.strip()
    
    def get_citation_for_image(self, image_bytes):
        """Main method to get citation for an image."""
        try:
            # Search image on Google
            source_urls = self.search_image(image_bytes)
            
            # Extract text from each URL
            extracted_texts = []
            for url in source_urls:
                text = self.extract_text_from_url(url)
                if text:
                    extracted_texts.append(text)
            
            # Generate citation using Claude
            if extracted_texts:
                return self.generate_citation(image_bytes, extracted_texts)
            return "Source information could not be determined."
            
        except Exception as e:
            return f"Error finding citation: {str(e)}"
    
    def __del__(self):
        """Cleanup Selenium driver."""
        if hasattr(self, 'driver'):
            self.driver.quit()

def get_image_citation(pdf_image_data, show_browser=False):  # Changed parameter name
    finder = ImageCitationFinder(show_browser=show_browser)  # Changed parameter name
    urls = finder.search_image(pdf_image_data)
    
    # Extract text from each URL
    extracted_texts = []
    for url in urls:
        text = finder.extract_text_from_url(url)
        if text:
            extracted_texts.append(text)
    citation = finder.generate_citation(pdf_image_data, extracted_texts) if extracted_texts else "Source information could not be determined."
    
    return {
        'citation': citation,
        'source_urls': urls
    }