#!/usr/bin/env python

# Imported from //depot/prj/plan9topam/master/code/plan9topam.py#4 on
# 2009-06-15.

"""Command line tool to convert from Plan 9 image format to PNG format.

Plan 9 image format description:
https://plan9.io/magic/man2html/6/image

Where possible this tool will use unbuffered read() calls,
so that when finished the file offset is exactly at the end of
the image data.
This is useful for Plan9 subfont files which place font metric
data immediately after the image.
"""

# Test materials

# asset/left.bit is a Plan 9 image file, a leftwards facing Glenda.
# Other materials have to be scrounged from the internet.
# https://plan9.io/sources/plan9/sys/games/lib/sokoban/images/cargo.bit

import array
import collections
import io

# http://www.python.org/doc/2.3.5/lib/module-itertools.html
import itertools
import os

# http://www.python.org/doc/2.3.5/lib/module-re.html
import re
import struct

# http://www.python.org/doc/2.3.5/lib/module-sys.html
import sys

# https://docs.python.org/3/library/tarfile.html
import tarfile


# https://pypi.org/project/pypng/
import png

# internal
import prix


class Error(Exception):
    """Some sort of Plan 9 image error."""


def block(s, n):
    return zip(*[iter(s)] * n)


def plan9_as_image(inp):
    """Represent a Plan 9 image file as a png.Image instance, so
    that it can be written as a PNG file.
    Works with compressed input files and may work with uncompressed files.
    """

    # Use inp.raw if available.
    # This avoids buffering and means that when the image is processed,
    # the resulting input stream is cued up exactly at the end
    # of the image.
    inp = getattr(inp, "raw", inp)

    info, blocks = plan9_open_image(inp)

    rows, infodict = plan9_image_rows(blocks, info)

    return png.Image(rows, infodict)


def plan9_open_image(inp):
    """Open a Plan9 image file (`inp` should be an already open
    file object), and return (`info`, `blocks`) pair.
    `info` should be a Plan9 5-tuple;
    `blocks` is the input, and it should yield (`row`, `data`)
    pairs (see :meth:`pixmeta`).
    """

    r = inp.read(11)
    if r == b"compressed\n":
        info, blocks = decompress(inp)
    else:
        # Since Python 3, there is a good chance that this path
        # doesn't work.
        info, blocks = glue(inp, r)

    return info, blocks


def glue(f, r):
    """Return (info, stream) pair, given `r` the initial portion of
    the metadata that has already been read from the stream `f`.
    """

    r = r + f.read(60 - len(r))
    return (meta(r), f)


def meta(r):
    """Convert 60 byte bytestring `r`, the metadata from an image file.
    Returns a 5-tuple (*chan*,*minx*,*miny*,*limx*,*limy*).
    5-tuples may settle into lists in transit.

    As per https://plan9.io/magic/man2html/6/image the metadata
    comprises 5 words separated by blanks.
    As it happens each word starts at an index that is a multiple of 12,
    but this routine does not care about that.
    """

    r = r.split()
    # :todo: raise FormatError
    if 5 != len(r):
        raise Error("Expected 5 space-separated words in metadata")
    r = [r[0]] + [int(x) for x in r[1:]]
    return r


def bitdepthof(chan):
    """Return the bitdepth for a Plan9 pixel format string."""

    maxd = 0
    for c in re.findall(rb"[a-z]\d*", chan):
        if c[0] != "x":
            maxd = max(maxd, int(c[1:]))
    return maxd


def maxvalof(chan):
    """Return the netpbm MAXVAL for a Plan9 pixel format string."""

    bitdepth = bitdepthof(chan)
    return (2 ** bitdepth) - 1


def plan9_image_rows(blocks, metadata):
    """
    Convert (uncompressed) Plan 9 image file to pair of (*rows*, *info*).
    This is intended to be used by PyPNG format.
    *info* is the image info (metadata) returned in a dictionary,
    *rows* is an iterator that yields each row in
    boxed row flat pixel format.

    `blocks`, should be an iterator of (`row`, `data`) pairs.
    """

    chan, minx, miny, limx, limy = metadata
    rows = limy - miny
    width = limx - minx
    nchans = len(re.findall(b"[a-wyz]", chan))
    alpha = b"a" in chan
    # Iverson's convention for the win!
    ncolour = nchans - alpha
    greyscale = ncolour == 1
    bitdepth = bitdepthof(chan)
    maxval = maxvalof(chan)

    # PNG style info dict.
    meta = dict(
        size=(width, rows),
        bitdepth=bitdepth,
        greyscale=greyscale,
        alpha=alpha,
        planes=nchans,
    )

    arraycode = "BH"[bitdepth > 8]

    return (
        map(
            lambda x: array.array(arraycode, itertools.chain(*x)),
            block(unpack(blocks, rows, width, chan, maxval), width),
        ),
        meta,
    )


def unpack(f, rows, width, chan, maxval):
    """Unpack `f` into pixels.
    `chan` describes the pixel format using
    the Plan9 syntax ("k8", "r8g8b8", and so on).
    Assumes the pixel format has a total channel bit depth
    that is either a multiple or a divisor of 8
    (the Plan9 image specification requires this).
    `f` should be an iterator that returns blocks of input such that
    each block contains a whole number of pixels.
    The return value is an iterator that yields each pixel as an n-tuple.
    """

    def mask(w):
        """An integer, to be used as a mask, with bottom `w` bits set to 1."""

        return (1 << w) - 1

    def deblock(f, depth, width):
        """A "packer" used to convert multiple bytes into single pixels.
        `depth` is the pixel depth in bits (>= 8), `width` is the row width in
        pixels.
        """

        w = depth // 8
        i = 0
        for block in f:
            for i in range(len(block) // w):
                p = block[w * i : w * (i + 1)]
                i += w
                # Convert little-endian p to integer x
                x = 0
                s = 1  # scale
                for j in p:
                    x += s * j
                    s <<= 8
                yield x

    def bitfunge(f, depth, width):
        """A "packer" used to convert single bytes into multiple pixels.
        Depth is the pixel depth (< 8), width is the row width in pixels.
        """

        assert 8 / depth == 8 // depth

        for block in f:
            col = 0
            for x in block:
                for j in range(8 // depth):
                    yield x >> (8 - depth)
                    col += 1
                    if col == width:
                        # A row-end forces a new byte even if
                        # we haven't consumed all of the current byte.
                        # Effectively rows are bit-padded to make
                        # a whole number of bytes.
                        col = 0
                        break
                    x <<= depth

    # number of bits in each channel
    bits = [int(d) for d in re.findall(rb"\d+", chan)]
    # colr of each channel
    # (r, g, b, k for actual colours, and
    # a, m, x for alpha, map-index, and unused)
    colr = re.findall(b"[a-z]", chan)

    depth = sum(bits)

    # Select a "packer" that either:
    # - gathers multiple bytes into a single pixel (for depth >= 8); or,
    # - splits bytes into several pixels (for depth < 8).
    if depth >= 8:
        assert depth % 8 == 0
        packer = deblock
    else:
        assert 8 % depth == 0
        packer = bitfunge

    for x in packer(f, depth, width):
        # x is the pixel as an unsigned integer
        o = []
        # This is a bit yucky.
        # Extract each channel from the _most_ significant part of x.
        for b, col in zip(bits, colr):
            v = (x >> (depth - b)) & mask(b)
            x <<= b
            if col != "x":
                # scale to maxval
                v = v * float(maxval) / mask(b)
                v = int(v + 0.5)
                o.append(v)
        yield o


def decompress(f):
    """Decompress a Plan 9 image file.
    The input `f` should be a binary file object that
    is already cued past the initial 'compressed\n' string.
    The return result is (`info`, `blocks`);
    `info` is a 5-tuple of the Plan 9 image metadata;
    `blocks` is an iterator that yields a (row, data) pair
    for each block of data.
    """

    r = meta(f.read(60))
    return r, decomprest(f, r[4])


def decomprest(f, rows):
    """Iterator that decompresses the rest of a file once the metadata
    have been consumed."""

    row = 0
    while row < rows:
        row, o = deblock(f)
        yield o


def deblock(f):
    """Decompress a single block from a compressed Plan 9 image file.
    Each block starts with 2 decimal strings of 12 bytes each.
    Yields a sequence of (row, data) pairs where
    `row` is the total number of rows processed
    (according to the file format) and
    `data` is the decompressed data for this block.
    """

    row = int(f.read(12))
    size = int(f.read(12))
    if not (0 <= size <= 6000):
        raise Error("block has invalid size; not a Plan 9 image file?")

    # Since each block is at most 6000 bytes we may as well read it all in
    # one go.
    d = f.read(size)
    i = 0
    o = []

    while i < size:
        x = d[i]
        i += 1
        if x & 0x80:
            x = (x & 0x7F) + 1
            lit = d[i : i + x]
            i += x
            o.extend(lit)
            continue
        # x's high-order bit is 0
        length = (x >> 2) + 3
        # Offset is made from bottom 2 bits of x and 8 bits of next byte.
        #     MSByte                                 LSByte
        #   +---------------------+-------------------------+
        #   | - - - - - - | x1 x0 | d7 d6 d5 d4 d3 d2 d1 d0 |
        #   +-----------------------------------------------+
        # Had to discover by inspection which way round the bits go,
        # because https://plan9.io/magic/man2html/6/image doesn't say.
        # that x's 2 bits are most significant.
        offset = (x & 3) << 8
        offset |= d[i]
        i += 1
        # Note: complement operator neatly maps (0 to 1023) to (-1 to
        # -1024).  Adding len(o) gives a (non-negative) offset into o from
        # which to start indexing.
        offset = ~offset + len(o)
        if offset < 0:
            raise Error(
                "byte offset indexes off the begininning of "
                "the output buffer; not a Plan 9 image file?"
            )
        for j in range(length):
            o.append(o[offset + j])
    return row, bytes(o)


FontChar = collections.namedtuple("FontChar", "x top bottom left width")


def font_copy(inp, image, out, control):
    """
    Convert a Plan 9 font (`inp`, `image`) to a series of PNG images,
    and write them out as a tar file to the file object `out`.
    Write a text control file out to the file object `control`.

    Each valid glyph in the font becomes a single PNG image;
    the output is a tar file of all the images.

    A Plan 9 font consists of a Plan 9 image immediately
    followed by font data.
    The image for the font should be the `image` argument,
    the file containing the rest of the font data should be the
    file object `inp` which should be cued up to the start of
    the font data that immediately follows the image.

    https://plan9.io/magic/man2html/6/font
    """

    # The format is a little unusual, and isn't completely
    # clearly documented.
    # Each 6-byte structure (see FontChar above) defines
    # a rectangular region of the image that is used for each
    # glyph.
    # The source image region that is used may be strictly
    # smaller than the rectangle for the target glyph.
    # This seems like a micro-optimisation.
    # For each glyph,
    # rows above `top` and below `bottom` will not be copied
    # from the source (they can be assumed to be blank).
    # No space is saved in the source image, since the rows must
    # be present.
    # `x` is always non-decreasing, so the glyphs appear strictly
    # left-to-image in the source image.
    # The x of the next glyph is used to
    # infer the width of the source rectangle.
    # `top` and `bottom` give the y-coordinate of the top- and
    # bottom- sides of the rectangle in both source and targets.
    # `left` is the x-coordinate of the left-side of the
    # rectangle in the target glyph. (equivalently, the amount
    # of padding that should be added on the left).
    # `width` is the advance-width of the glyph; by convention
    # it is 0 for an undefined glyph.

    name = getattr(inp, "name", "*subfont*name*not*supplied*")

    header = inp.read(36)
    n, height, ascent = [int(x) for x in header.split()]
    print("baseline", name, ascent, file=control, sep=",")

    chs = []
    for i in range(n + 1):
        bs = inp.read(6)
        ch = FontChar(*struct.unpack("<HBBBB", bs))
        chs.append(ch)

    tar = tarfile.open(mode="w|", fileobj=out)

    # Start at 0, increment for every image output
    # (recall that not every input glyph has an output image)
    output_index = 0
    for i in range(n):
        ch = chs[i]
        if ch.width == 0:
            continue

        print("png", "index", output_index, "glyph", name, i, file=control, sep=",")

        info = dict(image.info, size=(ch.width, height))
        target = new_image(info)

        source_width = chs[i + 1].x - ch.x
        rect = ((ch.left, ch.top), (ch.left + source_width, ch.bottom))
        image_draw(target, rect, image, (ch.x, ch.top))

        # :todo: add source, glyph, and baseline data here (as a
        # private tag?)
        o = io.BytesIO()
        target.write(o)
        binary_size = o.tell()
        o.seek(0)

        tarinfo = tar.gettarinfo(arcname="%s/glyph%d.png" % (name, i), fileobj=inp)
        tarinfo.size = binary_size
        tar.addfile(tarinfo, fileobj=o)

        output_index += 1

    tar.close()


def new_image(info):
    """Return a fresh png.Image instance."""

    width, height = info["size"]
    vpr = width * info["planes"]
    row = lambda: [0] * vpr
    rows = [row() for _ in range(height)]
    return png.Image(rows, info)


def image_draw(target, rect, source, point):
    """The point `point` in the source image is aligned with the
    top-left of rect in the target image, and then the rectangle
    in target is replaced with the pixels from `source`.

    This routine assumes that both source and target can have
    their rows objects indexed (not streamed).
    """

    # :todo: there is no attempt to do clipping or channel or
    # colour conversion. But maybe later?

    if target.info["planes"] != source.info["planes"]:
        raise NotImplementedError(
            "source and target must have the same number of planes"
        )

    if target.info["bitdepth"] != source.info["bitdepth"]:
        raise NotImplementedError("source and target must have the same bitdepth")

    tl, br = rect
    left, top = tl
    right, bottom = br
    height = bottom - top

    planes = source.info["planes"]

    vpr = (right - left) * planes
    source_left, source_top = point

    source_l = source_left * planes
    source_r = source_l + vpr

    target_l = left * planes
    target_r = target_l + vpr

    for y in range(height):
        row = source.rows[y + source_top]
        row = row[source_l:source_r]
        target.rows[top + y][target_l:target_r] = row


def main(argv=None):
    import argparse

    parser = argparse.ArgumentParser(description="Convert Plan9 image to PNG")
    parser.add_argument(
        "input",
        nargs="?",
        default="-",
        type=png.cli_open,
        metavar="image",
        help="image file in Plan 9 format",
    )
    parser.add_argument(
        "--control",
        default=os.path.devnull,
        type=argparse.FileType("w"),
        metavar="ControlCSV",
        help="(when using --font) write a control CSV file to named file",
    )
    parser.add_argument(
        "--font",
        action="store_true",
        help="process as Plan 9 subfont: output a tar file of PNGs",
    )

    args = parser.parse_args()

    image = plan9_as_image(args.input)
    image.stream()

    if not args.font:
        image.write(png.binary_stdout())
    else:
        font_copy(args.input, image, png.binary_stdout(), args.control)


if __name__ == "__main__":
    sys.exit(main())