from __future__ import annotations
from typing import Any
from pathlib import Path
import io
import pymupdf
from PIL import Image
from .metrics import compute_all_metrics_from_pil

MetricsOpt = tuple[float | None, float | None, float | None, float | None, float | None]
LOSSLESS_FALLBACK_EXTS = {"jpx", "jp2", "j2k", "jpf", "jpm"}
DIRECT_SAVE_EXTS = {"png", "jpg", "jpeg", "tif", "tiff", "bmp", "gif", "webp"}


def _decode_with_pillow(encoded_bytes: bytes) -> Image.Image | None:
    try:
        with Image.open(io.BytesIO(encoded_bytes)) as pil:
            pil.load()
            return pil.copy()
    except Exception:
        return None


def _decode_with_pymupdf(doc: pymupdf.Document, xref: int) -> Image.Image | None:
    if xref <= 0:
        return None

    try:
        pix = pymupdf.Pixmap(doc, xref)
    except Exception:
        return None

    try:
        needs_rgb = (
            pix.colorspace is None
            or pix.colorspace.n not in (1, 3)
            or pix.alpha
        )
        if needs_rgb:
            pix = pymupdf.Pixmap(pymupdf.csRGB, pix)
        return Image.open(io.BytesIO(pix.tobytes("png"))).copy()
    except Exception:
        return None


def _decode_for_processing(encoded_bytes: bytes, doc: pymupdf.Document, xref: int) -> Image.Image | None:
    pil = _decode_with_pillow(encoded_bytes)
    if pil is not None:
        return pil
    return _decode_with_pymupdf(doc, xref)


def _compute_metrics_from_pil(pil: Image.Image | None) -> MetricsOpt:
    if pil is None:
        return (None, None, None, None, None)

    try:
        sh, hi, co, cf, cr = compute_all_metrics_from_pil(pil)
        return (sh, hi, co, cf, cr)
    except Exception:
        return (None, None, None, None, None)


def _encode_png_from_pil(pil: Image.Image) -> bytes:
    buf = io.BytesIO()
    pil.save(buf, format="PNG")
    return buf.getvalue()


def _safe_dpi(px: int, inches: float) -> float | None:
    if px <= 0 or inches <= 0.0:
        return None
    return round(px / inches, 1)


def extract_all_images(page: pymupdf.Page, page_index: int, out_dir: Path) -> list[dict[str, Any]]:
    out_dir.mkdir(parents=True, exist_ok=True)
    results: list[dict[str, Any]] = []

    raw = page.get_text("rawdict")
    blocks: list[dict[str, Any]] = raw.get("blocks", [])

    occ = 0
    for b in blocks:
        if b.get("type") != 1:
            continue

        bbox = b.get("bbox")
        if not bbox or len(bbox) != 4:
            continue
        x0, y0, x1, y1 = (float(bbox[0]), float(bbox[1]), float(bbox[2]), float(bbox[3]))
        w_pt = max(0.0, x1 - x0)
        h_pt = max(0.0, y1 - y0)
        w_in = (w_pt / 72.0) if w_pt > 0 else 0.0
        h_in = (h_pt / 72.0) if h_pt > 0 else 0.0

        w_px = int(b.get("width", 0))
        h_px = int(b.get("height", 0))
        dx = _safe_dpi(w_px, w_in)
        dy = _safe_dpi(h_px, h_in)

        data: bytes = b.get("image", b"")
        ext = (b.get("ext") or "png").lower()
        xref = int(b.get("xref", 0) or 0)
        if xref > 0:
            try:
                rec = page.parent.extract_image(xref)
                data = rec.get("image", data) or data
                ext = (rec.get("ext") or ext).lower()
            except Exception:
                pass

        if not data:
            continue

        occ += 1

        base_id = f"p{page_index}_Im{occ}"
        pil_for_processing = _decode_for_processing(data, page.parent, xref)

        save_as_png = ext in LOSSLESS_FALLBACK_EXTS or ext not in DIRECT_SAVE_EXTS
        if save_as_png and pil_for_processing is not None:
            out_name = f"{base_id}.png"
            out_bytes = _encode_png_from_pil(pil_for_processing)
        else:
            out_name = f"{base_id}.{ext}"
            out_bytes = data

        out_path = out_dir / out_name
        out_path.write_bytes(out_bytes)

        metrics = _compute_metrics_from_pil(pil_for_processing)

        results.append({
            "page": page_index,
            "image_id": base_id,
            "bbox_pt": (x0, y0, x1, y1),
            "pixels": (w_px, h_px),
            "dpi_pair": (dx, dy),
            "saved": str(out_path),
            "metrics": metrics,
        })

    return results
