#!/usr/bin/env python3
"""Screenshot-to-structured-Markdown pipeline for crawled web pages.

The crawler uses this module only when --markdown true is enabled. The HTML and
plain .txt crawl outputs remain available, but the .md/.json output is generated
from screenshots through a vision model so visual hierarchy is preserved.
"""
from __future__ import annotations

import hashlib
import json
import mimetypes
import os
import re
import sys
import threading
import time
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

try:
    from google import genai
    from google.genai import errors as genai_errors
    from google.genai import types as genai_types
except ImportError:  # pragma: no cover - optional dependency guard
    genai = None
    genai_types = None

    class _MissingGenAIError(Exception):
        pass

    class _MissingGenAIErrors:
        ClientError = _MissingGenAIError
        ServerError = _MissingGenAIError
        APIError = _MissingGenAIError

    genai_errors = _MissingGenAIErrors()

WEBPAGE_SCHEMA_VERSION = "webpage/v1"
WEBPAGE_VISION_CACHE_VERSION = "webpage-vision-cache/v1"
DEFAULT_GEMINI_MODEL = "google_gemini.gemini-3.1-flash-lite-preview"
DEFAULT_MAX_ATTEMPTS = 6
DEFAULT_REQUEST_PAUSE_SECONDS = 0.2
GEMINI_PRICING_USD_PER_MILLION = {
    "models/gemini-3.1-flash-lite-preview": {"input": 0.25, "output": 1.50},
}

WEB_BLOCK_TYPES = [
    "heading",
    "paragraph",
    "list",
    "table",
    "code",
    "image",
    "navigation",
    "hero",
    "button",
    "form",
    "card",
    "footer",
    "link_list",
    "modal",
    "breadcrumb",
    "testimonial",
    "pricing",
    "call_to_action",
]

PAGE_SCHEMA: dict[str, Any] = {
    "type": "object",
    "additionalProperties": False,
    "required": ["screenshot_number", "page_title", "url", "viewport", "blocks"],
    "properties": {
        "screenshot_number": {"type": "integer", "minimum": 1},
        "page_title": {"type": ["string", "null"]},
        "url": {"type": "string"},
        "viewport": {
            "type": "object",
            "additionalProperties": False,
            "required": ["width", "height"],
            "properties": {
                "width": {"type": "integer", "minimum": 1},
                "height": {"type": "integer", "minimum": 1},
            },
        },
        "blocks": {
            "type": "array",
            "items": {
                "type": "object",
                "additionalProperties": False,
                "required": ["type"],
                "properties": {
                    "type": {"type": "string", "enum": WEB_BLOCK_TYPES},
                    "level": {"type": ["integer", "null"], "minimum": 1, "maximum": 6},
                    "text": {"type": ["string", "null"]},
                    "items": {"type": ["array", "null"], "items": {"type": "string"}},
                    "ordered": {"type": ["boolean", "null"]},
                    "markdown": {"type": ["string", "null"]},
                    "language": {"type": ["string", "null"]},
                    "alt": {"type": ["string", "null"]},
                    "description": {"type": ["string", "null"]},
                    "role": {"type": ["string", "null"]},
                    "href": {"type": ["string", "null"]},
                    "aria_label": {"type": ["string", "null"]},
                    "position": {
                        "type": ["object", "null"],
                        "additionalProperties": False,
                        "required": ["region", "order"],
                        "properties": {
                            "region": {"type": ["string", "null"]},
                            "order": {"type": ["integer", "null"], "minimum": 1},
                        },
                    },
                    "fields": {
                        "type": ["array", "null"],
                        "items": {
                            "type": "object",
                            "additionalProperties": False,
                            "required": ["label", "type", "placeholder", "required", "value", "help_text"],
                            "properties": {
                                "label": {"type": ["string", "null"]},
                                "type": {"type": ["string", "null"]},
                                "placeholder": {"type": ["string", "null"]},
                                "required": {"type": ["boolean", "null"]},
                                "value": {"type": ["string", "null"]},
                                "help_text": {"type": ["string", "null"]},
                            },
                        },
                    },
                    "links": {
                        "type": ["array", "null"],
                        "items": {
                            "type": "object",
                            "additionalProperties": False,
                            "required": ["label", "href", "role"],
                            "properties": {
                                "label": {"type": ["string", "null"]},
                                "href": {"type": ["string", "null"]},
                                "role": {"type": ["string", "null"]},
                            },
                        },
                    },
                    "continues_from_previous_screenshot": {"type": "boolean"},
                    "continues_to_next_screenshot": {"type": "boolean"},
                    "bounding_box": {
                        "type": ["object", "null"],
                        "description": "Pixel bounding box of this image in the screenshot (x,y left-top, w,h width-height)",
                        "additionalProperties": False,
                        "required": ["x", "y", "w", "h"],
                        "properties": {
                            "x": {"type": "integer", "minimum": 0},
                            "y": {"type": "integer", "minimum": 0},
                            "w": {"type": "integer", "minimum": 1},
                            "h": {"type": "integer", "minimum": 1},
                        },
                    },
                    "suggested_filename": {
                        "type": ["string", "null"],
                        "description": "Hint for matching this image to a downloaded file (e.g. product-name or logo)",
                    },
                    "image_type": {
                        "type": ["string", "null"],
                        "enum": ["product_image", "diagram", "process_description", "other_image"],
                        "description": "Classification of the image content",
                    },
                    "image_description": {
                        "type": ["string", "null"],
                        "description": "Concise description of what the image shows. For product_image: what product, angle, setting. For diagram: what it illustrates. For process_description: what process/flow. For other_image: general description.",
                    },
                },
            },
        },
    },
}

_thread_local = threading.local()


@dataclass(slots=True)
class UsageStats:
    prompt_tokens: int = 0
    output_tokens: int = 0
    total_tokens: int = 0

    def add(self, other: "UsageStats") -> None:
        self.prompt_tokens += other.prompt_tokens
        self.output_tokens += other.output_tokens
        self.total_tokens += other.total_tokens


def utc_now_iso() -> str:
    return datetime.now(UTC).isoformat()


def normalize_space(text: str) -> str:
    return re.sub(r"\s+", " ", text or "").strip()


def sanitize_text_block(text: str) -> str:
    lines = [line.rstrip() for line in (text or "").splitlines()]
    value = "\n".join(lines).strip()
    value = re.sub(r"\n{3,}", "\n\n", value)
    return value


def strip_list_marker(text: str) -> str:
    value = normalize_space(text)
    value = re.sub(r"^[-*•]\s+", "", value)
    value = re.sub(r"^\d+[.)]\s+", "", value)
    return value.strip()


def slug_anchor(value: str) -> str:
    slug = normalize_space(value).lower()
    slug = re.sub(r"[^a-z0-9\- _]+", "", slug)
    slug = re.sub(r"[\s_]+", "-", slug).strip("-")
    return slug or "section"


def canonical_gemini_model_name(model: str) -> str:
    value = (model or "").strip()
    mapping = {
        "google_gemini.gemini-3.1-flash-lite-preview": "models/gemini-3.1-flash-lite-preview",
        "google_gemini.gemini-3-flash-preview": "models/gemini-3-flash-preview",
        "gemini-3.1-flash-lite-preview": "models/gemini-3.1-flash-lite-preview",
        "gemini-3-flash-preview": "models/gemini-3-flash-preview",
    }
    return mapping.get(value, value)


def estimate_gemini_cost_usd(model: str, usage: UsageStats) -> float | None:
    pricing = GEMINI_PRICING_USD_PER_MILLION.get(canonical_gemini_model_name(model))
    if pricing is None:
        return None
    input_cost = (usage.prompt_tokens / 1_000_000.0) * float(pricing["input"])
    output_cost = (usage.output_tokens / 1_000_000.0) * float(pricing["output"])
    return input_cost + output_cost


def image_mime_type(path: Path) -> str:
    guessed, _ = mimetypes.guess_type(path.name)
    return guessed or "image/png"


def get_genai_client(api_key: str) -> genai.Client:
    if genai is None:
        raise ImportError("google-genai is required for --markdown true. Install with: pip install google-genai")
    client = getattr(_thread_local, "genai_client", None)
    if client is None:
        client = genai.Client(api_key=api_key)
        _thread_local.genai_client = client
    return client


def usage_from_response(response: Any) -> UsageStats:
    usage_metadata = getattr(response, "usage_metadata", None)
    return UsageStats(
        prompt_tokens=int(getattr(usage_metadata, "prompt_token_count", 0) or 0),
        output_tokens=int(
            (getattr(usage_metadata, "candidates_token_count", 0) or 0)
            + (getattr(usage_metadata, "thoughts_token_count", 0) or 0)
        ),
        total_tokens=int(getattr(usage_metadata, "total_token_count", 0) or 0),
    )


def log_retry(provider: str, code: int | None, attempt: int, delay: float, message: str) -> None:
    print(
        f"Retrying {provider} request after error code={code!r} attempt={attempt} delay={delay:.1f}s: {message}",
        file=sys.stderr,
        flush=True,
    )


def write_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temp_path = path.with_suffix(path.suffix + ".tmp")
    with temp_path.open("w", encoding="utf-8", newline="\n") as handle:
        json.dump(payload, handle, ensure_ascii=False, indent=2)
        handle.write("\n")
    temp_path.replace(path)


def write_text(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temp_path = path.with_suffix(path.suffix + ".tmp")
    temp_path.write_text(text, encoding="utf-8", newline="\n")
    temp_path.replace(path)


def screenshot_signature(path: Path) -> dict[str, Any]:
    digest = hashlib.sha256(path.read_bytes()).hexdigest()
    return {"sha256": digest, "size_bytes": path.stat().st_size}


def _metadata_line(metadata: dict[str, Any]) -> str:
    viewport = metadata.get("viewport") if isinstance(metadata.get("viewport"), dict) else {}
    return json.dumps(
        {
            "screenshot_number": metadata.get("screenshot_number"),
            "url": metadata.get("url"),
            "page_title": metadata.get("page_title"),
            "viewport": {
                "width": int(viewport.get("width") or 0),
                "height": int(viewport.get("height") or 0),
            },
            "scroll_y": metadata.get("scroll_y"),
            "tile_index": metadata.get("tile_index"),
            "captured_at": metadata.get("captured_at"),
            "screenshot_path": metadata.get("screenshot_path"),
        },
        ensure_ascii=False,
        indent=2,
    )


def extract_screenshot_with_gemini(
    *,
    api_key: str,
    model: str,
    screenshot_path: Path,
    screenshot_metadata: dict[str, Any],
    previous_screenshot_metadata: dict[str, Any] | None = None,
    next_screenshot_metadata: dict[str, Any] | None = None,
    text_hint: str | None = None,
    max_attempts: int = DEFAULT_MAX_ATTEMPTS,
) -> tuple[dict[str, Any], UsageStats]:
    """Extract structured webpage blocks from one screenshot using Gemini.

    The screenshot pixels are the source of truth. Optional text_hint should only
    be used to resolve tiny or ambiguous visible text and must not introduce
    hidden content.
    """
    current_bytes = screenshot_path.read_bytes()
    current_mime = image_mime_type(screenshot_path)
    screenshot_number = int(screenshot_metadata.get("screenshot_number") or 1)
    viewport = screenshot_metadata.get("viewport") if isinstance(screenshot_metadata.get("viewport"), dict) else {}
    url = str(screenshot_metadata.get("url") or "")
    page_title = screenshot_metadata.get("page_title")

    prompt_parts: list[Any] = [
        (
            "You convert live webpage screenshots into stable structured JSON for AI workflows. "
            "Return blocks only for the CURRENT screenshot. If there are previous or next screenshot metadata entries, use them only to decide continuation flags across vertical page tiles. "
            "Use screenshot pixels as the source of truth. Preserve visible text exactly where possible. Do not summarize marketing copy, translate, or invent hidden text.\n\n"
            "Rules:\n"
            "- Capture visual hierarchy: navigation, breadcrumb, hero, headings, paragraphs, cards, forms, buttons, CTAs, pricing blocks, testimonials, tables, code, and footer.\n"
            "- Emit heading blocks only for visibly prominent page headings. Use levels 1 to 6 when clear.\n"
            "- Use navigation/link_list blocks for visible groups of links. Put visible link labels into links/items. Set href only when the target URL is visibly printed or explicitly present in the provided metadata; otherwise use null.\n"
            "- Use button or call_to_action blocks for visible buttons. Preserve the button label exactly.\n"
            "- Use form blocks for visible forms. Include visible fields, labels, placeholders, required markers, help text, and submit buttons. Do not invent hidden form fields.\n"
            "- Use table blocks only when there is actual visible tabular structure. Put valid Markdown table text into markdown.\n"
            "- Use code blocks for visible source code, command snippets, API examples, or configuration snippets. Preserve indentation and punctuation.\n"
            "- Ignore browser chrome, cookie consent if it was accepted/closed, decorative lines, and empty loading placeholders.\n"
            "- If text is uncertain, keep the closest visible reading and add '(uncertain)' in the text or description. Do not guess.\n"
            "- If a paragraph/list/table/code/card/form clearly starts in the previous screenshot tile, set continues_from_previous_screenshot=true.\n"
            "- If it clearly continues below the current screenshot tile, set continues_to_next_screenshot=true.\n"
            "- The JSON schema requires all fields. Use null or [] when a field does not apply.\n\n"
            f"Schema version: {WEBPAGE_SCHEMA_VERSION}\n"
            f"Current screenshot metadata:\n{_metadata_line(screenshot_metadata)}\n"
            f"Current URL: {url}\n"
            f"Current page title from browser: {page_title or '(unknown)'}\n"
            f"Current viewport: {int(viewport.get('width') or 0)}x{int(viewport.get('height') or 0)}\n"
        )
    ]
    if previous_screenshot_metadata is not None:
        prompt_parts.append(f"Previous screenshot metadata for continuation context only:\n{_metadata_line(previous_screenshot_metadata)}")
    if next_screenshot_metadata is not None:
        prompt_parts.append(f"Next screenshot metadata for continuation context only:\n{_metadata_line(next_screenshot_metadata)}")
    if text_hint:
        prompt_parts.append(
            "Optional visible-text hint from the crawler. Use only to resolve tiny visible text. "
            "If it conflicts with the screenshot, trust the screenshot and do not include non-visible text:\n"
            + text_hint[:5000]
        )
    prompt_parts.extend(
        [
            "Current webpage screenshot. Extract blocks from this screenshot only:",
            genai_types.Part.from_bytes(data=current_bytes, mime_type=current_mime),
        ]
    )

    delay_seconds = 2.0
    canonical_model = canonical_gemini_model_name(model)
    for attempt in range(1, max_attempts + 1):
        try:
            client = get_genai_client(api_key)
            response = client.models.generate_content(
                model=canonical_model,
                contents=prompt_parts,
                config=genai_types.GenerateContentConfig(
                    temperature=0.1,
                    top_p=0.95,
                    max_output_tokens=16384,
                    response_mime_type="application/json",
                    response_json_schema=PAGE_SCHEMA,
                    thinking_config=genai_types.ThinkingConfig(thinking_budget=0),
                ),
            )
            if not response.text:
                raise RuntimeError("Gemini returned an empty response")
            data = json.loads(response.text)
            if not isinstance(data, dict) or not isinstance(data.get("blocks"), list):
                raise RuntimeError("Gemini returned an invalid screenshot payload")
            data["screenshot_number"] = int(data.get("screenshot_number") or screenshot_number)
            data["url"] = str(data.get("url") or url)
            return data, usage_from_response(response)
        except json.JSONDecodeError:
            # Don't retry malformed JSON — same input produces same broken output
            raise RuntimeError(f"Gemini returned malformed JSON: {response.text[:200] if response else 'no response'}")
        except (genai_errors.ClientError, genai_errors.ServerError, genai_errors.APIError) as exc:
            code = getattr(exc, "code", None)
            if code in {404, 408, 429, 500, 502, 503, 504} and attempt < max_attempts:
                log_retry("Gemini", code, attempt, delay_seconds, str(exc))
                time.sleep(delay_seconds)
                delay_seconds = min(delay_seconds * 2.0, 60.0)
                continue
            raise
        except (OSError, TimeoutError, ValueError, RuntimeError) as exc:
            if attempt < max_attempts:
                log_retry("Gemini", None, attempt, delay_seconds, str(exc))
                time.sleep(delay_seconds)
                delay_seconds = min(delay_seconds * 2.0, 60.0)
                continue
            raise
    raise RuntimeError("Unreachable Gemini retry path")


def _clean_links(value: Any) -> list[dict[str, str | None]]:
    if not isinstance(value, list):
        return []
    links: list[dict[str, str | None]] = []
    for raw in value:
        if not isinstance(raw, dict):
            continue
        label = normalize_space(str(raw.get("label") or ""))
        href = normalize_space(str(raw.get("href") or "")) or None
        role = normalize_space(str(raw.get("role") or "")) or None
        if label or href:
            links.append({"label": label or None, "href": href, "role": role})
    return links


def _clean_fields(value: Any) -> list[dict[str, Any]]:
    if not isinstance(value, list):
        return []
    fields: list[dict[str, Any]] = []
    for raw in value:
        if not isinstance(raw, dict):
            continue
        field = {
            "label": normalize_space(str(raw.get("label") or "")) or None,
            "type": normalize_space(str(raw.get("type") or "")) or None,
            "placeholder": normalize_space(str(raw.get("placeholder") or "")) or None,
            "required": raw.get("required") if isinstance(raw.get("required"), bool) else None,
            "value": normalize_space(str(raw.get("value") or "")) or None,
            "help_text": normalize_space(str(raw.get("help_text") or "")) or None,
        }
        if any(value is not None for value in field.values()):
            fields.append(field)
    return fields


def _clean_position(value: Any, fallback_region: str | None = None, fallback_order: int | None = None) -> dict[str, Any]:
    if isinstance(value, dict):
        region = normalize_space(str(value.get("region") or "")) or fallback_region
        try:
            order = int(value.get("order") or fallback_order or 0) or None
        except Exception:
            order = fallback_order
        return {"region": region, "order": order}
    return {"region": fallback_region, "order": fallback_order}


def sanitize_page_blocks(raw_blocks: list[dict[str, Any]], screenshot_rel: str) -> list[dict[str, Any]]:
    sanitized: list[dict[str, Any]] = []
    for index, raw_block in enumerate(raw_blocks, start=1):
        if not isinstance(raw_block, dict):
            continue
        block_type = str(raw_block.get("type") or "").strip().lower()
        if block_type not in WEB_BLOCK_TYPES:
            continue

        position = _clean_position(raw_block.get("position"), fallback_region=block_type, fallback_order=index)
        base = {
            "type": block_type,
            "role": normalize_space(str(raw_block.get("role") or "")) or None,
            "href": normalize_space(str(raw_block.get("href") or "")) or None,
            "aria_label": normalize_space(str(raw_block.get("aria_label") or "")) or None,
            "position": position,
            "continues_from_previous_screenshot": bool(raw_block.get("continues_from_previous_screenshot", False)),
            "continues_to_next_screenshot": bool(raw_block.get("continues_to_next_screenshot", False)),
        }

        if block_type == "heading":
            text = normalize_space(str(raw_block.get("text") or ""))
            try:
                level = int(raw_block.get("level") or 1)
            except Exception:
                level = 1
            level = max(1, min(6, level))
            if text:
                sanitized.append({**base, "level": level, "text": text})
            continue

        if block_type == "paragraph":
            text = sanitize_text_block(str(raw_block.get("text") or ""))
            if text:
                sanitized.append({**base, "text": text})
            continue

        if block_type == "list":
            items = [strip_list_marker(str(item)) for item in list(raw_block.get("items") or [])]
            items = [item for item in items if item]
            if items:
                sanitized.append({**base, "ordered": bool(raw_block.get("ordered", False)), "items": items})
            continue

        if block_type == "table":
            markdown = str(raw_block.get("markdown") or raw_block.get("text") or "").strip()
            if markdown:
                sanitized.append({**base, "markdown": markdown})
            continue

        if block_type == "code":
            text = str(raw_block.get("text") or raw_block.get("markdown") or "").rstrip()
            if text:
                sanitized.append({**base, "language": normalize_space(str(raw_block.get("language") or "")), "text": text})
            continue

        if block_type == "image":
            description = normalize_space(str(raw_block.get("description") or ""))
            alt = normalize_space(str(raw_block.get("alt") or "")) or "Visible image"
            if description or alt:
                sanitized.append({**base, "src": screenshot_rel, "alt": alt, "description": description})
            continue

        text = sanitize_text_block(str(raw_block.get("text") or ""))
        items = [normalize_space(str(item)) for item in list(raw_block.get("items") or []) if normalize_space(str(item))]
        markdown = str(raw_block.get("markdown") or "").strip()
        description = normalize_space(str(raw_block.get("description") or "")) or None
        links = _clean_links(raw_block.get("links"))
        fields = _clean_fields(raw_block.get("fields"))
        if text or items or markdown or description or links or fields:
            sanitized.append(
                {
                    **base,
                    "text": text or None,
                    "items": items,
                    "markdown": markdown or None,
                    "description": description,
                    "links": links,
                    "fields": fields,
                }
            )
    return sanitized


def merge_adjacent_paragraphs(blocks: list[dict[str, Any]]) -> list[dict[str, Any]]:
    merged: list[dict[str, Any]] = []
    for block in blocks:
        if block.get("type") != "paragraph":
            merged.append(block)
            continue
        if merged and merged[-1].get("type") == "paragraph" and not block.get("continues_from_previous_screenshot"):
            merged[-1]["text"] = f"{str(merged[-1].get('text') or '').rstrip()}\n{str(block.get('text') or '').lstrip()}".strip()
            merged[-1]["continues_to_next_screenshot"] = block.get("continues_to_next_screenshot", False)
            continue
        merged.append(dict(block))
    return merged


def merge_continued_blocks(blocks: list[dict[str, Any]]) -> list[dict[str, Any]]:
    merged: list[dict[str, Any]] = []
    for block in blocks:
        if not merged:
            merged.append(dict(block))
            continue
        previous = merged[-1]
        if previous.get("type") == "heading" or block.get("type") == "heading":
            merged.append(dict(block))
            continue
        if not previous.get("continues_to_next_screenshot") or not block.get("continues_from_previous_screenshot"):
            merged.append(dict(block))
            continue
        if previous.get("type") != block.get("type"):
            merged.append(dict(block))
            continue
        block_type = str(block.get("type") or "")
        if block_type in {"paragraph", "code"}:
            previous["text"] = f"{str(previous.get('text') or '').rstrip()}\n{str(block.get('text') or '').lstrip()}".strip()
            previous["continues_to_next_screenshot"] = block.get("continues_to_next_screenshot", False)
            continue
        if block_type == "list" and bool(previous.get("ordered")) == bool(block.get("ordered")):
            previous["items"] = list(previous.get("items") or []) + list(block.get("items") or [])
            previous["continues_to_next_screenshot"] = block.get("continues_to_next_screenshot", False)
            continue
        if block_type == "table":
            previous["markdown"] = f"{str(previous.get('markdown') or '').rstrip()}\n{str(block.get('markdown') or '').lstrip()}".strip()
            previous["continues_to_next_screenshot"] = block.get("continues_to_next_screenshot", False)
            continue
        if block_type in {"navigation", "link_list", "breadcrumb", "form", "card"}:
            previous["items"] = list(previous.get("items") or []) + list(block.get("items") or [])
            previous["links"] = list(previous.get("links") or []) + list(block.get("links") or [])
            previous["fields"] = list(previous.get("fields") or []) + list(block.get("fields") or [])
            if block.get("text"):
                previous["text"] = "\n".join(part for part in [previous.get("text"), block.get("text")] if part)
            previous["continues_to_next_screenshot"] = block.get("continues_to_next_screenshot", False)
            continue
        merged.append(dict(block))
    return merge_adjacent_paragraphs(merged)


def strip_internal_block_fields(block: dict[str, Any]) -> dict[str, Any]:
    cleaned = dict(block)
    cleaned.pop("continues_from_previous_screenshot", None)
    cleaned.pop("continues_to_next_screenshot", None)
    return cleaned


def _title_from_structural_block(block: dict[str, Any]) -> str:
    block_type = str(block.get("type") or "").strip()
    role = normalize_space(str(block.get("role") or ""))
    region = ""
    if isinstance(block.get("position"), dict):
        region = normalize_space(str(block["position"].get("region") or ""))
    mapping = {
        "navigation": "Navigation",
        "breadcrumb": "Breadcrumb",
        "hero": "Hero",
        "footer": "Footer",
        "pricing": "Pricing",
        "testimonial": "Testimonials",
        "call_to_action": "Call to Action",
        "form": "Form",
        "modal": "Modal",
        "link_list": "Links",
    }
    if role:
        return role.replace("_", " ").title()
    if region and region not in {block_type, "main", "content"}:
        return region.replace("_", " ").title()
    return mapping.get(block_type, "Content")


def assemble_sections_from_blocks(blocks: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    frontmatter_blocks: list[dict[str, Any]] = []
    sections: list[dict[str, Any]] = []
    stack: list[tuple[int, str]] = []
    structural_starters = {"navigation", "breadcrumb", "hero", "footer", "pricing", "testimonial", "call_to_action", "form", "modal"}

    def start_section(heading: str, level: int = 1) -> dict[str, Any]:
        nonlocal stack
        level = max(1, min(6, level))
        while stack and stack[-1][0] >= level:
            stack.pop()
        stack.append((level, heading))
        section = {"id": "", "heading": heading, "level": level, "path": [item[1] for item in stack], "blocks": []}
        sections.append(section)
        return section

    for block in blocks:
        block_type = str(block.get("type") or "")
        if block_type == "heading":
            heading = str(block.get("text") or "").strip()
            if not heading:
                continue
            level = max(1, min(6, int(block.get("level") or 1)))
            start_section(heading, level)
            continue

        cleaned = strip_internal_block_fields(block)
        if block_type in structural_starters:
            heading = _title_from_structural_block(block)
            current_heading = normalize_space(str(sections[-1].get("heading") or "")) if sections else ""
            if not sections or current_heading != heading:
                start_section(heading, 1)
            sections[-1]["blocks"].append(cleaned)
            continue

        if sections:
            sections[-1]["blocks"].append(cleaned)
        else:
            frontmatter_blocks.append(cleaned)
    return frontmatter_blocks, sections


@dataclass(slots=True)
class TocEntry:
    id: str
    level: int
    title: str
    path: list[str]


def build_toc_entries(title: str, sections: list[dict[str, Any]]) -> list[TocEntry]:
    entries: list[TocEntry] = []
    seen_ids: set[str] = set()
    for section in sections:
        heading = normalize_space(str(section.get("heading") or ""))
        if not heading:
            continue
        level = max(1, min(6, int(section.get("level") or 1)))
        anchor = slug_anchor(heading)
        if anchor in seen_ids:
            suffix = 2
            while f"{anchor}-{suffix}" in seen_ids:
                suffix += 1
            anchor = f"{anchor}-{suffix}"
        seen_ids.add(anchor)
        section["id"] = anchor
        entries.append(TocEntry(id=anchor, level=level, title=heading, path=list(section.get("path") or [])))
    return entries


def _render_links(links: list[dict[str, Any]]) -> list[str]:
    lines: list[str] = []
    for link in links:
        label = normalize_space(str(link.get("label") or ""))
        href = normalize_space(str(link.get("href") or ""))
        role = normalize_space(str(link.get("role") or ""))
        if not label and not href:
            continue
        rendered = f"[{label}]({href})" if label and href else (label or href)
        if role:
            rendered = f"{rendered} ({role})"
        lines.append(f"- {rendered}")
    return lines


def _render_fields(fields: list[dict[str, Any]]) -> list[str]:
    lines: list[str] = []
    for field in fields:
        parts: list[str] = []
        label = normalize_space(str(field.get("label") or ""))
        if label:
            parts.append(label)
        field_type = normalize_space(str(field.get("type") or ""))
        if field_type:
            parts.append(f"type: {field_type}")
        placeholder = normalize_space(str(field.get("placeholder") or ""))
        if placeholder:
            parts.append(f"placeholder: {placeholder}")
        if field.get("required") is True:
            parts.append("required")
        value = normalize_space(str(field.get("value") or ""))
        if value:
            parts.append(f"value: {value}")
        help_text = normalize_space(str(field.get("help_text") or ""))
        if help_text:
            parts.append(f"help: {help_text}")
        if parts:
            lines.append("- " + "; ".join(parts))
    return lines


def render_block(block: dict[str, Any]) -> str:
    block_type = block.get("type")
    if block_type == "paragraph":
        return str(block.get("text") or "").strip()
    if block_type == "list":
        ordered = bool(block.get("ordered", False))
        items = list(block.get("items") or [])
        if ordered:
            return "\n".join(f"1. {item}" for item in items if str(item).strip())
        return "\n".join(f"- {item}" for item in items if str(item).strip())
    if block_type == "table":
        return str(block.get("markdown") or "").strip()
    if block_type == "code":
        language = str(block.get("language") or "").strip()
        body = str(block.get("text") or "").rstrip()
        return f"```{language}\n{body}\n```".rstrip()
    if block_type == "image":
        alt = str(block.get("alt") or "Visible image").strip()
        src = str(block.get("src") or "").strip()
        description = str(block.get("description") or "").strip()
        image_type = str(block.get("image_type") or "").strip()
        image_description = str(block.get("image_description") or "").strip()
        lines: list[str] = []
        if src:
            lines.append(f"![{alt}]({src})")
        if image_type:
            lines.append(f"*Type: {image_type}*")
        if image_description:
            lines.append(f"*{image_description}*")
        elif description:
            lines.append(f"*{description}*")
        return "\n".join(lines).strip()

    if block_type in {"navigation", "link_list", "breadcrumb"}:
        title = _title_from_structural_block(block)
        lines = [f"**{title}:**"]
        text = str(block.get("text") or "").strip()
        if text:
            lines.append(text)
        links = _render_links(list(block.get("links") or []))
        if links:
            lines.extend(links)
        else:
            items = [normalize_space(str(item)) for item in list(block.get("items") or []) if normalize_space(str(item))]
            lines.extend(f"- {item}" for item in items)
        return "\n".join(lines).strip()

    if block_type in {"button", "call_to_action"}:
        label = str(block.get("text") or "").strip()
        href = str(block.get("href") or "").strip()
        prefix = "Call to action" if block_type == "call_to_action" else "Button"
        if label and href:
            return f"**{prefix}:** [{label}]({href})"
        if label:
            return f"**{prefix}:** {label}"
        return ""

    if block_type == "form":
        lines = ["**Form:**"]
        text = str(block.get("text") or "").strip()
        if text:
            lines.append(text)
        field_lines = _render_fields(list(block.get("fields") or []))
        if field_lines:
            lines.append("Visible fields:")
            lines.extend(field_lines)
        link_lines = _render_links(list(block.get("links") or []))
        if link_lines:
            lines.append("Visible links/buttons:")
            lines.extend(link_lines)
        items = [normalize_space(str(item)) for item in list(block.get("items") or []) if normalize_space(str(item))]
        if items:
            lines.append("Visible form text:")
            lines.extend(f"- {item}" for item in items)
        return "\n".join(lines).strip()

    if block_type in {"hero", "card", "footer", "modal", "testimonial", "pricing"}:
        title = _title_from_structural_block(block)
        lines = [f"**{title}:**"]
        text = str(block.get("text") or "").strip()
        if text:
            lines.append(text)
        description = str(block.get("description") or "").strip()
        if description:
            lines.append(f"*Description: {description}*")
        items = [normalize_space(str(item)) for item in list(block.get("items") or []) if normalize_space(str(item))]
        lines.extend(f"- {item}" for item in items)
        link_lines = _render_links(list(block.get("links") or []))
        lines.extend(link_lines)
        markdown = str(block.get("markdown") or "").strip()
        if markdown:
            lines.append(markdown)
        return "\n".join(lines).strip()

    return str(block.get("text") or "").strip()


def render_frontmatter(frontmatter_blocks: list[dict[str, Any]]) -> str:
    parts = [render_block(block) for block in frontmatter_blocks]
    return "\n\n".join(part for part in parts if part).strip()


def render_toc(entries: list[TocEntry]) -> str:
    lines = ["## Table of Contents"]
    for entry in entries:
        indent = "  " * max(entry.level - 1, 0)
        lines.append(f"{indent}- [{entry.title}](#{entry.id})")
    return "\n".join(lines).strip()


def render_sections(sections: list[dict[str, Any]]) -> str:
    parts: list[str] = []
    for section in sections:
        heading = str(section.get("heading") or "").strip()
        level = max(1, min(6, int(section.get("level") or 1)))
        blocks = [render_block(block) for block in list(section.get("blocks") or [])]
        section_text = "\n\n".join(part for part in blocks if part).strip()
        if heading:
            prefix = "#" * level
            if section_text:
                parts.append(f"{prefix} {heading}\n\n{section_text}".strip())
            else:
                parts.append(f"{prefix} {heading}")
        elif section_text:
            parts.append(section_text)
    return "\n\n".join(part for part in parts if part).strip()


def choose_document_title(url: str, screenshot_payloads: list[dict[str, Any]], sections: list[dict[str, Any]]) -> str:
    for payload in screenshot_payloads:
        title = normalize_space(str(payload.get("page_title") or ""))
        if title:
            return title
    for section in sections:
        heading = normalize_space(str(section.get("heading") or ""))
        if heading and heading.lower() not in {"navigation", "breadcrumb", "hero"}:
            return heading
    parsed = re.sub(r"[-_]+", " ", Path(url.split("?", 1)[0]).stem).strip()
    return parsed.title() if parsed else "Webpage"


def build_markdown_output(
    *,
    source_url: str,
    title: str,
    screenshot_metadata: list[dict[str, Any]],
    frontmatter_blocks: list[dict[str, Any]],
    toc_entries: list[TocEntry],
    sections: list[dict[str, Any]],
) -> str:
    escaped_title = title.replace('"', '\\"')
    frontmatter_lines = [
        "---",
        f"schema: {WEBPAGE_SCHEMA_VERSION}",
        f'title: "{escaped_title}"',
        f"source_url: {source_url}",
        f"screenshot_count: {len(screenshot_metadata)}",
        f"generated_at: {utc_now_iso()}",
        "---",
    ]
    parts = ["\n".join(frontmatter_lines)]
    frontmatter = render_frontmatter(frontmatter_blocks)
    if frontmatter:
        parts.append(frontmatter)
    if toc_entries:
        parts.append(render_toc(toc_entries))
    body = render_sections(sections)
    if body:
        parts.append(body)
    return "\n\n".join(part for part in parts if part).strip() + "\n"


def build_json_output(
    *,
    source_url: str,
    title: str,
    screenshot_metadata: list[dict[str, Any]],
    screenshot_payloads: list[dict[str, Any]],
    frontmatter_blocks: list[dict[str, Any]],
    toc_entries: list[TocEntry],
    sections: list[dict[str, Any]],
    model: str,
    usage: UsageStats,
) -> dict[str, Any]:
    primary_viewport: dict[str, Any] = {}
    if screenshot_metadata:
        viewport = screenshot_metadata[0].get("viewport")
        if isinstance(viewport, dict):
            primary_viewport = viewport
    captured_at = screenshot_metadata[0].get("captured_at") if screenshot_metadata else None
    return {
        "schema_version": WEBPAGE_SCHEMA_VERSION,
        "generated_at": utc_now_iso(),
        "source": {
            "url": source_url,
            "screenshot_paths": [str(item.get("screenshot_path") or "") for item in screenshot_metadata],
            "screenshots": screenshot_metadata,
            "viewport": primary_viewport,
            "captured_at": captured_at,
        },
        "models": {"page_vision": model},
        "usage": {
            "gemini": {
                "prompt_tokens": usage.prompt_tokens,
                "output_tokens": usage.output_tokens,
                "total_tokens": usage.total_tokens,
                "estimated_cost_usd": estimate_gemini_cost_usd(model, usage),
            }
        },
        "title": title,
        "frontmatter_blocks": frontmatter_blocks,
        "table_of_contents": [
            {"id": entry.id, "level": entry.level, "title": entry.title, "path": entry.path}
            for entry in toc_entries
        ],
        "sections": sections,
        "vision_pages": screenshot_payloads,
    }


def _detect_images_pass2(
    *,
    api_key: str,
    model: str,
    screenshot_path: Path,
    max_attempts: int = 3,
) -> list[dict]:
    """Second-pass Gemini call: find visible images and their bounding boxes.

    Uses a focused prompt like the lamp detection approach — separate from
    the structural block extraction for better image detection accuracy.
    Returns list of blocks with type='image', bounding_box, alt, suggested_filename.
    """
    try:
        from PIL import Image
    except ImportError:
        return []

    if genai is None:
        return []

    current_bytes = screenshot_path.read_bytes()
    current_mime = image_mime_type(screenshot_path)
    preview = Image.open(screenshot_path)
    pw, ph = preview.size

    prompt = f"""Find every visible product image, photograph, diagram, logo, and icon on this webpage screenshot.

Return ONLY a JSON array. Each item MUST be:
{{"type": "image", "alt": "short description", "bounding_box": {{"x": int, "y": int, "w": int, "h": int}}, "suggested_filename": "descriptive-name-for-matching"}}

Rules:
- bounding_box x,y is top-left pixel, w,h is width and height in pixels of the image ON THIS SCREENSHOT.
- The screenshot is {pw}x{ph} pixels. All coordinates must be within these bounds.
- Include ALL visible product images, hero banners, photos, diagrams, and the site logo.
- Skip tiny icons (less than 20x20 pixels), text, navigation elements, and the cookie banner.
- suggested_filename is a hint like "prolicht-logo" or "acoustic-panel-photo" to help match this to a downloaded image file.
- If there are no images visible, return [].
- Avoid duplicate boxes for the same image.""".strip()

    client = get_genai_client(api_key)
    canonical_model = canonical_gemini_model_name(model)

    for attempt in range(1, max_attempts + 1):
        try:
            response = client.models.generate_content(
                model=canonical_model,
                contents=[
                    prompt,
                    genai_types.Part.from_bytes(data=current_bytes, mime_type=current_mime),
                ],
                config=genai_types.GenerateContentConfig(
                    temperature=0.0,
                    max_output_tokens=4096,
                    response_mime_type="application/json",
                    thinking_config=genai_types.ThinkingConfig(thinking_budget=0),
                ),
            )
            if not response.text:
                continue
            raw = response.text.strip()
            if raw.startswith("```"):
                raw = raw.split("\n", 1)[1] if "\n" in raw else raw[3:]
                if raw.endswith("```"):
                    raw = raw[:-3]
                raw = raw.strip()
            data = json.loads(raw)
            # Accept both array and wrapped-object responses
            if isinstance(data, list):
                return data
            if isinstance(data, dict):
                for key in ("items", "detections", "images", "results"):
                    val = data.get(key)
                    if isinstance(val, list):
                        return val
                # Maybe it IS the block list already wrapped
                if data.get("type") == "image" and "bounding_box" in data:
                    return [data]
                return []
            return []
        except json.JSONDecodeError:
            if attempt < max_attempts:
                time.sleep(1.5 * attempt)
            continue
        except Exception:
            if attempt < max_attempts:
                time.sleep(2 * attempt)
            continue
    return []


def process_webpage_screenshots(
    *,
    screenshot_metadata: list[dict[str, Any]],
    output_md_path: Path,
    output_json_path: Path,
    api_key: str | None = None,
    model: str = DEFAULT_GEMINI_MODEL,
    max_attempts: int = DEFAULT_MAX_ATTEMPTS,
    text_hint: str | None = None,
    force: bool = False,
) -> tuple[Path, Path, dict[str, Any]]:
    """Run screenshot vision extraction and write structured Markdown + JSON."""
    if not screenshot_metadata:
        raise ValueError("No screenshot metadata supplied")

    if output_md_path.exists() and output_json_path.exists() and not force:
        payload = json.loads(output_json_path.read_text(encoding="utf-8"))
        return output_md_path, output_json_path, payload

    api_key = api_key or os.environ.get("GEMINI_API_KEY")
    if not api_key:
        raise EnvironmentError("GEMINI_API_KEY is not set; it is required for --markdown true")

    usage = UsageStats()
    screenshot_payloads: list[dict[str, Any]] = []
    flattened_blocks: list[dict[str, Any]] = []
    base_dir = output_md_path.parent.resolve()

    for index, metadata in enumerate(screenshot_metadata):
        screenshot_path = Path(str(metadata.get("screenshot_path") or "")).resolve()
        if not screenshot_path.exists():
            raise FileNotFoundError(f"Screenshot not found: {screenshot_path}")
        previous_metadata = screenshot_metadata[index - 1] if index > 0 else None
        next_metadata = screenshot_metadata[index + 1] if index + 1 < len(screenshot_metadata) else None
        payload, page_usage = extract_screenshot_with_gemini(
            api_key=api_key,
            model=model,
            screenshot_path=screenshot_path,
            screenshot_metadata=metadata,
            previous_screenshot_metadata=previous_metadata,
            next_screenshot_metadata=next_metadata,
            text_hint=text_hint,
            max_attempts=max_attempts,
        )
        usage.add(page_usage)
        screenshot_payloads.append(payload)
        try:
            screenshot_rel = screenshot_path.relative_to(base_dir).as_posix()
        except ValueError:
            screenshot_rel = screenshot_path.as_posix()
        flattened_blocks.extend(sanitize_page_blocks(list(payload.get("blocks") or []), screenshot_rel))
        time.sleep(DEFAULT_REQUEST_PAUSE_SECONDS)

    # Pass 2: dedicated image bounding-box detection on full-page screenshot
    if flattened_blocks and screenshot_metadata:
        first_screenshot = Path(str(screenshot_metadata[0].get("screenshot_path") or "")).resolve()
        if first_screenshot.exists():
            try:
                image_blocks = _detect_images_pass2(
                    api_key=api_key,
                    model=model,
                    screenshot_path=first_screenshot,
                    max_attempts=max_attempts,
                )
                if image_blocks:
                    try:
                        screenshot_rel = first_screenshot.relative_to(base_dir).as_posix()
                    except ValueError:
                        screenshot_rel = first_screenshot.as_posix()
                    flattened_blocks.extend(sanitize_page_blocks(image_blocks, screenshot_rel))
            except Exception as exc:
                import logging
                logging.getLogger("crawler").warning("Image detection pass 2 failed: %s", exc)

    # Match image blocks to local downloaded files via perceptual hashing
    if flattened_blocks and screenshot_metadata:
        first_screenshot = Path(str(screenshot_metadata[0].get("screenshot_path") or "")).resolve()
        if first_screenshot.exists():
            images_dir_check = base_dir / "images"
            try:
                flattened_blocks = match_vision_images(flattened_blocks, first_screenshot, images_dir_check)
            except Exception as exc:
                import logging
                logging.getLogger("crawler").warning("Vision image matching failed: %s", exc)

    merged_blocks = merge_continued_blocks(flattened_blocks)
    frontmatter_blocks, sections = assemble_sections_from_blocks(merged_blocks)
    source_url = str(screenshot_metadata[0].get("url") or "")
    title = choose_document_title(source_url, screenshot_payloads, sections)
    toc_entries = build_toc_entries(title, sections)
    markdown_output = build_markdown_output(
        source_url=source_url,
        title=title,
        screenshot_metadata=screenshot_metadata,
        frontmatter_blocks=frontmatter_blocks,
        toc_entries=toc_entries,
        sections=sections,
    )
    json_output = build_json_output(
        source_url=source_url,
        title=title,
        screenshot_metadata=screenshot_metadata,
        screenshot_payloads=screenshot_payloads,
        frontmatter_blocks=frontmatter_blocks,
        toc_entries=toc_entries,
        sections=sections,
        model=model,
        usage=usage,
    )
    write_text(output_md_path, markdown_output)
    write_json(output_json_path, json_output)
    return output_md_path, output_json_path, json_output


# ---------------------------------------------------------------------------
# Image matching via OpenCV template matching + perceptual hash fallback
# ---------------------------------------------------------------------------
def _compute_dhash(img, hash_size: int = 8) -> str:
    """Compute difference hash (dHash). Returns empty string on failure."""
    try:
        from PIL import Image
    except ImportError:
        return ""
    if img is None:
        return ""
    try:
        resized = img.convert("L").resize((hash_size + 1, hash_size), Image.LANCZOS)
        pixels = list(resized.getdata())
        diff_bits = []
        for row in range(hash_size):
            for col in range(hash_size):
                left = pixels[row * (hash_size + 1) + col]
                right = pixels[row * (hash_size + 1) + col + 1]
                diff_bits.append("1" if left > right else "0")
        h = int("".join(diff_bits), 2)
        return hex(h)[2:].zfill(16)
    except Exception:
        return ""


def _ham64(h1: str, h2: str) -> int:
    """Hamming distance via XOR + bit_count for 64-bit hex hashes."""
    try:
        return (int(h1, 16) ^ int(h2, 16)).bit_count()
    except Exception:
        return 999


def _preload_image_hashes(images_dir: Path) -> dict[str, str]:
    """Pre-compute dHash for all images. Returns {path_str: hash}."""
    try:
        from PIL import Image
    except ImportError:
        return {}
    hashes: dict[str, str] = {}
    if not images_dir.exists():
        return hashes
    for img_path in sorted(images_dir.iterdir()):
        if img_path.suffix.lower() not in {".png", ".jpg", ".jpeg", ".gif", ".webp", ".avif", ".bmp", ".svg"}:
            continue
        try:
            with Image.open(img_path) as img:
                img = img.convert("RGB")
                h = _compute_dhash(img)
                if h:
                    hashes[str(img_path)] = h
        except Exception:
            pass
    return hashes


def match_vision_images(
    blocks: list[dict],
    fullpage_path: Path,
    images_dir: Path,
) -> list[dict]:
    """Crop Gemini bounding boxes, hash-match to images/ or save to images/vision/."""
    import logging
    log = logging.getLogger("crawler")
    try:
        from PIL import Image
    except ImportError:
        return blocks

    image_hashes = _preload_image_hashes(images_dir)
    log.debug("match_vision_images: loaded %d image hashes, %d blocks, fullpage=%s", len(image_hashes), len(blocks), fullpage_path.name)

    vision_dir = images_dir / "vision"

    page_stem = fullpage_path.stem.replace("__fullpage", "").replace("__viewport_initial", "")
    seq = 0

    for block in blocks:
        if block.get("type") != "image":
            continue
        bbox = block.get("bounding_box")
        if not isinstance(bbox, dict):
            continue
        x, y, w, h = int(bbox.get("x", 0)), int(bbox.get("y", 0)), int(bbox.get("w", 0)), int(bbox.get("h", 0))
        if w < 16 or h < 16:
            continue
        try:
            full_img = Image.open(fullpage_path)
            crop = full_img.crop((x, y, x + w, y + h))
            seq += 1

            # Try hash matching against downloaded images first
            crop_hash = _compute_dhash(crop)
            matched = None
            if crop_hash and image_hashes:
                best_path = None
                best_distance = 13
                for path, img_hash in image_hashes.items():
                    dist = _ham64(crop_hash, img_hash)
                    if dist < best_distance:
                        best_distance = dist
                        best_path = path
                if best_path and best_distance <= 10:
                    try:
                        matched = f"images/{Path(best_path).relative_to(images_dir).as_posix()}"
                    except ValueError:
                        matched = f"images/{Path(best_path).name}"

            if matched:
                block["src"] = matched
                log.info("Vision image matched: %s -> %s", block.get("alt", "?"), matched)
            else:
                vision_dir.mkdir(parents=True, exist_ok=True)
                crop_name = f"{page_stem}__vision_{seq:03d}.png"
                crop_path = vision_dir / crop_name
                crop.save(str(crop_path), "PNG")
                block["src"] = f"images/vision/{crop_name}"
                log.info("Vision image saved to vision/: %s", crop_name)
        except Exception as exc:
            import logging
            logging.getLogger("crawler").debug("Vision image crop/match failed for bbox %s: %s", bbox, exc)

    return blocks
