AutoAPI Gemini Rules

"""
Gemini-powered webpage analyzer that generates scraper recipes.
"""
import os
import json
from typing import Any, Optional
from io import BytesIO

from dotenv import load_dotenv
from google import genai
from google.genai import types
from PIL import Image
from playwright.async_api import async_playwright

from .models import ScraperRecipe
from .executor import RecipeExecutor
from .explorer import explore_page


SYSTEM_PROMPT = """You are a web scraping expert. Your job is to analyze a webpage and generate a JSON recipe for extracting data.

You will receive:
1. A screenshot of the webpage
2. The HTML source
3. A user prompt describing what data to extract

Your task is to generate a ScraperRecipe JSON that will extract the requested data.

Available action types:
- goto: Navigate to a URL
- wait_for: Wait for a CSS selector to appear
- click: Click an element (useful for "load more", modals, etc.)
- scroll: Scroll the page (directions: down, up, bottom)
- extract_single: Extract one item with multiple fields from the full page
- extract_many: Extract multiple items from repeated container elements

Field extraction types:
- text: Get inner text of the element
- attribute: Get an attribute value (specify "attribute" field)
- exists: Boolean, whether element exists
- html: Get inner HTML
- nested: Extract a nested object with its own fields
- nested_many: Extract an array of nested objects

Field transformations:
- number: Extract numbers from text (e.g. "$12.99" -> 12.99)
- boolean: Convert to boolean
- trim: Remove whitespace
- lowercase: Convert to lowercase

NESTED EXTRACTION:
For nested objects, use type "nested" with:
- selector: CSS selector for the parent element
- fields: Object with nested field configurations

For nested arrays, use type "nested_many" with:
- selector: CSS selector for the parent element
- container: CSS selector for repeated items (relative to parent)
- fields: Object with nested field configurations

Example with nested data:
{
  "title": {"selector": "h1", "type": "text"},
  "author": {
    "selector": ".author-section",
    "type": "nested",
    "fields": {
      "name": {"selector": ".name", "type": "text"},
      "bio": {"selector": ".bio", "type": "text"}
    }
  },
  "comments": {
    "selector": ".comments-section",
    "type": "nested_many",
    "container": ".comment",
    "fields": {
      "author": {"selector": ".comment-author", "type": "text"},
      "text": {"selector": ".comment-text", "type": "text"}
    }
  }
}

CRITICAL RULES:
- ONLY use CSS selectors. NEVER use XPath. Our executor does not support XPath.
- All selectors must be valid CSS selectors (e.g. ".class", "#id", "tag", "[attr=val]", "tag.class > child")
- For extract_many, the "container" selector must select elements that CONTAIN all the fields you want to extract
- Field selectors inside extract_many are RELATIVE to the container element
- For nested fields, selectors are RELATIVE to their parent element
- Use stable selectors (prefer classes, IDs, data attributes over positional selectors)
- Always start with a goto action
- Add wait_for actions before extraction to ensure content is loaded
- The "output_schema" can now include nested objects and arrays (e.g. {"title": "string", "author": {"name": "string"}, "tags": ["string"]})

Example recipe for extracting products with nested data:
{
  "url": "https://example.com/products",
  "output_schema": {
    "title": "string",
    "price": "number",
    "image": "string",
    "seller": {
      "name": "string",
      "rating": "number"
    },
    "reviews": [
      {
        "author": "string",
        "text": "string",
        "rating": "number"
      }
    ]
  },
  "steps": [
    {
      "action": "goto",
      "url": "https://example.com/products"
    },
    {
      "action": "wait_for",
      "selector": ".product-grid"
    },
    {
      "action": "extract_many",
      "container": ".product-card",
      "fields": {
        "title": {
          "selector": "h3",
          "type": "text"
        },
        "price": {
          "selector": ".price",
          "type": "text",
          "transform": "number"
        },
        "image": {
          "selector": "img",
          "type": "attribute",
          "attribute": "src"
        },
        "seller": {
          "selector": ".seller-info",
          "type": "nested",
          "fields": {
            "name": {"selector": ".seller-name", "type": "text"},
            "rating": {"selector": ".seller-rating", "type": "text", "transform": "number"}
          }
        },
        "reviews": {
          "selector": ".reviews-section",
          "type": "nested_many",
          "container": ".review",
          "fields": {
            "author": {"selector": ".review-author", "type": "text"},
            "text": {"selector": ".review-text", "type": "text"},
            "rating": {"selector": ".review-rating", "type": "text", "transform": "number"}
          }
        }
      }
    }
  ],
  "pagination": {
    "next_button": ".pagination .next",
    "max_pages": 3
  }
}

Generate ONLY the JSON recipe, no other text."""


RETRY_PROMPT = """The recipe you generated had errors during execution. Here is the recipe you generated:

{recipe_json}

Here are the errors that occurred:
{errors}

Here are some sample results (note the null values where extraction failed):
{sample_results}

Please fix the recipe. Remember:
- ONLY use CSS selectors. NEVER use XPath.
- All selectors must be valid CSS (e.g. ".class", "#id", "tag > child", "[attr=val]")
- If a field's data is in a sibling element outside the container, use a broader container that wraps both, or restructure the extraction approach
- Check the HTML source carefully for the correct selectors

Generate the corrected JSON recipe."""


class WebpageAnalyzer:
    """Analyzes webpages and generates scraper recipes using Gemini."""

    def __init__(self, api_key: Optional[str] = None):
        load_dotenv()
        self.api_key = api_key or os.getenv("GEMINI_API_KEY")
        if not self.api_key:
            raise ValueError("GEMINI_API_KEY not set")

        self.client = genai.Client(api_key=self.api_key)

    async def capture_page(self, url: str) -> tuple[bytes, str]:
        """Capture screenshot and HTML from a URL.
        
        DEPRECATED: This method is kept for backward compatibility but is no longer
        used in the main flow. The explorer handles page capture now.
        """
        async with async_playwright() as p:
            browser = await p.chromium.launch(headless=True)
            page = await browser.new_page(viewport={"width": 1280, "height": 800})

            # Use domcontentloaded instead of networkidle for speed
            await page.goto(url, wait_until="domcontentloaded", timeout=15000)

            # Give a brief moment for any immediate JS to run
            await page.wait_for_timeout(1000)

            screenshot_bytes = await page.screenshot(full_page=False)
            html = await page.content()

            await browser.close()

            return screenshot_bytes, html

    def _call_gemini(
        self,
        user_prompt: str,
        screenshot: bytes,
    ) -> ScraperRecipe:
        """Call Gemini and parse the response into a ScraperRecipe."""
        response = self.client.models.generate_content(
            model="gemini-3-flash-preview",
            contents=[
                types.Content(
                    role="user",
                    parts=[
                        types.Part(text=user_prompt),
                        types.Part(inline_data=types.Blob(
                            mime_type="image/png",
                            data=screenshot
                        ))
                    ]
                )
            ],
            config=types.GenerateContentConfig(
                system_instruction=SYSTEM_PROMPT,
                temperature=0.1,
                response_mime_type="application/json",
            )
        )

        recipe_dict = json.loads(response.text)
        return ScraperRecipe(**recipe_dict)

    def generate_recipe(
        self,
        url: str,
        prompt: str,
        screenshot: bytes,
        html: str,
        desired_schema: Optional[dict[str, str]] = None,
        exploration_context: Optional[dict] = None
    ) -> ScraperRecipe:
        """Generate a scraper recipe from page analysis.
        
        Args:
            exploration_context: Optional dict with 'screenshots' and 'observations' from exploration
        """
        user_prompt = f"URL: {url}\n\nUser request: {prompt}\n\n"

        if desired_schema:
            user_prompt += f"Desired schema: {json.dumps(desired_schema, indent=2)}\n\n"

        # Add exploration observations if available
        if exploration_context and exploration_context.get("observations"):
            user_prompt += "Exploration observations:\n"
            for obs in exploration_context["observations"]:
                user_prompt += f"- {obs}\n"
            user_prompt += "\n"

        user_prompt += f"HTML (first 15000 chars):\n{html[:15000]}\n\n"
        user_prompt += "Generate a ScraperRecipe JSON to extract the requested data."

        # Use the first screenshot from exploration if available, otherwise use provided
        screenshot_to_use = screenshot
        if exploration_context and exploration_context.get("screenshots"):
            screenshot_to_use = exploration_context["screenshots"][0]

        return self._call_gemini(user_prompt, screenshot_to_use)

    def fix_recipe(
        self,
        recipe: ScraperRecipe,
        errors: list[str],
        results: list[dict],
        screenshot: bytes,
    ) -> ScraperRecipe:
        """Send errors back to Gemini and get a fixed recipe."""
        sample = results[:3] if results else []
        user_prompt = RETRY_PROMPT.format(
            recipe_json=json.dumps(recipe.model_dump(), indent=2),
            errors="\n".join(f"- {e}" for e in errors),
            sample_results=json.dumps(sample, indent=2),
        )

        return self._call_gemini(user_prompt, screenshot)

    async def analyze_url(
        self,
        url: str,
        prompt: str,
        desired_schema: Optional[dict[str, str]] = None,
        max_retries: int = 2,
        on_status: Any = None,
        broadcast: Any = None,
    ) -> tuple[ScraperRecipe, dict]:
        """Full pipeline: explore, analyze, execute, retry on errors.

        Args:
            on_status: Optional callback(message: str) for progress updates.
            broadcast: Optional async callback for streaming events to WebSocket.

        Returns:
            Tuple of (final_recipe, execution_result)
        """
        def status(msg: str):
            import logging
            logger = logging.getLogger("uvicorn")
            logger.info(f"[autoAPI] {msg}")
            if on_status:
                on_status(msg)

        # Phase 1: Exploration
        status("Exploring page with AI agent...")
        page, browser, exploration_context = await explore_page(
            url=url,
            gemini_client=self.client,
            max_steps=5,
            broadcast=broadcast
        )
        
        # Get HTML from the explored page
        html = await page.content()
        
        # Get a final screenshot after exploration
        final_screenshot = await page.screenshot(full_page=False)
        exploration_context["screenshots"].append(final_screenshot)
        
        status(f"✓ Exploration complete with {len(exploration_context['observations'])} observations")

        # Phase 2: Recipe Generation
        if broadcast:
            await broadcast({
                "type": "status",
                "message": "Generating scraper recipe from exploration...",
                "phase": "generate"
            })
        
        status("Generating scraper recipe...")
        recipe = self.generate_recipe(
            url, prompt, final_screenshot, html, desired_schema, exploration_context
        )
        status(f"✓ Generated recipe with {len(recipe.steps)} steps")
        
        if broadcast:
            await broadcast({
                "type": "recipe",
                "recipe": recipe.model_dump(),
                "phase": "generate"
            })

        # Phase 3: Execution with retries
        for attempt in range(max_retries):
            if broadcast:
                await broadcast({
                    "type": "status",
                    "message": "Executing recipe...",
                    "phase": "execute"
                })
            
            status("Executing recipe...")

            async with RecipeExecutor(page=page, broadcast=broadcast) as executor:
                result = await executor.execute_recipe(recipe)

            errors = result.get("errors", [])
            null_count = sum(
                1 for item in result["data"]
                for v in item.values()
                if v is None
            )

            if not errors and null_count == 0:
                status(f"✓ Recipe executed successfully! Extracted {len(result['data'])} items")
                
                # Clean up browser
                await browser.close()
                if hasattr(browser, '_playwright'):
                    await browser._playwright.stop()
                
                return recipe, result

            if attempt < max_retries - 1:
                status(
                    f"⚠ Got {len(errors)} error(s) and {null_count} null value(s). "
                    f"Asking Gemini to regenerate recipe..."
                )
                # Use the final screenshot for retry
                recipe = self.fix_recipe(recipe, errors, result["data"], final_screenshot)
                status(f"✓ Received updated recipe")
            else:
                status(
                    f"⚠ Returning result with {len(errors)} error(s) and {null_count} null(s)."
                )

        # Clean up browser
        await browser.close()
        if hasattr(browser, '_playwright'):
            await browser._playwright.stop()

        return recipe, result


async def analyze_url(
    url: str,
    prompt: str,
    desired_schema: Optional[dict[str, str]] = None
) -> tuple[ScraperRecipe, dict]:
    """Convenience function to analyze a URL."""
    analyzer = WebpageAnalyzer()
    return await analyzer.analyze_url(url, prompt, desired_schema)