from pathlib import Path

import click

from loguru import logger

import png
import zlib

@click.command()
@click.argument("pngfile", type=click.File('rb'))
def detach(pngfile):
    reader = png.Reader(file=pngfile)
    data_cur = 0
    data_buf = None
    compressed = False
    attachment_name = None
    decoded_attachment = None
    for chunk_type, data in reader.chunks():
        if chunk_type.decode() == "atCh":
            fn_bytes = bytearray()
            for i, b in enumerate(data):
                if b != 0:
                    fn_bytes.append(b)
                else:
                    data_cur = i+1
                    break
            attachment_name = fn_bytes.decode()
            compressed = data[data_cur] == 1
            data_buf = data[data_cur+1:]
    if attachment_name:
        logger.success(f"Found attachment named {attachment_name}. {'Not compressed.' if not compressed else 'Compressed.'}")
    else:
        logger.success(f"No attachment found.")
        return

    def inflate(buf):
        decompress = zlib.decompressobj()
        inflated = decompress.decompress(buf)
        inflated += decompress.flush()
        return inflated

    if compressed:
        decoded_attachment = inflate(data_buf)
    else:
        decoded_attachment = data_buf

    outpath = Path(attachment_name)

    if not outpath.exists():
        with open(outpath, 'wb') as out:
            out.write(decoded_attachment)
        logger.success(f"Decoded attachment written to {outpath}")
    else:
        logger.error(f"Cowardly refusing to write attachment to existing file {outpath}")