# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

# src/pdftl/operations/images_roundtrip.py

import hashlib
import json
import logging
import sys
from pathlib import Path

import pdftl.core.constants as c
from pdftl.core.core_types import OpResult
from pdftl.core.registry import register_operation
from pdftl.exceptions import InvalidArgumentError
from pdftl.utils.hooks import from_result_meta
from pdftl.utils.io_helpers import smart_open_maybe_dash
from pdftl.utils.keyval_parser import parse_keyval_list
from pdftl.utils.page_specs import page_numbers_matching_page_specs, is_valid_page_spec
from pdftl.utils.string_utils import compact_json_string
from pdftl.utils.images.selectors import extract_to_pil
from pdftl.operations.helpers.image_processor import (
    encode_and_update_pdf_image,
    get_orig_stream_size,
    ImageContext,
)

logger = logging.getLogger(__name__)


# ============================================================================
# User Documentation & Examples
# ============================================================================

_EXPORT_IMAGES_LONG_DESC = """
The `export_images` operation extracts all embedded images from the PDF, saves them to a
specified directory, and generates a structured JSON manifest file.

This allows you to bulk-export images, edit them in external software (like GIMP or
Photoshop), and then easily re-integrate them back into the PDF using the `import_images`
command.

The JSON manifest groups images by their underlying stream data. This means that if a
single image is drawn multiple times across different pages, it will only be exported once.

### Arguments

* `<directory>`: **(Required)** The directory where the extracted images will be saved.
* `[output <manifest.json>]`: The file path to save the JSON manifest. If omitted, the
  manifest is automatically saved as `manifest.json` inside the specified directory.
  Specify `output -` to print the manifest directly to standard output.
* `[page_range]`: Optional page ranges to limit the extraction.
"""

_EXPORT_IMAGES_EXAMPLES = [
    {
        "cmd": "in.pdf export_images my_images",
        "desc": (
            "Extract all images to 'my_images/' and save the manifest to "
            "'my_images/manifest.json'."
        ),
    },
    {
        "cmd": "in.pdf export_images 1-5 edited_assets output manifest.json",
        "desc": (
            "Extract images from pages 1-5 into 'edited_assets/' and write "
            "manifest to 'manifest.json'."
        ),
    },
]

_IMPORT_IMAGES_LONG_DESC = """
The `import_images` operation reads a JSON manifest (generated by `export_images`) and
overwrites the corresponding internal PDF image streams with the edited files from a
directory.

If an image file in the directory has not been modified (its MD5 hash matches the original
export hash), it is safely skipped to prevent unnecessary re-encoding and generation loss.

### Arguments

* `<directory>`: **(Required)** The directory containing the edited images.
* `[manifest=<manifest.json>]`: Path to the JSON manifest. If omitted, defaults to
  `<directory>/manifest.json`. Specify `manifest=-` to read the manifest from standard input.
* `[quality=<q>]`: The JPEG compression quality (1-100) used when writing back lossy
  images. (Default: 75)
"""

_IMPORT_IMAGES_EXAMPLES = [
    {
        "cmd": "in.pdf import_images my_images output out.pdf",
        "desc": (
            "Import edited images from 'my_images/' using 'my_images/manifest.json' and save."
        ),
        "test_setup": {
            "mkdirs": ["my_images"],
            "create_files": {"my_images/manifest.json": '{ "images": [] }'},
        },
    },
    {
        "cmd": "in.pdf import_images edited_assets manifest=custom_manifest.json output out.pdf",
        "desc": "Import edited images using a custom manifest location.",
        "test_setup": {
            "mkdirs": ["edited_assets"],
            "create_files": {"custom_manifest.json": '{ "images": [] }'},
        },
    },
]


# ============================================================================
# Shared Private Helpers
# ============================================================================


def _file_hash(filepath: Path) -> str:
    """Calculates the MD5 hash of a file to track modifications."""
    hasher = hashlib.md5()
    try:
        with open(filepath, "rb") as f:
            buf = f.read()
            hasher.update(buf)
    except OSError as exc:
        # Explanatory comment: Return empty hash if the file is unreadable,
        # which will safely force a re-import if attempted.
        logger.debug("Failed to calculate hash for %s: %s", filepath, exc)
        return ""
    return hasher.hexdigest()


def _classify_tokens(tokens: list[str]) -> tuple[str | None, list[str]]:
    """Separates the directory argument from page specifications based on typical naming
    formats."""
    if not tokens:
        return None, []
    if len(tokens) == 1:
        return tokens[0], []

    directory = None
    page_specs = []
    for t in reversed(tokens):  # last dir-like takes priority
        if directory is None and not is_valid_page_spec(t):
            directory = t
        else:
            page_specs.append(t)

    # Fallback if everything matched the page spec pattern
    if directory is None and page_specs:
        directory = page_specs.pop(0)

    return directory, page_specs


# ============================================================================
# Export Images Pipeline Logic
# ============================================================================


def _get_target_pages(pdf, page_specs: list[str]) -> list[int]:
    """Resolves target page numbers from raw page specifications."""
    num_pages = len(pdf.pages)
    if page_specs:
        return sorted(list(page_numbers_matching_page_specs(page_specs, num_pages)))
    return list(range(1, num_pages + 1))


def _export_single_image(xobj, img: dict, out_dir: Path) -> tuple[str, str] | None:
    """Extracts a PDF XObject to a PIL Image and saves it to disk."""
    pil_img = extract_to_pil(xobj)
    if pil_img is None:
        logger.warning("Could not extract PIL image for object %s", xobj.objgen)
        return None

    fmt = img.get("format", "unknown")
    ext = "jpg" if fmt == "dctdecode" else "png"
    obj_id = f"{xobj.objgen[0]}_{xobj.objgen[1]}"
    filename = f"obj_{obj_id}.{ext}"
    filepath = out_dir / filename

    try:
        if pil_img.mode == "CMYK" and ext == "png":
            save_img = pil_img.convert("RGB")
        else:
            save_img = pil_img
        save_img.save(filepath)
    except (OSError, ValueError) as exc:
        # Explanatory comment: Gracefully handle Pillow save and file I/O failures.
        logger.warning("Failed to save image %s: %s", filepath, exc)
        return None

    return filename, _file_hash(filepath)


def _build_manifest(pdf, target_pages: list[int], out_dir: Path) -> dict:
    """Iterates over embedded PDF images and builds the manifest dictionary."""
    from pdftl.utils.images.finders import extract_pdf_images

    image_data = extract_pdf_images(pdf, target_pages)
    manifest: dict = {"image_streams": {}}

    for img in image_data:
        xobj = img.pop("xobj")

        # Inline Collision Guard: Skip direct objects to prevent mapping crashes
        if not xobj.is_indirect:
            logger.warning(
                "Skipping inline/direct image object on page %s. "
                "Only indirect image objects are supported.",
                img.get("page", "?"),
            )
            continue

        obj_id = f"{xobj.objgen[0]}_{xobj.objgen[1]}"

        placement = {
            "page": img["page"],
            "name": img["name"],
            "bbox": img["bbox"],
            "ppi_x": img["ppi_x"],
            "ppi_y": img["ppi_y"],
        }

        if obj_id not in manifest["image_streams"]:
            export_info = _export_single_image(xobj, img, out_dir)
            if export_info is None:
                continue

            filename, file_hash = export_info
            manifest["image_streams"][obj_id] = {
                "format": img.get("format", "unknown"),
                "width_px": img["width_px"],
                "height_px": img["height_px"],
                "colorspace": img["colorspace"],
                "bits": img["bits"],
                "export_file": filename,
                "file_hash": file_hash,
                "placements": [placement],
            }
        else:
            manifest["image_streams"][obj_id]["placements"].append(placement)

    return manifest


def export_images_cli_hook(result: OpResult, stage, _pipeline):
    """Writes the manifest data to stdout or a file in JSON."""
    if not result.success:
        return

    output_file = from_result_meta(result, c.META_OUTPUT_FILE)

    with smart_open_maybe_dash(output_file) as file:
        json_string = compact_json_string(json.dumps(result.data, indent=2), fold_dicts=False)
        file.write(json_string + "\n")


@register_operation(
    "export_images",
    tags=["info", "metadata", "images", "export"],
    type="single input operation",
    desc="Export images and a JSON manifest for external editing",
    long_desc=_EXPORT_IMAGES_LONG_DESC,
    examples=_EXPORT_IMAGES_EXAMPLES,
    cli_hook=export_images_cli_hook,
    usage="<input> export_images [<spec>...] <dir> [output <output>]",
    args=([c.INPUT_PDF, c.OPERATION_ARGS], {"output_file": c.OUTPUT}),
    skip_pipeline_save=True,
)
def export_images(pdf, specs, output_file=None) -> OpResult:
    """Export embedded images to a directory with a JSON manifest."""
    raw_tokens = []
    parse_keyval_list(
        specs or [],
        bare_tokens=raw_tokens,
        allowed_keys=[],
        context="export_images",
    )

    out_dir_str, page_specs = _classify_tokens(raw_tokens)
    if not out_dir_str:
        raise InvalidArgumentError("Missing required directory argument.")

    out_dir = Path(out_dir_str)
    out_dir.mkdir(parents=True, exist_ok=True)

    target_pages = _get_target_pages(pdf, page_specs)
    manifest = _build_manifest(pdf, target_pages, out_dir)

    resolved_output = output_file if output_file is not None else str(out_dir / "manifest.json")

    return OpResult(
        success=True,
        data=manifest,
        meta={c.META_OUTPUT_FILE: resolved_output},
    )


# ============================================================================
# Import Images Pipeline Logic
# ============================================================================


def _parse_quality(quality_str: str) -> int:
    """Parses and validates JPEG compression quality bounds."""
    try:
        quality = int(quality_str)
        if not (1 <= quality <= 100):
            raise ValueError
        return quality
    except ValueError as exc:
        raise InvalidArgumentError(
            f"Invalid quality: '{quality_str}'. Must be an integer between 1 and 100."
        ) from exc


def _read_manifest(manifest_path: Path) -> dict:
    """Loads JSON manifest from the file system or standard input."""
    if str(manifest_path) == "-":
        try:
            return json.load(sys.stdin)
        except json.JSONDecodeError as exc:
            raise InvalidArgumentError(f"Invalid JSON manifest from stdin: {exc}") from exc

    if not manifest_path.is_file():
        raise InvalidArgumentError(f"Manifest file not found: {manifest_path}")

    try:
        with open(manifest_path, encoding="utf-8") as f:
            return json.load(f)
    except json.JSONDecodeError as exc:
        raise InvalidArgumentError(f"Invalid JSON manifest: {exc}") from exc


def _get_pdf_image_objects(pdf) -> dict:
    """Crawls active internal PDF stream resources to index target XObjects."""
    import pikepdf

    objgen_map = {}
    for obj in pdf.objects:
        if (
            isinstance(obj, pikepdf.Stream)
            and obj.get(pikepdf.Name.Type) == pikepdf.Name.XObject
            and obj.get(pikepdf.Name.Subtype) == pikepdf.Name.Image
        ):
            obj_id = f"{obj.objgen[0]}_{obj.objgen[1]}"
            objgen_map[obj_id] = obj
    return objgen_map


def _check_dimensions(filename: str, pil_img, stream_info: dict) -> None:
    """Checks for dimension shifts that might cause layout distortions."""
    orig_w = stream_info.get("width_px")
    orig_h = stream_info.get("height_px")
    if orig_w and orig_h and (pil_img.width != orig_w or pil_img.height != orig_h):
        logger.warning(
            "Imported image %s dimensions (%dx%d) differ from original (%dx%d). "
            "This may cause visual distortion in the PDF layout.",
            filename,
            pil_img.width,
            pil_img.height,
            orig_w,
            orig_h,
        )


def _resolve_import_filepath(filename: str, out_dir: Path) -> Path | None:
    """Checks and resolves format suffix swaps (e.g. .jpg -> .png conversions)."""
    # Security: Prevent Path Traversal attacks via custom filenames inside the manifest.
    try:
        filepath = Path(out_dir / filename).resolve()
    except (OSError, ValueError) as exc:
        logger.debug("Failed to resolve filename path %s: %s", filename, exc)
        return None

    if not filepath.is_relative_to(out_dir.resolve()):
        raise InvalidArgumentError(
            f"Security violation: path traversal detected in filename '{filename}'"
        )

    if filepath.is_file():
        return filepath

    # Extended Extension Fallback Cycle
    supported_exts = {".png", ".jpg", ".jpeg", ".webp"}
    for ext in supported_exts:
        alt_filepath = filepath.with_suffix(ext)
        if alt_filepath.is_file():
            logger.info("Using alternative format file: %s", alt_filepath.name)
            return alt_filepath

    logger.debug("Skipping %s, file not found in directory.", filename)
    return None


def _import_single_image(
    obj_id: str,
    stream_info: dict,
    out_dir: Path,
    objgen_map: dict,
    quality: int,
) -> bool:
    """Re-encodes and updates a single image stream from the manifest inside the PDF."""
    from PIL import Image

    filename = stream_info.get("export_file")
    if not filename:
        return False

    filepath = _resolve_import_filepath(filename, out_dir)
    if not filepath:
        return False

    current_hash = _file_hash(filepath)
    if current_hash == stream_info.get("file_hash"):
        logger.debug("Skipping %s, file hash unchanged.", filename)
        return False

    xobj = objgen_map.get(obj_id)
    if not xobj:
        logger.warning("Target stream %s from manifest not found in PDF.", obj_id)
        return False

    try:
        pil_img = Image.open(filepath)
        pil_img.load()
    except (OSError, ValueError) as exc:
        # Explanatory comment: Gracefully log and bypass unopenable/corrupt edited files.
        logger.warning("Failed to read image %s: %s", filepath, exc)
        return False

    _check_dimensions(filepath.name, pil_img, stream_info)

    # Resolve real original stream context metadata to satisfy encoder constraints
    placements = stream_info.get("placements", [])
    page_num = placements[0].get("page", 0) if placements else 0

    ctx = ImageContext(
        xobj=xobj,
        smask_xobj=xobj.get("/SMask"),
        orig_size=get_orig_stream_size(xobj),
        img_dict={"name": filepath.name},
        page_num=page_num,
    )

    try:
        encode_and_update_pdf_image(ctx, pil_img, quality)
        logger.info("Imported and updated image stream %s from %s", obj_id, filepath.name)
        return True
    except (OSError, ValueError) as exc:
        # Explanatory comment: Log encoder mapping or layout write errors gracefully.
        logger.warning("Failed to encode and update PDF image stream %s: %s", obj_id, exc)
        return False


@register_operation(
    "import_images",
    tags=["in_place", "images", "import", "replace"],
    type="single input operation",
    desc="Import edited images from a directory using a JSON manifest",
    long_desc=_IMPORT_IMAGES_LONG_DESC,
    usage="<input> import_images <dir> [manifest=<manifest.json>] [quality=<q>] output <out.pdf>",
    examples=_IMPORT_IMAGES_EXAMPLES,
    args=([c.INPUT_PDF, c.OPERATION_ARGS], {}),
)
def import_images(pdf, specs) -> OpResult:
    """Import edited images from a directory and overwrite internal PDF streams."""
    raw_tokens = []
    kv = parse_keyval_list(
        specs or [],
        bare_tokens=raw_tokens,
        allowed_keys=["manifest", "quality"],
        context="import_images",
    )

    out_dir_str, _ = _classify_tokens(raw_tokens)
    if not out_dir_str:
        raise InvalidArgumentError("Missing required directory argument.")

    out_dir = Path(out_dir_str)
    if not out_dir.is_dir():
        raise InvalidArgumentError(f"Target directory does not exist: {out_dir}")

    quality = _parse_quality(kv.get("quality", "75"))
    manifest_str = kv.get("manifest")

    manifest_path = out_dir / "manifest.json" if not manifest_str else Path(manifest_str)
    manifest = _read_manifest(manifest_path)

    image_streams = manifest.get("image_streams", {})
    if not image_streams:
        logger.info("Manifest contains no image streams. Skipping import.")
        return OpResult(success=True, pdf=pdf)

    objgen_map = _get_pdf_image_objects(pdf)
    updated_count = 0

    for obj_id, stream_info in image_streams.items():
        if _import_single_image(obj_id, stream_info, out_dir, objgen_map, quality):
            updated_count += 1

    logger.info("Successfully imported %d image(s).", updated_count)
    return OpResult(success=True, pdf=pdf)
