components/google_grounder.py

"""
Google Search grounding for failed citation validation.
"""

from google import genai
from google.genai import types
import json
import logging
from typing import Dict, List, Any

logger = logging.getLogger(__name__)


class GoogleGrounder:
    """
    Grounds failed citations using Google Search.
    Provides AI-powered verification with web search for citations that couldn't be validated via APIs.
    """

    def __init__(self, client: genai.Client, model: str = "gemini-2.5-flash-lite"):
        """
        Initialize the Google grounder.

        Args:
            client: Google GenAI client instance
            model: Model name to use for grounding
        """
        self.client = client
        self.model = model

    def ground_with_search(self, api_results: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        Validates references using Google Gen AI with grounding via Google Search.
        This method extracts and displays Google Search grounding metadata.

        Args:
            api_results: The results from the API validators.

        Returns:
            A dictionary containing:
              - "results": The validated results from Gemini API.
              - "candidates_grounding_metadata": Google Search grounding metadata.
        """

        prompt = f""" **Objective:** Verify the existence and accuracy of academic references using the provided API search results and Google Search for cross-validation.

                **Input:** You will receive `api_results`, a JSON list containing potential matches found by APIs (like CrossRef, arXiv, OpenAlex). Each entry includes:
                *   `query`: The original reference details searched for (title, author, year, journal).
                *   `exists`: A boolean indicating if the API found a likely match based on its criteria.
                *   `source`: The API that found the match (e.g., 'crossref', 'arxiv').
                *   `matches`: A list of potential matches found by the API, ranked by similarity.
                *   `best_match`: The top match, including similarity scores (`overall_score`, `title_similarity`, etc.).

                **Verification Steps:**

                1.  **Analyze API Results:** For each reference entry in `api_results`:
                    *   If `exists` is `true` and the `best_match['overall_score']` is high (e.g., > 90%), consider the reference validated by the API (`source`). Note this in the explanation.
                    *   If `exists` is `false` OR the `best_match['overall_score']` is low/moderate, proceed to Google Search verification.

                2.  **Mandatory Google Search (Conditional):**
                    *   Use the Google Search tool ONLY if step 1 indicates verification is needed (i.e., `exists` is false or score is low).
                    *   Search for the reference using the `query` details (title, author, year).
                    *   Prioritize results from Google Scholar, ArXiv, official publisher websites, institutional repositories, and other credible academic sources.

                3.  **Determine Final Existence and Details:**
                    *   **Exact Match Found:** If Google Search finds an exact match (matching title, first author, year, journal if applicable) to the `query`, set `exists` to `true`. Provide the link found via search.
                    *   **Similar Match Found:** If Google Search finds a highly similar reference (e.g., minor title variation, slightly different year, correct author/topic) that seems to be the intended reference:
                        *   Set `exists` to `true`.
                        *   Provide the link to the found reference.
                        *   Clearly describe the discrepancies between the original `query` and the found reference in the `explanation` field.
                        *   List specific differences in the `issues` field (e.g., "Year mismatch: Query had 2020, found 2021 via Google Search", "Minor title variation noted").
                    *   **No Match Found:** If, after checking the API results and performing Google Search (if required), no credible source confirms the reference, set `exists` to `false`. Explain briefly why (e.g., "Not found via API search or Google Scholar/ArXiv").

                **Output Format:** Respond strictly in JSON format. For each reference processed, create a JSON object with the following fields:

                *   `title`: The title from the original `query`.
                *   `exists`: `true` or `false` based on your verification.
                *   `link`: The URL to the confirmed reference (if `exists` is `true`), otherwise `null`. Prefer DOIs or direct publication links.
                *   `explanation`: A brief summary of the verification process and findings (e.g., "Confirmed via Crossref", "Found similar reference on Google Scholar, year differs", "Reference not found").
                *   `first_author`: The first author from the original `query` (optional).
                *   `year`: The year from the original `query` (optional).
                *   `journal`: The journal from the original `query` (optional).
                *   `issues`: A list of strings describing discrepancies found (e.g., ["Year mismatch", "Title slightly different"]), or `null` if no issues.

                **Important:** Ensure the output is a valid JSON list containing one object for each reference processed from the input `api_results`. Output only the JSON.

                **API Results Input:**
                {json.dumps(api_results, indent=2)}
            """

        try:
            from google.genai import errors as genai_errors

            response = self.client.models.generate_content(
                model=self.model,
                contents=[prompt],
                config=types.GenerateContentConfig(
                    tools=[types.Tool(google_search=types.GoogleSearch())],
                    temperature=0
                    # Note: Cannot use response_mime_type with tools - Gemini API limitation
                )
            )

            # Extract JSON from response text (may have markdown code blocks)
            response_text = self._extract_text_from_response(response)

            if not response_text:
                logger.error("No text found in grounding response")
                return {"results": [], "candidates_grounding_metadata": []}

            # Parse JSON
            results = self._parse_json_response(response_text)

            # Extract grounding metadata
            grounding_metadata = self.extract_grounding_metadata(response)

            logger.info(f"Successfully grounded {len(results) if isinstance(results, list) else 1} references")

            return {
                "results": results,
                "candidates_grounding_metadata": grounding_metadata
            }

        except genai_errors.ClientError as e:
            # Handle rate limiting (429) specifically
            if e.status_code == 429:
                logger.error(f"Rate limit exceeded for Gemini API: {e}")
                # Extract retry delay if available
                retry_info = "Please try again later"
                if 'Please retry in' in str(e):
                    import re
                    match = re.search(r'retry in (\d+\.?\d*)s', str(e))
                    if match:
                        retry_seconds = int(float(match.group(1)))
                        retry_info = f"Please retry in {retry_seconds} seconds (~{retry_seconds//60} minutes)"

                return {
                    "results": [],
                    "candidates_grounding_metadata": [],
                    "error": "rate_limit_exceeded",
                    "message": f"API rate limit exceeded. {retry_info}. Upgrade your API plan for higher limits."
                }
            else:
                # Other API errors
                logger.error(f"Gemini API error {e.status_code}: {e}")
                return {
                    "results": [],
                    "candidates_grounding_metadata": [],
                    "error": "api_error",
                    "message": f"Gemini API error: {str(e)}"
                }

        except Exception as e:
            logger.exception(f"Unexpected error in Google Search grounding: {e}")
            return {"results": [], "candidates_grounding_metadata": []}

    def _extract_text_from_response(self, response) -> str:
        """Extract text from GenAI response, handling multiple parts."""
        try:
            # Check if response has expected structure
            if not response.candidates or not response.candidates[0].content:
                logger.error("Response has no candidates or content")
                return ""

            # Get the first part that has text
            for part in response.candidates[0].content.parts:
                if hasattr(part, 'text') and part.text:
                    return part.text

            return ""

        except Exception as e:
            logger.error(f"Error extracting text from response: {e}")
            return ""

    def _parse_json_response(self, response_text: str) -> List[Dict[str, Any]]:
        """Parse JSON from response text, handling markdown code blocks."""
        try:
            # Remove markdown code blocks if present
            if response_text.strip().startswith('```'):
                lines = response_text.strip().split('\n')
                # Skip first line (```json) and last line (```)
                json_lines = [line for line in lines[1:-1]]
                response_text = '\n'.join(json_lines)

            results = json.loads(response_text)
            return results if isinstance(results, list) else [results]

        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse grounded GenAI response as JSON: {e}")
            logger.error(f"Response text: {response_text[:500] if response_text else 'None'}...")
            return []

    def extract_grounding_metadata(self, response: types.GenerateContentResponse) -> List[Dict[str, Any]]:
        """
        Extracts grounding metadata from the GenAI response.

        Args:
            response: The GenerateContentResponse from GenAI

        Returns:
            List of grounding metadata dictionaries
        """
        grounding_metadata = []

        try:
            for candidate_idx, candidate in enumerate(response.candidates):
                if not hasattr(candidate, 'grounding_metadata') or not candidate.grounding_metadata:
                    continue

                candidate_metadata = {
                    "candidate_index": candidate_idx,
                    "generated_text": candidate.content.parts[0].text if candidate.content.parts else "",
                    "search_entry_point": None,
                    "grounding_chunks": [],
                    "grounding_supports": [],
                    "web_search_queries": []
                }

                grounding = candidate.grounding_metadata

                # Extract search entry point
                if hasattr(grounding, 'search_entry_point') and grounding.search_entry_point:
                    candidate_metadata["search_entry_point"] = {
                        "rendered_content": grounding.search_entry_point.rendered_content
                    }

                # Extract grounding chunks
                if hasattr(grounding, 'grounding_chunks') and grounding.grounding_chunks:
                    for chunk in grounding.grounding_chunks:
                        chunk_data = {}
                        if hasattr(chunk, 'web') and chunk.web:
                            chunk_data = {
                                "uri": chunk.web.uri,
                                "title": chunk.web.title
                            }
                        candidate_metadata["grounding_chunks"].append(chunk_data)

                # Extract grounding supports
                if hasattr(grounding, 'grounding_supports') and grounding.grounding_supports:
                    for support in grounding.grounding_supports:
                        support_data = {
                            "segment": {
                                "start_index": support.segment.start_index if hasattr(support, 'segment') and support.segment else None,
                                "end_index": support.segment.end_index if hasattr(support, 'segment') and support.segment else None,
                                "text": support.segment.text if hasattr(support, 'segment') and support.segment else None
                            },
                            "grounding_chunk_indices": support.grounding_chunk_indices if hasattr(support, 'grounding_chunk_indices') else [],
                            "confidence_scores": support.confidence_scores if hasattr(support, 'confidence_scores') else []
                        }
                        candidate_metadata["grounding_supports"].append(support_data)

                # Extract web search queries
                if hasattr(grounding, 'web_search_queries') and grounding.web_search_queries:
                    candidate_metadata["web_search_queries"] = grounding.web_search_queries

                grounding_metadata.append(candidate_metadata)

        except Exception as e:
            logger.error(f"Error extracting grounding metadata: {e}")

        return grounding_metadata