"""Extract images and surrounding text blocks from PDFs using pymupdf."""

import argparse
import base64
import io
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import fitz  # pymupdf
from PIL import Image as PILImage

from output_utils import output_to_file_or_stdout


def extract(pdf_path: str, include_b64: bool = False) -> list[dict]:
    """Return list of image dicts with page, bbox, surrounding text, and optional base64."""
    results = []
    with fitz.open(pdf_path) as doc:
        for i, page in enumerate(doc, start=1):
            images = page.get_images(full=True)
            text_blocks = page.get_text("blocks")  # (x0, y0, x1, y1, text, block_no, block_type)

            for img_idx, img_info in enumerate(images):
                xref = img_info[0]
                base_image = doc.extract_image(xref)
                if not base_image:
                    continue

                # Find image bbox on the page
                img_rects = page.get_image_rects(img_info)
                if not img_rects:
                    continue
                bbox = img_rects[0]

                # Find nearby text blocks (within 50pt vertical proximity)
                nearby_text = []
                for block in text_blocks:
                    if block[6] != 0:  # skip image blocks
                        continue
                    block_rect = fitz.Rect(block[:4])
                    if abs(block_rect.y0 - bbox.y1) < 50 or abs(bbox.y0 - block_rect.y1) < 50:
                        nearby_text.append(block[4].strip())

                entry = {
                    "page": i,
                    "image_index": img_idx + 1,
                    "bbox": [round(c, 1) for c in bbox],
                    "width": base_image["width"],
                    "height": base_image["height"],
                    "format": base_image["ext"],
                    "nearby_text": nearby_text,
                }

                if include_b64:
                    entry["base64"] = base64.b64encode(base_image["image"]).decode("ascii")

                results.append(entry)
    return results


def _images_to_text(images: list[dict]) -> list[dict]:
    """Convert image info into page-text dicts for txt/md output."""
    page_texts: dict[int, list[str]] = {}
    for img in images:
        lines = [
            f"Image {img['image_index']}: {img['width']}x{img['height']} ({img['format']})",
            f"  bbox: {img['bbox']}",
        ]
        if img["nearby_text"]:
            lines.append("  nearby text:")
            for t in img["nearby_text"]:
                lines.append(f"    {t}")
        page_texts.setdefault(img["page"], []).append("\n".join(lines))
    return [{"page": p, "text": "\n\n".join(texts)} for p, texts in sorted(page_texts.items())]


def main():
    parser = argparse.ArgumentParser(description="Extract images from PDFs")
    parser.add_argument("pdf", help="Path to input PDF")
    parser.add_argument("output", nargs="?", default=None, help="Output file path")
    parser.add_argument("--format", choices=["json", "md", "txt"], default="json",
                        dest="fmt", help="Output format (default: json)")
    parser.add_argument("--base64", action="store_true", help="Include base64-encoded image data")
    args = parser.parse_args()

    images = extract(args.pdf, include_b64=args.base64)

    if args.fmt == "json":
        data = {"extractor": "images", "images": images}
        output_to_file_or_stdout(data, "json", args.output)
    else:
        text_pages = _images_to_text(images)
        output_to_file_or_stdout(text_pages, args.fmt, args.output)


if __name__ == "__main__":
    main()
