#!/usr/bin/env python
"""Repair low-quality markdown pages: tiled image detection + hash matching + vision markdown.

1. Tiles the full-page screenshot into A4-proportioned slices with 50% overlap
2. Sends each tile to Gemini for image bounding-box detection (multi-threaded)
3. Crops detected regions, computes perceptual hash, matches against downloaded images
4. Generates structural blocks via Gemini vision
5. Merges structural + image blocks into __vision.md with local image links
"""

from __future__ import annotations

import concurrent.futures
import io
import json
import os
import sys
import time
from pathlib import Path
from typing import Any

from PIL import Image

try:
    from google import genai
    from google.genai import types as genai_types
except ImportError:
    genai = None
    genai_types = None

from webpage_vision_markdown import (
    DEFAULT_GEMINI_MODEL,
    DEFAULT_MAX_ATTEMPTS,
    _compute_dhash,
    _ham64,
    _preload_image_hashes,
    process_webpage_screenshots,
)

HERE = Path(__file__).resolve().parent
MODEL = "models/gemini-3-flash-preview"
TILE_WORKERS = 8
QUALITY_CONFIDENCES_TO_REPAIR = {"cannot_understand", "maybe"}


def find_low_quality_pages(out_dir: Path) -> list[dict]:
    pages = []
    for quality_path in sorted(out_dir.glob("*quality.json")):
        data = json.loads(quality_path.read_text(encoding="utf-8"))
        confidence = data.get("confidence", "cannot_understand")
        if confidence not in QUALITY_CONFIDENCES_TO_REPAIR:
            continue
        stem = data.get("_page_stem") or quality_path.stem.replace(".md_quality", "").replace(".md_quality", "").replace("_quality", "")
        if "__vision" in stem:
            continue
        crawl_path = out_dir / f"{stem}.crawl.json"
        source_url = None
        if crawl_path.exists():
            try:
                source_url = json.loads(crawl_path.read_text(encoding="utf-8")).get("url")
            except Exception:
                pass
        if not source_url:
            md_path = out_dir / f"{stem}.md"
            if md_path.exists():
                for line in md_path.read_text(encoding="utf-8").splitlines()[:10]:
                    if line.startswith("source_url:"):
                        source_url = line.split(":", 1)[1].strip().strip('"')
                        break
        if not source_url:
            continue
        pages.append({"stem": stem, "url": source_url, "confidence": confidence,
                       "value_alignment_score": data.get("value_alignment_score", 0),
                       "issues": data.get("issues", [])})
    return pages


def tile_fullpage(img: Image.Image, tile_h: int, overlap_ratio: float = 0.5) -> list[tuple[int, int, Image.Image]]:
    """Split image into vertical tiles with overlap. Returns (y0, y1, crop)."""
    w, h = img.size
    step = max(1, int(tile_h * (1 - overlap_ratio)))
    tiles = []
    y = 0
    while y < h:
        y1 = min(y + tile_h, h)
        tiles.append((y, y1, img.crop((0, y, w, y1))))
        y += step
    if tiles and tiles[-1][1] < h:
        y = max(0, h - tile_h)
        tiles.append((y, h, img.crop((0, y, w, h))))
    return tiles


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"])
    tw, th = tile_img.size

    prompt = f"""Find every visible IMAGE on this webpage 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 {tw}x{th} tile. Only x,y,w,h.
- image_type: "product_image","diagram","process_description","other_image"
- Include product photos, hero banners, diagrams, process illustrations, logos >40px.
- Skip text, navigation, cookie banners, icons <40px.
- If no images: return []."""

    buf = io.BytesIO()
    tile_img.save(buf, format="PNG")
    canonical = "models/gemini-3-flash-preview"

    for attempt in range(3):
        try:
            response = client.models.generate_content(
                model=canonical,
                contents=[prompt, genai_types.Part.from_bytes(data=buf.getvalue(), mime_type="image/png")],
                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().lstrip("```json").lstrip("```").rstrip("```").strip()
            data = json.loads(raw)
            items = data if isinstance(data, list) else []
            if isinstance(data, dict):
                for key in ("items", "detections", "images", "results"):
                    v = data.get(key)
                    if isinstance(v, list):
                        items = v
                        break

            valid = []
            seen = set()
            for item in items:
                if not isinstance(item, dict):
                    continue
                bbox = item.get("bounding_box", {})
                if not isinstance(bbox, dict):
                    continue
                try:
                    bx = bbox.get("x", bbox.get("translate", bbox.get("left", 0)))
                    by = bbox.get("y", bbox.get("label", bbox.get("top", 0)))
                    bw = bbox.get("w", bbox.get("width", 10))
                    bh = bbox.get("h", bbox.get("height", bbox.get("label", 10)))
                    x, y, bw, bh = int(bx), int(by), int(bw), int(bh)
                except (ValueError, TypeError):
                    continue
                if bw < 40 or bh < 40:
                    continue
                x = max(0, min(x, tw - 1)); y = max(0, min(y, th - 1))
                bw = min(bw, tw - x); bh = min(bh, th - y)
                key = (x // 10, (y + tile_y0) // 10, bw // 20, bh // 20)
                if key in seen:
                    continue
                seen.add(key)
                valid.append({
                    "type": "image",
                    "alt": str(item.get("alt", "image")),
                    "bounding_box": {"x": x, "y": y + tile_y0, "w": bw, "h": bh},
                    "image_type": str(item.get("image_type", "product_image")),
                    "image_description": str(item.get("image_description", "")),
                })
            return valid
        except json.JSONDecodeError:
            if attempt < 2: time.sleep(1)
        except Exception:
            if attempt < 2: time.sleep(1)
    return []


def detect_all_images(fullpage_path: Path, tile_workers: int = 8) -> list[dict]:
    """Tile the full-page screenshot, detect images in each tile via Gemini, deduplicate."""
    img = Image.open(fullpage_path)
    w, h = img.size
    tile_h = int(w * 1.414)
    tiles = tile_fullpage(img, tile_h)

    print(f"  Tiled into {len(tiles)} slices (A4 ratio, {tile_h}px, 50% overlap)")

    all_items = []
    api_sem = __import__("threading").Semaphore(tile_workers)

    def _worker(crop, y0):
        with api_sem:
            return detect_images_on_tile(crop, y0)

    with concurrent.futures.ThreadPoolExecutor(max_workers=tile_workers) as executor:
        futures = {executor.submit(_worker, crop, y0): i for i, (y0, y1, crop) in enumerate(tiles)}
        for future in concurrent.futures.as_completed(futures):
            try:
                results = future.result()
                all_items.extend(results)
            except Exception:
                pass

    # Global dedup
    final = []
    seen = set()
    for item in sorted(all_items, key=lambda b: (b["bounding_box"]["y"], b["bounding_box"]["x"])):
        bbox = item["bounding_box"]
        key = (bbox["x"] // 12, bbox["y"] // 12, bbox["w"] // 12, bbox["h"] // 12)
        if key in seen:
            continue
        seen.add(key)
        final.append(item)

    print(f"  Found {len(final)} unique images (raw: {len(all_items)})")
    return final


def match_and_localize_images(
    image_blocks: list[dict],
    fullpage_path: Path,
    images_dir: Path,
) -> list[dict]:
    """Crop each bounding box from full-page, hash-match to images/ or save to images/vision/."""
    image_hashes = _preload_image_hashes(images_dir)
    print(f"  Loaded {len(image_hashes)} image hashes from {images_dir}")

    vision_dir = images_dir / "vision"
    page_stem = fullpage_path.stem.replace("__fullpage", "").replace("__viewport_initial", "")
    full_img = Image.open(fullpage_path)

    matched = 0
    saved_vision = 0
    for idx, block in enumerate(image_blocks):
        bbox = block.get("bounding_box", {})
        x, y, bw, bh = int(bbox.get("x", 0)), int(bbox.get("y", 0)), int(bbox.get("w", 0)), int(bbox.get("h", 0))
        if bw < 16 or bh < 16:
            continue
        try:
            crop = full_img.crop((x, y, x + bw, y + bh))
            crop_hash = _compute_dhash(crop)

            if crop_hash and image_hashes:
                best_path = None
                best_dist = 13
                for path, img_hash in image_hashes.items():
                    dist = _ham64(crop_hash, img_hash)
                    if dist < best_dist:
                        best_dist = dist
                        best_path = path
                if best_path and best_dist <= 10:
                    try:
                        block["src"] = f"images/{Path(best_path).relative_to(images_dir).as_posix()}"
                    except ValueError:
                        block["src"] = f"images/{Path(best_path).name}"
                    matched += 1
                    continue

            # No match — save to vision folder
            vision_dir.mkdir(parents=True, exist_ok=True)
            crop_name = f"{page_stem}__vision_{idx + 1:03d}.png"
            crop.save(str(vision_dir / crop_name), "PNG")
            block["src"] = f"images/vision/{crop_name}"
            saved_vision += 1
        except Exception:
            pass

    print(f"  Image matching: {matched} matched to images/, {saved_vision} saved to images/vision/")
    return image_blocks


def build_vision_markdown(
    stem: str,
    source_url: str,
    structural_blocks: list[dict],
    image_blocks: list[dict],
    out_dir: Path,
) -> str:
    """Build complete vision markdown from structural + image blocks."""
    lines = []
    lines.append("---")
    lines.append('schema: webpage/v1')
    lines.append(f'title: "Vision Repair: {stem}"')
    lines.append(f"source_url: {source_url}")
    lines.append(f"generated_at: {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}")
    lines.append("---")
    lines.append("")

    # Merge: insert image blocks after the first heading in structural blocks
    merged = list(structural_blocks)
    insert_after = 0
    for i, b in enumerate(merged):
        if b.get("type") == "heading":
            insert_after = i + 1
            break
    for img in reversed(image_blocks):
        merged.insert(insert_after, img)

    # Render blocks
    for block in merged:
        bt = block.get("type", "")
        if bt == "heading":
            level = block.get("level") or 2
            text = block.get("text", "")
            if text:
                lines.append(f"{'#' * int(min(level, 6))} {text}")
                lines.append("")
        elif bt == "paragraph":
            text = block.get("text", "")
            if text:
                lines.append(text)
                lines.append("")
        elif bt == "list":
            items = block.get("items") or []
            ordered = block.get("ordered", False)
            for i, item in enumerate(items):
                prefix = f"{i + 1}." if ordered else "-"
                lines.append(f"{prefix} {item}")
            lines.append("")
        elif bt == "image":
            src = block.get("src", "")
            alt = block.get("alt", "image")
            img_type = block.get("image_type", "")
            desc = block.get("image_description", "")
            if src:
                lines.append(f"![{alt}]({src})")
            if img_type:
                lines.append(f"*{img_type}*")
            if desc:
                lines.append(f"*{desc}*")
            lines.append("")
        elif bt in ("navigation", "link_list"):
            title = block.get("text") or "Navigation"
            links = block.get("links") or []
            lines.append(f"**{title}:**")
            for link in links:
                label = link.get("label", "")
                href = link.get("href", "")
                if label and href:
                    lines.append(f"- [{label}]({href})")
                elif label:
                    lines.append(f"- {label}")
            lines.append("")
        elif bt == "form":
            lines.append("**Form:**")
            fields = block.get("fields") or []
            for f in fields:
                label = f.get("label", "")
                ftype = f.get("type", "")
                if label:
                    lines.append(f"- {label}; type: {ftype}")
            lines.append("")
        elif bt in ("code",):
            lines.append(f"```{block.get('language','')}")
            lines.append(block.get("text", ""))
            lines.append("```")
            lines.append("")
        elif bt in ("table",):
            md = block.get("markdown", "")
            if md:
                lines.append(md)
                lines.append("")
        elif bt == "button":
            label = block.get("text", "")
            href = block.get("href", "")
            lines.append(f"**[{label}]({href})**" if href else f"**{label}**")
            lines.append("")
        else:
            text = block.get("text") or block.get("description") or ""
            if text:
                lines.append(text)
                lines.append("")

    return "\n".join(lines)


def repair_page(page: dict, out_dir: Path) -> bool:
    """Full repair pipeline for one page."""
    stem = page["stem"]
    fullpage_path = out_dir / "screenshots" / f"{stem}__fullpage.png"
    if not fullpage_path.exists():
        print(f"  SKIP: no full-page screenshot")
        return False

    images_dir = out_dir / "images"

    # Step 1: Get structural blocks via existing vision pipeline
    print(f"  Pass 1: structural extraction via Gemini...")
    screenshot_metadata = [{
        "screenshot_number": 1,
        "screenshot_path": str(fullpage_path.resolve()),
        "local_screenshot_path": str(fullpage_path.resolve()),
        "source_url": page["url"], "url": page["url"], "page_title": "",
        "viewport": {"width": 1920, "height": 1080},
        "screenshot_pixel_size": {"width": 1904, "height": 985},
        "capture_timestamp": "", "captured_at": "",
        "scroll_y": 0, "scroll_position": {"x": 0, "y": 0},
        "tile_index": 1, "tile_count": 1, "screenshot_type": "fullpage",
    }]

    output_md = out_dir / f"{stem}__vision.md"
    output_json = out_dir / f"{stem}__vision.json"

    try:
        process_webpage_screenshots(
            screenshot_metadata=screenshot_metadata,
            output_md_path=output_md,
            output_json_path=output_json,
            model="google_gemini.gemini-3-flash-preview",
            max_attempts=DEFAULT_MAX_ATTEMPTS,
            text_hint=None, force=True,
        )
    except Exception as e:
        print(f"  FAILED structural: {e}")
        return False

    # Step 2: Tiled image detection + hash matching
    print(f"  Pass 2: tiled image detection via Gemini ({TILE_WORKERS} workers)...")
    image_blocks = detect_all_images(fullpage_path, TILE_WORKERS)

    if image_blocks:
        image_blocks = match_and_localize_images(image_blocks, fullpage_path, images_dir)

    # Step 3: Rebuild vision markdown with merged blocks
    try:
        vision_data = json.loads(output_json.read_text(encoding="utf-8"))
        structural_blocks = vision_data.get("vision_pages", [{}])[0].get("blocks", [])
    except Exception:
        structural_blocks = []

    markdown = build_vision_markdown(stem, page["url"], structural_blocks, image_blocks, out_dir)
    output_md.write_text(markdown, encoding="utf-8")

    # Update JSON with image blocks
    try:
        vision_data["vision_pages"][0]["blocks"] = structural_blocks + image_blocks
        output_json.write_text(json.dumps(vision_data, ensure_ascii=False, indent=2), encoding="utf-8")
    except Exception:
        pass

    print(f"  OK -> {output_md.name} ({len(image_blocks)} images, {len(structural_blocks)} structural blocks)")
    return True


def main() -> int:
    import argparse
    p = argparse.ArgumentParser(description="Repair low-quality markdown via tiled vision + image matching")
    p.add_argument("out_dir", help="Crawl output directory")
    args = p.parse_args()

    if not os.environ.get("GEMINI_API_KEY"):
        print("ERROR: GEMINI_API_KEY not set")
        return 1

    out_dir = Path(args.out_dir)
    pages = find_low_quality_pages(out_dir)
    if not pages:
        print("No low-quality pages found.")
        return 0

    print(f"Found {len(pages)} page(s) needing repair:")
    for page in pages:
        print(f"  [{page['confidence']}] {page['stem']} (align={page['value_alignment_score']}/5)")

    ok = failed = 0
    for page in pages:
        print(f"\n[{page['confidence']}] {page['stem']}")
        if repair_page(page, out_dir):
            ok += 1
        else:
            failed += 1

    print(f"\nDone: {ok} ok, {failed} failed")
    return 1 if failed else 0


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