#!/venv/bin/python3

import struct

import png


def write_pnm(file, plain, rows, meta):
    """
    Write a Netpbm PNM (or PAM) file.
    *file* output file object;
    *plain* (a bool) true if writing plain format (not possible for PAM);
    *rows* an iterator for the rows;
    *meta* the info dictionary.
    """

    meta = dict(meta)
    meta["maxval"] = 2 ** meta["bitdepth"] - 1
    meta["width"], meta["height"] = meta["size"]

    # Number of planes determines both image formats:
    # 1 : L to PGM
    # 2 : LA to PAM
    # 3 : RGB to PPM
    # 4 : RGBA to PAM
    planes = meta["planes"]

    # Assume inputs are from a PNG file.
    assert planes in (1, 2, 3, 4)
    if planes in (1, 3):
        if 1 == planes:
            # PGM
            # Even if maxval is 1 we use PGM instead of PBM,
            # to avoid converting data.
            magic = "P5"
            if plain:
                magic = "P2"
        else:
            # PPM
            magic = "P6"
            if plain:
                magic = "P3"
        header = "{magic} {width:d} {height:d} {maxval:d}\n".format(magic=magic, **meta)
    if planes in (2, 4):
        # PAM
        # See http://netpbm.sourceforge.net/doc/pam.html
        if plain:
            raise Exception("PAM (%d-plane) does not support plain format" % planes)
        if 2 == planes:
            tupltype = "GRAYSCALE_ALPHA"
        else:
            tupltype = "RGB_ALPHA"
        header = (
            "P7\nWIDTH {width:d}\nHEIGHT {height:d}\n"
            "DEPTH {planes:d}\nMAXVAL {maxval:d}\n"
            "TUPLTYPE {tupltype}\nENDHDR\n".format(tupltype=tupltype, **meta)
        )
    file.write(header.encode("ascii"))

    # Values per row
    vpr = planes * meta["width"]

    if plain:
        for row in rows:
            row_b = b" ".join([b"%d" % v for v in row])
            file.write(row_b)
            file.write(b"\n")
    else:
        # format for struct.pack
        fmt = ">%d" % vpr
        if meta["maxval"] > 0xFF:
            fmt = fmt + "H"
        else:
            fmt = fmt + "B"
        for row in rows:
            file.write(struct.pack(fmt, *row))

    file.flush()


def main(argv=None):
    import argparse

    parser = argparse.ArgumentParser(description="Convert PNG to PAM")
    parser.add_argument("--plain", action="store_true")
    parser.add_argument(
        "input", nargs="?", default="-", type=png.cli_open, metavar="PNG"
    )

    args = parser.parse_args()

    # Encode PNG to PNM (or PAM)
    image = png.Reader(file=args.input)
    _, _, rows, info = image.asDirect()
    write_pnm(png.binary_stdout(), args.plain, rows, info)


if __name__ == "__main__":
    import sys

    sys.exit(main())