#!/usr/bin/env python3
"""Test tiled image detection: A4-ratio tiles with 50% overlap, multi-threaded via Gemini."""
import json, os, sys, time
from concurrent.futures import ThreadPoolExecutor, as_completed
from io import BytesIO
from pathlib import Path
from PIL import Image
from google import genai
from google.genai import types

FULLPAGE = Path(r"D:\repos\crawler\sites_crawled\www.prolicht.at\screenshots\prolicht__at__de__produktkonfigurator__produktkonfiguration__fullpage.png")
MODEL = "models/gemini-3-flash-preview"
TILE_WORKERS = 8
API_SEMAPHORE = None  # set by main()


def tile_image(img: Image.Image, tile_h: int, overlap_ratio: float = 0.5) -> list[tuple[int, int, Image.Image]]:
    """Split image into vertical tiles with 50% overlap. Each tile covers tile_h new content."""
    w, h = img.size
    step = int(tile_h * (1 - overlap_ratio))
    tiles = []
    y = 0
    while y < h:
        y1 = min(y + tile_h, h)
        crop = img.crop((0, y, w, y1))
        tiles.append((y, y1, crop))
        y += step
        if y >= h:
            break
    # Ensure last region is covered
    if tiles and tiles[-1][1] < h:
        y = max(0, h - tile_h)
        crop = img.crop((0, y, w, h))
        tiles.append((y, h, crop))
    return tiles


def _validate_boxes(items: list, tile_y0: int, tile_w: int, tile_h: int) -> list[dict]:
    """Validate and fix bbox fields, map to full-page coordinates."""
    valid = []
    seen_hashes = set()
    for item in items:
        if not isinstance(item, dict):
            continue
        bbox = item.get("bounding_box")
        if not isinstance(bbox, dict):
            continue
        x = int(bbox.get("x", bbox.get("translate", bbox.get("left", 0))))
        y = int(bbox.get("y", bbox.get("label", bbox.get("top", 0))))
        w = int(bbox.get("w", bbox.get("width", 10)))
        h = int(bbox.get("h", bbox.get("height", bbox.get("label", 10))))
        if w < 24 or h < 24:
            continue
        # Clamp to tile bounds
        x = max(0, min(x, tile_w - 1))
        y = max(0, min(y, tile_h - 1))
        w = min(w, tile_w - x)
        h = min(h, tile_h - y)
        # Deduplicate by rough position hash
        pos_hash = f"{x//20}_{(y+tile_y0)//20}_{w//20}_{h//20}"
        if pos_hash in seen_hashes:
            continue
        seen_hashes.add(pos_hash)
        item["bounding_box"] = {"x": x, "y": y + tile_y0, "w": w, "h": h}
        item.setdefault("image_type", "product_image")
        item.setdefault("alt", "image")
        item.setdefault("image_description", "")
        item.setdefault("type", "image")
        valid.append(item)
    return valid


def detect_images_on_tile(tile_img: Image.Image, tile_y0: int) -> list[dict]:
    """Send one tile to Gemini, return validated image blocks with full-page coordinates."""
    client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
    w, h = tile_img.size

    prompt = f"""Find every visible IMAGE on this webpage screenshot section.

Return ONLY a JSON array. Each item:
{{"alt":"short name","bounding_box":{{"x":int,"y":int,"w":int,"h":int}},"image_type":"product_image","image_description":"concise description"}}

RULES:
- bounding_box: x=left, y=top, w=width, h=height in PIXELS on this {w}x{h} tile.
  Do NOT use translate/label/left/top as field names — only x,y,w,h.
- image_type: "product_image", "diagram", "process_description", or "other_image"
- alt: visible product name or short label. image_description: one sentence describing the image.
- Include product photos, hero banners, diagrams, process illustrations, logos.
- Skip text, navigation, cookie banners, browser chrome, icons <24px.
- If no images: return []."""

    buf = BytesIO()
    tile_img.save(buf, format="PNG")

    for attempt in range(3):
        try:
            with API_SEMAPHORE:
                response = client.models.generate_content(
                    model=MODEL,
                    contents=[prompt, types.Part.from_bytes(data=buf.getvalue(), mime_type="image/png")],
                    config=types.GenerateContentConfig(
                        temperature=0.0, max_output_tokens=4096,
                        response_mime_type="application/json",
                        thinking_config=types.ThinkingConfig(thinking_budget=0),
                    ),
                )
            if not response.text:
                continue
            raw = response.text.strip().lstrip("```json").lstrip("```").rstrip("```").strip()
            data = json.loads(raw)
            items = data if isinstance(data, list) else data.get("items", data.get("detections", data.get("images", data.get("results", []))))
            if isinstance(items, list):
                return _validate_boxes(items, tile_y0, w, h)
            return []
        except json.JSONDecodeError:
            if attempt < 2:
                time.sleep(1)
        except Exception:
            if attempt < 2:
                time.sleep(1)
    return []


def main() -> int:
    global API_SEMAPHORE
    if not os.environ.get("GEMINI_API_KEY"):
        print("Set GEMINI_API_KEY", file=sys.stderr)
        return 1

    API_SEMAPHORE = __import__("threading").Semaphore(TILE_WORKERS)

    img = Image.open(FULLPAGE)
    w, h = img.size
    tile_h = int(w * 1.414)
    tiles = tile_image(img, tile_h)
    print(f"Full page: {w}x{h}  |  {len(tiles)} tiles (A4={tile_h}px, 50% overlap)  |  {TILE_WORKERS} workers")
    print("=" * 70)

    all_images: list[dict] = []
    with ThreadPoolExecutor(max_workers=TILE_WORKERS) as executor:
        futures = {executor.submit(detect_images_on_tile, crop, y0): (idx, y0, y1) for idx, (y0, y1, crop) in enumerate(tiles, 1)}
        for future in as_completed(futures):
            idx, y0, y1 = futures[future]
            try:
                results = future.result()
            except Exception as e:
                print(f"Tile {idx:>2}: y={y0}-{y1} ERROR: {e}")
                continue
            if results:
                print(f"Tile {idx:>2}: y={y0:>5}-{y1:<5} -> {len(results)} image(s)")
                for r in results[:3]:
                    bbox = r["bounding_box"]
                    print(f"    [{r.get('image_type','?')}] {r.get('alt','?'):30s} ({bbox['x']},{bbox['y']},{bbox['w']}x{bbox['h']})")
                if len(results) > 3:
                    print(f"    ... and {len(results)-3} more")
            else:
                print(f"Tile {idx:>2}: y={y0:>5}-{y1:<5} -> 0 images")
            all_images.extend(results)

    # Global dedup across tiles (position-based)
    final = []
    seen = set()
    for img_block in sorted(all_images, key=lambda b: (b["bounding_box"]["y"], b["bounding_box"]["x"])):
        bbox = img_block["bounding_box"]
        key = (bbox["x"] // 15, bbox["y"] // 15, bbox["w"] // 15, bbox["h"] // 15)
        if key in seen:
            continue
        seen.add(key)
        final.append(img_block)

    print("=" * 70)
    print(f"Total: {len(final)} unique images across {len(tiles)} tiles (raw: {len(all_images)} before dedup)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
