#  zeromodel/png_text.py
"""Utilities for reading and writing PNG text chunks.

This module provides minimal helpers to insert, remove and query text chunks
within PNG images without relying on external libraries.
"""
from __future__ import annotations

import struct
import zlib
from typing import Iterator, Optional, Tuple

_PNG_SIG = b"\x89PNG\r\n\x1a\n"


def _crc32(chunk_type: bytes, data: bytes) -> int:
    return zlib.crc32(chunk_type + data) & 0xFFFFFFFF


def _build_chunk(chunk_type: bytes, data: bytes) -> bytes:
    return (
        struct.pack(">I", len(data))
        + chunk_type
        + data
        + struct.pack(">I", _crc32(chunk_type, data))
    )


def _iter_chunks(png: bytes) -> Iterator[Tuple[bytes, int, int, bytes]]:
    """Iterate over PNG chunks.

    Yields tuples of (type, start_offset, end_offset, data). ``start_offset``
    points to the 4-byte length field of the chunk and ``end_offset`` points
    right after the CRC (i.e. the start of the next chunk).
    """
    if not png.startswith(_PNG_SIG):
        raise ValueError("Not a PNG: bad signature")
    i = len(_PNG_SIG)
    n = len(png)
    while i + 8 <= n:
        length = struct.unpack(">I", png[i : i + 4])[0]
        ctype = png[i + 4 : i + 8]
        data_start = i + 8
        data_end = data_start + length
        crc_end = data_end + 4
        if crc_end > n:
            break  # Truncated/corrupt; stop parsing gracefully
        data = png[data_start:data_end]
        yield ctype, i, crc_end, data
        i = crc_end
        if ctype == b"IEND":
            break


def _find_iend_offset(png: bytes) -> int:
    # Return byte offset where IEND chunk starts; insert before this
    for ctype, start, _end, _data in _iter_chunks(png):
        if ctype == b"IEND":
            return start
    raise ValueError("PNG missing IEND chunk")


def _remove_text_chunks_with_key(png: bytes, key: str) -> bytes:
    """Remove existing iTXt/tEXt/zTXt chunks that match ``key``."""
    key_bytes = key.encode("latin-1", "ignore")
    out = bytearray()
    out.extend(png[: len(_PNG_SIG)])
    pos = len(_PNG_SIG)
    for ctype, start, end, data in _iter_chunks(png):
        out.extend(png[pos:start])  # normally empty
        keep = True
        if ctype in (b"tEXt", b"iTXt", b"zTXt"):
            try:
                nul = data.find(b"\x00")
                k = data[:nul] if nul != -1 else b""
            except Exception:
                k = b""
            if k == key_bytes:
                keep = False
        if keep:
            out.extend(png[start:end])
        pos = end
    out.extend(png[pos:])
    return bytes(out)


def _encode_text_chunk(
    key: str, text: str, use_itxt: bool = True, *, compress: bool = False
) -> bytes:
    """Build a tEXt or iTXt chunk."""
    if use_itxt:
        # iTXt layout:
        # keyword\0 compression_flag(1) compression_method(1)
        # language_tag\0 translated_keyword\0 text(UTF-8)
        keyword = key.encode("latin-1", "ignore")[:79]
        comp_flag = b"\x01" if compress else b"\x00"
        comp_method = b"\x00"  # zlib
        language_tag = b""  # empty
        translated_keyword = b""  # empty
        text_bytes = text.encode("utf-8", "strict")
        if compress:
            text_bytes = zlib.compress(text_bytes)
        data = (
            keyword
            + b"\x00"
            + comp_flag
            + comp_method
            + language_tag
            + b"\x00"
            + translated_keyword
            + b"\x00"
            + text_bytes
        )
        return _build_chunk(b"iTXt", data)
    else:
        # tEXt: keyword\0 text (both Latin-1)
        keyword = key.encode("latin-1", "ignore")[:79]
        text_bytes = text.encode("latin-1", "replace")
        data = keyword + b"\x00" + text_bytes
        return _build_chunk(b"tEXt", data)


def _decode_text_chunk(
    ctype: bytes, data: bytes
) -> Tuple[Optional[str], Optional[str]]:
    """Returns (keyword, text) from a tEXt/iTXt/zTXt chunk."""
    try:
        if ctype == b"tEXt":
            nul = data.find(b"\x00")
            if nul == -1:
                return (None, None)
            key = data[:nul].decode("latin-1", "ignore")
            txt = data[nul + 1 :].decode("latin-1", "ignore")
            return key, txt
        if ctype == b"iTXt":
            p = 0
            nul = data.find(b"\x00", p)
            if nul == -1:
                return (None, None)
            key = data[p:nul]
            p = nul + 1
            comp_flag = data[p]
            p += 1
            comp_method = data[p]
            p += 1
            nul = data.find(b"\x00", p)
            if nul == -1:
                return (None, None)
            lang = data[p:nul]
            p = nul + 1
            nul = data.find(b"\x00", p)
            if nul == -1:
                return (None, None)
            trkw = data[p:nul]
            p = nul + 1
            txt_bytes = data[p:]
            if comp_flag == 1:
                txt_bytes = zlib.decompress(txt_bytes)
            key = key.decode("latin-1", "ignore")
            txt = txt_bytes.decode("utf-8", "ignore")
            return key, txt
        if ctype == b"zTXt":
            nul = data.find(b"\x00")
            if nul == -1 or len(data) < nul + 2:
                return (None, None)
            key = data[:nul].decode("latin-1", "ignore")
            comp_method = data[nul + 1]
            comp = data[nul + 2 :]
            if comp_method != 0:
                return key, None
            txt = zlib.decompress(comp).decode("latin-1", "ignore")
            return key, txt
    except Exception:
        pass
    return (None, None)


def png_read_text_chunk(png_bytes: bytes, key: str) -> Optional[str]:
    """Read the text value for ``key`` from iTXt/tEXt/zTXt chunks."""
    if not png_bytes.startswith(_PNG_SIG):
        raise ValueError("Not a PNG")
    want_key = key
    found_text = None
    itxt_text = None
    for ctype, _s, _e, data in _iter_chunks(png_bytes):
        if ctype in (b"tEXt", b"iTXt", b"zTXt"):
            k, v = _decode_text_chunk(ctype, data)
            if k == want_key and v is not None:
                if ctype == b"iTXt":
                    itxt_text = v  # prefer iTXt
                elif found_text is None:
                    found_text = v
    return itxt_text if itxt_text is not None else found_text


def png_write_text_chunk(
    png_bytes: bytes,
    key: str,
    text: str,
    *,
    use_itxt: bool = True,
    compress: bool = False,
    replace_existing: bool = True,
) -> bytes:
    """Insert (or replace) a text chunk with ``(key, text)``.

    ``use_itxt`` selects the iTXt chunk (UTF-8 capable). ``compress`` applies
    zlib compression to the iTXt payload. ``replace_existing`` removes any
    prior chunks with the same key.
    """
    if not png_bytes.startswith(_PNG_SIG):
        raise ValueError("Not a PNG")
    png2 = (
        _remove_text_chunks_with_key(png_bytes, key) if replace_existing else png_bytes
    )
    new_chunk = _encode_text_chunk(key, text, use_itxt=use_itxt, compress=compress)
    iend_off = _find_iend_offset(png2)
    out = png2[:iend_off] + new_chunk + png2[iend_off:]
    return out


__all__ = ["png_read_text_chunk", "png_write_text_chunk"]

