# -*- coding: utf-8 -*-

"""Embed and extract YAML source data in PNG metadata (iTXT chunks)."""

from pathlib import Path
from typing import Optional, Tuple

from PIL import Image
from PIL.PngImagePlugin import PngInfo


PNG_KEY_YAML = "wireviz_yaml"
PNG_KEY_PREPEND = "wireviz_prepend_yaml"


def save_yaml_to_png(
    png_path: Path,
    yaml_input: str,
    prepend_input: str = "",
) -> None:
    """Save YAML source as compressed iTXT metadata in a PNG file."""
    png_path = Path(png_path)
    if not png_path.suffix == ".png":
        png_path = png_path.with_suffix(".png")
    if not png_path.exists():
        return

    with Image.open(fp=png_path) as im:
        txt = PngInfo()
        txt.add_itxt(PNG_KEY_YAML, yaml_input, zip=True)
        if prepend_input:
            txt.add_itxt(PNG_KEY_PREPEND, prepend_input, zip=True)
        im.save(fp=png_path, pnginfo=txt)


def read_yaml_from_png(png_path: Path) -> Tuple[str, Optional[str]]:
    """Extract YAML source from a PNG file's iTXT metadata.

    Returns (yaml_input, prepend_input) where prepend_input may be None.
    """
    png_path = Path(png_path)
    if not png_path.suffix == ".png":
        png_path = png_path.with_suffix(".png")

    with Image.open(fp=png_path) as im:
        im.load()
        yaml_input = im.text.get(PNG_KEY_YAML, "")
        prepend_input = im.text.get(PNG_KEY_PREPEND)

    return yaml_input, prepend_input


def has_yaml_metadata(png_path: Path) -> bool:
    """Check if a PNG file contains embedded WireViz YAML data."""
    png_path = Path(png_path)
    if not png_path.suffix == ".png":
        png_path = png_path.with_suffix(".png")
    if not png_path.exists():
        return False

    with Image.open(fp=png_path) as im:
        im.load()
        return PNG_KEY_YAML in im.text