png_parser.py

"""
PNG Parser for Character Card extraction
Supports Character Card V2 format embedded in PNG files
"""

import struct
import zlib
import json
import base64
from pathlib import Path
from typing import Optional, Dict, Any, List
from PIL import Image

from card_models import CharacterCard, parse_character_card_v2


class PNGParser:
    """Parser for PNG files with embedded Character Card data"""

    # PNG chunk signatures
    PNG_SIGNATURE = b'\x89PNG\r\n\x1a\n'

    # Chunk types that may contain character card data
    TEXT_CHUNKS = {
        b'tEXt',  # Latin-1 text
        b'zTXt',  # Compressed text
        b'iTXt',  # International text (UTF-8)
    }

    # Chunk types that carry binary image data and should never contain
    # character-card JSON.  Skipping them in the fallback scan avoids
    # trying to decode megabytes of compressed pixel data as UTF-8/Latin-1.
    BINARY_CHUNKS = {
        b'IDAT',  # compressed image data (may be megabytes)
        b'IEND',  # end-of-file marker, no data
        b'IHDR',  # image header, 13 bytes of integers
        b'PLTE',  # palette, raw RGB triples
        b'tRNS',  # transparency, raw bytes
        b'gAMA',  # gamma, 4-byte int
        b'cHRM',  # chromaticity, 32 bytes
        b'sRGB',  # sRGB intent, 1 byte
        b'sBIT',  # significant bits, raw bytes
        b'bKGD',  # background colour, raw bytes
        b'hIST',  # histogram, raw shorts
        b'pHYs',  # pixel dimensions, 9 bytes
        b'sPLT',  # suggested palette, binary
        b'tIME',  # image last-modification time, 7 bytes
    }

    # Known keywords for character card data
    CARD_KEYWORDS = [
        b'chara',  # Main keyword for SillyTavern cards
    ]

    def __init__(self, file_path: Path):
        self.file_path = file_path
        self.chunks = []
        self.card_data = None

    def parse(self) -> Optional[CharacterCard]:
        """
        Parse PNG file and extract Character Card data

        Returns:
            CharacterCard object or None if not found
        """
        try:
            with open(self.file_path, 'rb') as f:
                # Check PNG signature
                signature = f.read(8)
                if signature != self.PNG_SIGNATURE:
                    print(f"Not a valid PNG file: {self.file_path}")
                    return None

                # Read chunks
                self.chunks = []
                while True:
                    chunk_data = self._read_chunk(f)
                    if chunk_data is None:
                        break
                    self.chunks.append(chunk_data)

                    # Check for IEND (end of file)
                    if chunk_data[0] == b'IEND':
                        break

                # Extract character card data from chunks
                card_json = self._extract_card_data()
                if card_json:
                    return parse_character_card_v2(card_json, self.file_path)

                return None

        except Exception as e:
            print(f"Error parsing PNG {self.file_path}: {e}")
            return None

    # Maximum chunk size we will read into memory.  A legitimate PNG chunk
    # rarely exceeds a few MB; 64 MB is a generous cap that still prevents
    # a corrupt/malicious length field from allocating gigabytes of RAM.
    MAX_CHUNK_SIZE = 64 * 1024 * 1024  # 64 MB

    def _read_chunk(self, f) -> Optional[tuple]:
        """
        Read a PNG chunk

        Returns:
            Tuple of (chunk_type, chunk_data) or None at EOF
        """
        # Chunk length (4 bytes)
        length_bytes = f.read(4)
        if len(length_bytes) < 4:
            return None

        length = struct.unpack('>I', length_bytes)[0]

        # Chunk type (4 bytes)
        chunk_type = f.read(4)
        if len(chunk_type) < 4:
            return None

        # Guard against corrupt / malicious length values.
        # Skip the oversized chunk data + CRC instead of aborting the whole
        # parse — a large chunk before 'chara' would otherwise hide the card.
        if length > self.MAX_CHUNK_SIZE:
            f.seek(length + 4, 1)  # skip data + CRC
            return (chunk_type, None)  # caller checks for None data

        # Chunk data
        chunk_data = f.read(length)

        # CRC (4 bytes) - skip
        f.read(4)

        return (chunk_type, chunk_data)

    def _extract_card_data(self) -> Optional[Dict[str, Any]]:
        """
        Extract Character Card JSON from PNG chunks

        Returns:
            Dictionary with card data or None
        """
        # Method 1: Look for known keywords in tEXt/zTXt/iTXt chunks
        for chunk_type, chunk_data in self.chunks:
            if chunk_data is None:  # oversized chunk was skipped
                continue
            if chunk_type in self.TEXT_CHUNKS:
                card_data = self._parse_text_chunk(chunk_type, chunk_data)
                if card_data:
                    return card_data

        # Method 2: Look for JSON-like content in any non-binary chunk.
        # FIX #6: skip known binary chunk types (IDAT, IHDR, PLTE, etc.)
        # to avoid decoding megabytes of compressed pixel data as text.
        for chunk_type, chunk_data in self.chunks:
            if chunk_data is None:  # oversized chunk was skipped
                continue
            if chunk_type in self.BINARY_CHUNKS:
                continue
            if chunk_type in self.TEXT_CHUNKS:
                continue  # already handled in Method 1
            if self._looks_like_character_card(chunk_data):
                try:
                    if isinstance(chunk_data, bytes):
                        try:
                            text = chunk_data.decode('utf-8')
                        except UnicodeDecodeError:
                            text = chunk_data.decode('latin-1')
                    else:
                        text = chunk_data

                    result = self._try_parse_json(text)
                    # FIX #3: verify spec field to avoid false positives from
                    # arbitrary JSON stored in custom chunks.
                    if result is not None and result.get('spec') == 'chara_card_v2':
                        return result
                except Exception:
                    pass

            # Also try base64 decode for non-text chunks — some encoders
            # store the card payload as raw base64 outside a tEXt chunk.
            if isinstance(chunk_data, bytes) and len(chunk_data) > 0:
                try:
                    decoded = base64.b64decode(chunk_data, validate=True)
                    result = self._try_parse_json(decoded)
                    if result is not None and result.get('spec') == 'chara_card_v2':
                        return result
                except Exception:
                    pass

        return None

    def _parse_text_chunk(self, chunk_type: bytes, chunk_data: bytes) -> Optional[Dict[str, Any]]:
        """
        Parse a text chunk (tEXt, zTXt, iTXt) and extract character card data

        Args:
            chunk_type: Type of chunk (tEXt, zTXt, iTXt)
            chunk_data: Raw chunk data (may be None for oversized skipped chunks)

        Returns:
            Dictionary with card data or None
        """
        if chunk_data is None:
            return None
        try:
            if chunk_type == b'tEXt':
                # Format: keyword\0text
                null_idx = chunk_data.find(b'\x00')
                if null_idx > 0:
                    keyword = chunk_data[:null_idx]
                    text = chunk_data[null_idx + 1:]

                    if self._is_card_keyword(keyword):
                        # Try base64 first (validate=True raises on bad data)
                        try:
                            decoded = base64.b64decode(text, validate=True)
                            result = self._try_parse_json(decoded)
                            if result is not None:
                                return result
                        except Exception:
                            pass
                        # Fallback: try direct JSON
                        return self._try_parse_json(text)

            elif chunk_type == b'zTXt':
                # Format: keyword\0compression_method\0compressed_text
                null_idx = chunk_data.find(b'\x00')
                if null_idx > 0:
                    keyword = chunk_data[:null_idx]
                    remaining = chunk_data[null_idx + 1:]

                    if len(remaining) > 0:
                        compressed_text = remaining[1:]  # skip compression_method byte

                        if self._is_card_keyword(keyword):
                            try:
                                text = zlib.decompress(compressed_text).decode('utf-8')
                                return self._try_parse_json(text)
                            except Exception:
                                pass

            elif chunk_type == b'iTXt':
                # FIX: parse iTXt manually instead of using split(),
                # which misaligns fields and causes parts[5] to be out of range.
                #
                # iTXt binary layout:
                #   keyword \x00
                #   compression_flag (1 byte: 0=uncompressed, 1=compressed)
                #   compression_method (1 byte: always 0 = deflate)
                #   language_tag \x00
                #   translated_keyword \x00
                #   text
                result = self._parse_itxt_chunk(chunk_data)
                if result is not None:
                    return result

        except Exception as e:
            print(f"Error parsing text chunk: {e}")

        return None

    def _parse_itxt_chunk(self, chunk_data: bytes) -> Optional[Dict[str, Any]]:
        """
        Parse an iTXt chunk with correct field layout handling.

        The old split(b'\\x00', 5) approach failed when compression_flag=1
        because the compression_method byte (b'\\x00') was consumed as a
        separator, shifting all subsequent parts by one index.
        """
        try:
            # Find keyword (everything before first \x00)
            null_idx = chunk_data.find(b'\x00')
            if null_idx < 0:
                return None

            keyword = chunk_data[:null_idx]
            if not self._is_card_keyword(keyword):
                return None

            pos = null_idx + 1

            # compression_flag: 1 byte
            if pos >= len(chunk_data):
                return None
            compression_flag = chunk_data[pos]
            pos += 1

            # compression_method: 1 byte (always 0, skip)
            if pos >= len(chunk_data):
                return None
            pos += 1

            # language_tag: up to next \x00
            lang_end = chunk_data.find(b'\x00', pos)
            if lang_end < 0:
                return None
            pos = lang_end + 1

            # translated_keyword: up to next \x00
            trans_end = chunk_data.find(b'\x00', pos)
            if trans_end < 0:
                return None
            pos = trans_end + 1

            # text: everything remaining
            text_bytes = chunk_data[pos:]

            if compression_flag == 1:
                try:
                    text = zlib.decompress(text_bytes).decode('utf-8')
                except Exception:
                    text = text_bytes.decode('utf-8', errors='ignore')
            else:
                text = text_bytes.decode('utf-8', errors='ignore')

            return self._try_parse_json(text)

        except Exception as e:
            print(f"Error parsing iTXt chunk: {e}")
            return None

    def _is_card_keyword(self, keyword: bytes) -> bool:
        """Check if keyword is related to character cards"""
        keyword_lower = keyword.lower()
        for card_kw in self.CARD_KEYWORDS:
            if card_kw in keyword_lower:
                return True
        return False

    def _try_parse_json(self, data) -> Optional[Dict[str, Any]]:
        """Try to parse data as JSON. Accepts both str and bytes."""
        if isinstance(data, str):
            try:
                return json.loads(data)
            except json.JSONDecodeError:
                return None

        # bytes path: try UTF-8 first, then latin-1
        try:
            return json.loads(data.decode('utf-8'))
        except (UnicodeDecodeError, json.JSONDecodeError):
            try:
                return json.loads(data.decode('latin-1'))
            except Exception:
                return None

    def _looks_like_character_card(self, data: bytes) -> bool:
        """Check if data looks like a character card JSON"""
        try:
            if isinstance(data, bytes):
                try:
                    text = data.decode('utf-8')
                except Exception:
                    text = data.decode('latin-1')
            else:
                text = str(data)

            indicators = [
                '"spec"',
                '"chara_card_v2"',
                '"description"',
                '"first_mes"',
                '"personality"',
            ]

            text_lower = text.lower()
            matches = sum(1 for ind in indicators if ind in text_lower)
            return matches >= 2

        except Exception:
            return False

    def get_image(self) -> Optional[Image.Image]:
        """Load the PNG image using PIL"""
        try:
            return Image.open(self.file_path)
        except Exception as e:
            print(f"Error loading image: {e}")
            return None


def extract_character_card(file_path: Path) -> Optional[CharacterCard]:
    """
    Extract Character Card from PNG file

    Args:
        file_path: Path to the PNG file

    Returns:
        CharacterCard object or None
    """
    parser = PNGParser(file_path)
    return parser.parse()


def scan_directory(directory: Path, recursive: bool = True) -> List[CharacterCard]:
    """
    Scan directory for Character Card PNG files

    Args:
        directory: Directory to scan
        recursive: Whether to scan subdirectories

    Returns:
        List of CharacterCard objects
    """
    cards = []

    pattern = '**/*.png' if recursive else '*.png'

    for png_file in directory.glob(pattern):
        if png_file.is_file():
            card = extract_character_card(png_file)
            if card:
                cards.append(card)
                print(f"Found: {card.data.name} ({png_file.name})")
            else:
                print(f"No card data found in: {png_file.name}")

    return cards