#!/usr/bin/env python
"""Cookie-consent state detection via Gemini Flash Lite.

The crawler uses this module as a verification helper, not as the primary page
parser. It answers two questions from a viewport screenshot:

1. Is a cookie-consent window/banner currently visible?
2. If yes, where is the primary accept-all button and what visible text does it
   show?

The returned button coordinates are normalized to the screenshot viewport
coordinate system. The crawler can first try a DOM text click with the returned
text and then fall back to nodriver's viewport pixel click.
"""

from __future__ import annotations

import json
import logging
import os
import sys
import time
from pathlib import Path
from typing import Any, Sequence

from google import genai
from google.genai import types

log = logging.getLogger("crawler.gemini")

_MODEL = "models/gemini-3.1-flash-lite-preview"

_CONSENT_SCHEMA: dict[str, Any] = {
    "type": "object",
    "additionalProperties": False,
    "required": [
        "cookie_consent_shown",
        "accept_button_visible",
        "accept_button_text",
        "accept_button_box_2d",
        "recommended_action",
        "confidence",
        "reason",
    ],
    "properties": {
        "cookie_consent_shown": {"type": "boolean"},
        "accept_button_visible": {"type": "boolean"},
        "accept_button_text": {"type": ["string", "null"]},
        "accept_button_box_2d": {
            "type": ["array", "null"],
            "items": {"type": "number"},
            "minItems": 4,
            "maxItems": 4,
        },
        "recommended_action": {
            "type": "string",
            "enum": ["continue", "click_accept_button", "manual_review"],
        },
        "confidence": {"type": "number", "minimum": 0, "maximum": 1},
        "reason": {"type": "string"},
    },
}

_CHALLENGE_SCHEMA: dict[str, Any] = {
    "type": "object",
    "additionalProperties": False,
    "required": ["is_challenge", "provider", "confidence", "visible_text"],
    "properties": {
        "is_challenge": {"type": "boolean"},
        "provider": {
            "type": "string",
            "enum": ["cloudflare", "akamai", "other", "none"],
        },
        "confidence": {"type": "number", "minimum": 0, "maximum": 1},
        "visible_text": {"type": "string"},
    },
}

_CHALLENGE_PROMPT = """
You inspect a webpage VIEWPORT screenshot for bot-protection or security
challenge pages.

Answer whether the screenshot shows a Cloudflare, Akamai, or similar
bot-protection/security verification challenge instead of normal site content.
Look for text or branding such as: Verify you are human, Security verification,
Performing security verification, Checking your browser, Cloudflare, Akamai,
bot protection, human verification, CAPTCHA, or a security service page.

Rules:
- Use the screenshot pixels as source of truth.
- A normal page with a small Cloudflare footer link is not a challenge.
- A normal cookie-consent banner is not a bot-protection challenge.
- If the page is normal content, is_challenge must be false, provider must be
  none, and confidence should reflect confidence that it is not blocked.
- If the provider is unclear but the page is clearly a security challenge, use
  provider other.
- visible_text should contain the most relevant visible security/challenge text.
- Return only JSON matching the schema.
""".strip()

_PROMPT = """
You inspect a webpage VIEWPORT screenshot for cookie consent.

Answer whether a cookie-consent banner, modal, overlay, or privacy settings
window is visible. If a consent window is visible, identify the primary button
that accepts all cookies or dismisses the banner with consent. Prefer buttons
with visible labels like:

English: Accept All Cookies, Accept All, Accept, Allow All, Agree, I Agree,
OK, Got It, Continue, I Understand.
German: Alle Akzeptieren, Alle Cookies Akzeptieren, Akzeptieren, Zulassen,
Alle Zulassen, Zustimmen, Einverstanden, OK, Verstanden, Cookies Akzeptieren,
Erlauben, Allen Zustimmen, Ich Stimme Zu.

Also recognize cookie/privacy consent context when the popup or banner contains
text like: Datenschutzerklärung, Datenschutzeinstellungen, Datenschutz,
Cookie-Richtlinie, Cookie-Einstellungen, Privatsphäre, Wir verwenden Cookies,
Diese Website verwendet Cookies, Privacy Policy, Cookie Policy.
These indicate a consent UI even if the button text differs from above.

Rules:
- Use the screenshot pixels as source of truth.
- Do not report normal page CTAs as cookie buttons.
- If no cookie consent UI is visible, cookie_consent_shown must be false,
  accept_button_visible must be false, accept_button_text and
  accept_button_box_2d must be null, and recommended_action must be continue.
- If cookie consent is visible but no accept-all button is readable,
  recommended_action must be manual_review.
- accept_button_box_2d must be [x0, y0, x1, y1] normalized to 0..1000 relative
  to the current screenshot viewport, top-left origin.
- The box must cover the clickable button area, not only the text.
- Return only JSON matching the schema.
""".strip()


def _parse_json_response(text: str) -> dict[str, Any]:
    text = (text or "").strip()
    if text.startswith("```"):
        lines = text.splitlines()
        if lines and lines[0].startswith("```"):
            lines = lines[1:]
        while lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        text = "\n".join(lines).strip()
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        start = text.find("{")
        end = text.rfind("}")
        if start >= 0 and end > start:
            return json.loads(text[start:end + 1])
        raise


def _clamp_confidence(value: Any) -> float:
    try:
        return max(0.0, min(1.0, float(value)))
    except Exception:
        return 0.0


def _sanitize_challenge_payload(data: dict[str, Any]) -> dict[str, Any]:
    provider = str(data.get("provider") or "none").strip().lower()
    if provider not in {"cloudflare", "akamai", "other", "none"}:
        provider = "other" if provider else "none"

    is_challenge = bool(data.get("is_challenge", False))
    if not is_challenge:
        provider = "none"

    visible_text = str(data.get("visible_text") or "").strip()
    if len(visible_text) > 2000:
        visible_text = visible_text[:2000]

    return {
        "is_challenge": is_challenge,
        "provider": provider,
        "confidence": _clamp_confidence(data.get("confidence")),
        "visible_text": visible_text,
    }


def _sanitize_consent_payload(data: dict[str, Any]) -> dict[str, Any]:
    shown = bool(data.get("cookie_consent_shown", False))
    button_visible = bool(data.get("accept_button_visible", False))
    text = data.get("accept_button_text")
    text = str(text).strip() if text is not None and str(text).strip() else None

    box = data.get("accept_button_box_2d")
    clean_box: list[float] | None = None
    if isinstance(box, list) and len(box) == 4:
        try:
            values = [max(0.0, min(1000.0, float(v))) for v in box]
            x0, y0, x1, y1 = values
            if x1 > x0 and y1 > y0:
                clean_box = values
        except Exception:
            clean_box = None

    action = str(data.get("recommended_action") or "").strip()
    if not shown:
        return {
            "cookie_consent_shown": False,
            "accept_button_visible": False,
            "accept_button_text": None,
            "accept_button_box_2d": None,
            "recommended_action": "continue",
            "confidence": float(data.get("confidence") or 0.0),
            "reason": str(data.get("reason") or "No cookie consent UI visible."),
        }

    if action not in {"continue", "click_accept_button", "manual_review"}:
        action = "click_accept_button" if button_visible and (text or clean_box) else "manual_review"

    return {
        "cookie_consent_shown": True,
        "accept_button_visible": button_visible,
        "accept_button_text": text,
        "accept_button_box_2d": clean_box,
        "recommended_action": action,
        "confidence": float(data.get("confidence") or 0.0),
        "reason": str(data.get("reason") or "Cookie consent UI visible."),
    }


def analyze_cookie_consent(
    screenshot_path: str | Path,
    *,
    model: str = _MODEL,
    api_key: str | None = None,
    verbose: bool = False,
    max_attempts: int = 3,
) -> dict[str, Any]:
    """Return cookie-consent state for one viewport screenshot."""
    api_key = api_key or os.getenv("GEMINI_API_KEY", "")
    if not api_key:
        raise RuntimeError("GEMINI_API_KEY environment variable not set")

    screenshot_path = Path(screenshot_path)
    if not screenshot_path.exists():
        raise FileNotFoundError(f"Screenshot not found: {screenshot_path}")

    client = genai.Client(api_key=api_key)
    image_bytes = screenshot_path.read_bytes()
    mime_type = "image/png"
    if screenshot_path.suffix.lower() in {".jpg", ".jpeg"}:
        mime_type = "image/jpeg"
    elif screenshot_path.suffix.lower() == ".webp":
        mime_type = "image/webp"

    config = types.GenerateContentConfig(
        temperature=0.0,
        response_mime_type="application/json",
        response_json_schema=_CONSENT_SCHEMA,
    )

    delay_seconds = 1.5
    for attempt in range(1, max_attempts + 1):
        try:
            response = client.models.generate_content(
                model=model,
                contents=[
                    _PROMPT,
                    types.Part.from_bytes(data=image_bytes, mime_type=mime_type),
                ],
                config=config,
            )
            raw = response.text or ""
            if verbose:
                log.debug("Gemini consent raw response: %s", raw[:500])
            data = _parse_json_response(raw)
            return _sanitize_consent_payload(data)
        except Exception as exc:
            if attempt >= max_attempts:
                raise
            log.warning("Gemini consent check attempt %d failed: %s", attempt, exc)
            time.sleep(delay_seconds)
            delay_seconds = min(delay_seconds * 2, 10.0)

    raise RuntimeError("Unreachable Gemini consent retry path")


def detect_cloudflare_challenge(
    screenshot_path: str | Path,
    *,
    model: str = _MODEL,
    api_key: str | None = None,
    verbose: bool = False,
    max_attempts: int = 3,
) -> dict[str, Any]:
    """Return bot-protection challenge state for one viewport screenshot.

    The function name keeps the current crawler terminology, but the detector is
    intentionally generic and can return cloudflare, akamai, other, or none.
    """
    api_key = api_key or os.getenv("GEMINI_API_KEY", "")
    if not api_key:
        raise RuntimeError("GEMINI_API_KEY environment variable not set")

    screenshot_path = Path(screenshot_path)
    if not screenshot_path.exists():
        raise FileNotFoundError(f"Screenshot not found: {screenshot_path}")

    client = genai.Client(api_key=api_key)
    image_bytes = screenshot_path.read_bytes()
    mime_type = "image/png"
    if screenshot_path.suffix.lower() in {".jpg", ".jpeg"}:
        mime_type = "image/jpeg"
    elif screenshot_path.suffix.lower() == ".webp":
        mime_type = "image/webp"

    config = types.GenerateContentConfig(
        temperature=0.0,
        response_mime_type="application/json",
        response_json_schema=_CHALLENGE_SCHEMA,
    )

    delay_seconds = 1.5
    for attempt in range(1, max_attempts + 1):
        try:
            response = client.models.generate_content(
                model=model,
                contents=[
                    _CHALLENGE_PROMPT,
                    types.Part.from_bytes(data=image_bytes, mime_type=mime_type),
                ],
                config=config,
            )
            raw = response.text or ""
            if verbose:
                log.debug("Gemini challenge raw response: %s", raw[:500])
            data = _parse_json_response(raw)
            return _sanitize_challenge_payload(data)
        except Exception as exc:
            if attempt >= max_attempts:
                raise
            log.warning("Gemini challenge check attempt %d failed: %s", attempt, exc)
            time.sleep(delay_seconds)
            delay_seconds = min(delay_seconds * 2, 10.0)

    raise RuntimeError("Unreachable Gemini challenge retry path")


_OVERLAY_SCHEMA: dict[str, Any] = {
    "type": "object",
    "additionalProperties": False,
    "required": [
        "has_blocking_overlay",
        "close_button_visible",
        "close_button_box_2d",
        "overlay_coverage",
        "confidence",
    ],
    "properties": {
        "has_blocking_overlay": {"type": "boolean"},
        "close_button_visible": {"type": "boolean"},
        "close_button_box_2d": {
            "type": ["array", "null"],
            "items": {"type": "number"},
            "minItems": 4,
            "maxItems": 4,
        },
        "overlay_coverage": {
            "type": "string",
            "enum": ["full", "partial", "none"],
        },
        "confidence": {"type": "number", "minimum": 0, "maximum": 1},
    },
}

_OVERLAY_PROMPT = """
You inspect a webpage VIEWPORT screenshot for blocking overlays or modals.

Look at the screenshot purely VISUALLY. Determine whether a popup, modal,
overlay, or floating panel is sitting on top of the main page content and
preventing the user from reading or interacting with the page below.

Blocking overlays typically look like:
- A white or colored box centered on screen with a dimmed/darkened/blurred
  background behind it.
- A slide-in panel from the side or bottom covering page content.
- A fullscreen takeover hiding the real page.

These overlays often (but not always) have a dismissal element. The dismissal
element can take many forms:
- A small "X" or "x" icon, typically at the top-right corner of the overlay.
- A circle with an X, or a similar icon-based close affordance.
- A text button or link such as: Close, Dismiss, Ignore, No thanks, Skip,
  Cancel, Not now, Maybe later, Schließen, Abbrechen, Nein danke, Ignorieren,
  Nicht jetzt, Später, Überspringen.
- A secondary/less-prominent link below or beside the main call-to-action.
- The close element may appear at top-right, top-left, bottom-center, or
  outside the overlay box.

Rules:
- Use only the visual pixels as source of truth.
- Do NOT rely on text content to classify the overlay. Any popup that blocks
  the page counts, regardless of what it says.
- Cookie consent banners at the top or bottom that do NOT dim/block the full
  page are NOT blocking overlays.
- A small chat widget or floating action button is NOT a blocking overlay.
- close_button_box_2d must be [x0, y0, x1, y1] normalized to 0..1000 relative
  to the screenshot viewport, top-left origin.
- The box must tightly cover the clickable close icon area.
- If no close icon is visible, close_button_visible must be false and
  close_button_box_2d must be null.
- overlay_coverage: full = content completely hidden, partial = some content
  visible around/behind the overlay, none = no overlay.
- Return only JSON matching the schema.
""".strip()


def _sanitize_overlay_payload(data: dict[str, Any]) -> dict[str, Any]:
    has_overlay = bool(data.get("has_blocking_overlay", False))
    close_visible = bool(data.get("close_button_visible", False))

    box = data.get("close_button_box_2d")
    clean_box: list[float] | None = None
    if isinstance(box, list) and len(box) == 4:
        try:
            values = [max(0.0, min(1000.0, float(v))) for v in box]
            x0, y0, x1, y1 = values
            if x1 > x0 and y1 > y0:
                clean_box = values
        except Exception:
            clean_box = None

    if not has_overlay:
        return {
            "has_blocking_overlay": False,
            "close_button_visible": False,
            "close_button_box_2d": None,
            "overlay_coverage": "none",
            "confidence": _clamp_confidence(data.get("confidence")),
        }

    coverage = str(data.get("overlay_coverage") or "partial").strip().lower()
    if coverage not in {"full", "partial", "none"}:
        coverage = "partial"

    return {
        "has_blocking_overlay": True,
        "close_button_visible": close_visible and clean_box is not None,
        "close_button_box_2d": clean_box if close_visible else None,
        "overlay_coverage": coverage,
        "confidence": _clamp_confidence(data.get("confidence")),
    }


def detect_blocking_overlay(
    screenshot_path: str | Path,
    *,
    model: str = _MODEL,
    api_key: str | None = None,
    verbose: bool = False,
    max_attempts: int = 3,
) -> dict[str, Any]:
    """Detect blocking overlays/modals and locate the close button visually."""
    api_key = api_key or os.getenv("GEMINI_API_KEY", "")
    if not api_key:
        raise RuntimeError("GEMINI_API_KEY environment variable not set")

    screenshot_path = Path(screenshot_path)
    if not screenshot_path.exists():
        raise FileNotFoundError(f"Screenshot not found: {screenshot_path}")

    client = genai.Client(api_key=api_key)
    image_bytes = screenshot_path.read_bytes()
    mime_type = "image/png"
    if screenshot_path.suffix.lower() in {".jpg", ".jpeg"}:
        mime_type = "image/jpeg"
    elif screenshot_path.suffix.lower() == ".webp":
        mime_type = "image/webp"

    config = types.GenerateContentConfig(
        temperature=0.0,
        response_mime_type="application/json",
        response_json_schema=_OVERLAY_SCHEMA,
    )

    delay_seconds = 1.5
    for attempt in range(1, max_attempts + 1):
        try:
            response = client.models.generate_content(
                model=model,
                contents=[
                    _OVERLAY_PROMPT,
                    types.Part.from_bytes(data=image_bytes, mime_type=mime_type),
                ],
                config=config,
            )
            raw = response.text or ""
            if verbose:
                log.debug("Gemini overlay raw response: %s", raw[:500])
            data = _parse_json_response(raw)
            return _sanitize_overlay_payload(data)
        except Exception as exc:
            if attempt >= max_attempts:
                raise
            log.warning("Gemini overlay check attempt %d failed: %s", attempt, exc)
            time.sleep(delay_seconds)
            delay_seconds = min(delay_seconds * 2, 10.0)

    raise RuntimeError("Unreachable Gemini overlay retry path")


# Backwards-compatible helper name used by older crawler versions.
def detect_consent_button(
    screenshot_paths: Sequence[str | Path],
    viewport_width: int = 1920,
    viewport_height: int = 1080,
    model: str = _MODEL,
    verbose: bool = False,
) -> dict[str, Any] | None:
    """Return an old-style click target for the first screenshot with consent."""
    for screenshot_path in screenshot_paths:
        result = analyze_cookie_consent(screenshot_path, model=model, verbose=verbose)
        if not result.get("cookie_consent_shown"):
            continue
        box = result.get("accept_button_box_2d")
        if not isinstance(box, list) or len(box) != 4:
            continue
        x0, y0, x1, y1 = [float(v) for v in box]
        return {
            "x": int(round(((x0 + x1) / 2.0 / 1000.0) * viewport_width)),
            "y": int(round(((y0 + y1) / 2.0 / 1000.0) * viewport_height)),
            "box_2d": box,
            "accept_button_text": result.get("accept_button_text"),
            "raw": result,
            "screenshot": str(screenshot_path),
        }
    return None


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Verify cookie consent or bot-protection state from viewport screenshots")
    parser.add_argument("screenshots", nargs="+")
    parser.add_argument("--challenge", action="store_true", help="Detect Cloudflare/Akamai-style bot protection instead of cookie consent")
    parser.add_argument("--verbose", "-v", action="store_true")
    args = parser.parse_args()

    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)
    for path in args.screenshots:
        if args.challenge:
            result = detect_cloudflare_challenge(path, verbose=args.verbose)
        else:
            result = analyze_cookie_consent(path, verbose=args.verbose)
        print(json.dumps(result, ensure_ascii=False, indent=2))
