# ctswap.py - program to swap palettes in BMP images
#
# for usage information see http://www.ugcs.caltech.edu/~q/code/ctswap
# 
# Copyright (c) 2005, Tom Quetchenbach (q at ugcs dot caltech dot edu)
# All rights reserved.

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 
#     * Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
# 
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import os
import sys
import re
import array
import stat
import getopt

BMP_MAGIC = 0x4d42
BMP_MAGIC_SWAPPED = 0x424d
BMP_PALETTE_LENGTH = 0x2e
BMP_BIT_COUNT = 0x1c
BMP_PALETTE_START = 0x36
BMP_24BIT_NUM_COLORS = 2**24

SEEK_RELATIVE = 1

EXIT_PALETTE_LENGTH_MISMATCH = 1
EXIT_INVALID_BMP = 2
EXIT_INVALID_OPTIONS = 3
EXIT_IO_ERROR = 4

COLOR_TABLE_ENTRY_LENGTH = 4

class BMPError(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)

class BMPValidError(BMPError):
    pass

class BMPTypeError(BMPError):
    pass

def get_color_table(fn):
    try:
        f = open(fn, 'rb')
        ret = get_bmp_color_table(f) 
        f.close()
        return ret;
    except BMPValidError, e:
        # we tried to read a palette from something that wasn't a BMP file.
        # maybe it's a text file instead. This isn't a great way of doing
        # things, because we might want to be able to read from gif files etc
        # in the future.
        eprint(e);
        return readpal(fn);
    except BMPTypeError, e:
        eprint("Cannot read new palette from a 24-bit BMP file.")
        sys.exit(EXIT_INVALID_BMP)
    except IOError, e:
        eprint(e)
        sys.exit(EXIT_IO_ERROR)
    except:
        raise

def get_bmp_color_table(f):
    n = get_num_colors(f)
    eprint("Reading new palette from BMP file...")
    f.seek(BMP_PALETTE_START)
    table = array.array('B')
    table.fromfile(f, n * COLOR_TABLE_ENTRY_LENGTH)
    return table

def readpal(fn):
    """Read a palette file into a byte array representing a BMP color table. 

    Each line of the input file should have three numbers separated by spaces
    representing the red, green, and blue color levels and ranging from 0 to
    255. For example:

    0 0 0
    255 0 0
    0 255 0
    ...and so on

    A BMP color table is an array of four-byte RGBQUAD elements. The first
    byte in the RGBQUAD is the blue value, the second is the green, the third
    is the red, and the fourth is always zero."""

    palfile = open(fn, 'r')
    palette = array.array('B')
    for line in palfile:
        matches = re.findall("(\d+)\s+(\d+)\s+(\d+)", line)
        if matches:
            rgb = [int(x) for x in matches[0]]
            rgb.reverse()
            palette += array.array('B', rgb + [0])
    palfile.close()
    return palette

def get_bmp_swap(imgfile):
    """Determine whether BMP byte order matches native byte order

    Returns True if we need to swap bytes when reading from BMP fields (i.e if
    we are running on a little-endian machine.
    
    Raises BMPValidError if the file does not contain the correct BMP magic.
    
    Note that this is not a thorough check, as any text file starting with BM
    or MB will be detected as a BMP."""

    imgfile.seek(0) 
    magic = readint(imgfile, 'H', False)
    if magic == BMP_MAGIC:
        return False
    elif magic == BMP_MAGIC_SWAPPED:
        #running on little-endian machine
        return True
    else:
        raise BMPValidError("Not a BMP file!")

def get_num_colors(imgfile):
    """Get the size of the color table from a bmp file"""

    byteswap = get_bmp_swap(imgfile)        
    imgfile.seek(BMP_PALETTE_LENGTH)
    n = readint(imgfile, 'L', byteswap)
    if n == 0:
        # we have to calculate the length of the color table from the
        # bits-per-pixel value
        imgfile.seek(BMP_BIT_COUNT)
        n = 2**readint(imgfile, 'H', byteswap)
    if n == BMP_24BIT_NUM_COLORS:
        raise BMPTypeError("File %n is a 24-bit bitmap!")
    return n

def readint(f, type, swap):
    """Read an integer (with big-endian byte order) from the current
    position in the file f"""

    arr = array.array(type)
    arr.fromfile(f, 1);
    if swap: arr.byteswap()
    return arr[0]

def replace_palette(palette, imgfile, newimgfile):
    """write a new color table to the file fn and write to newimgfile"""

    # go back to beginning of file
    imgfile.seek(0)

    # copy header
    newimgfile.write(imgfile.read(BMP_PALETTE_START))

    #write new palette
    palette.tofile(newimgfile)

    #copy rest of file
    imgfile.seek(len(palette), SEEK_RELATIVE)
    newimgfile.write(imgfile.read())

def eprint(s):
    """write a line to standard error"""
    
    sys.stderr.write(str(s) + "\n")
    
def usage():
    """print a usage message to standard error"""
    
    sys.stderr.write("usage: %s [-F] [-o outfile] palette infile\n" %
            sys.argv[0])
    sys.stderr.write("\n    -F: force even if color table sizes differ or")
    sys.stderr.write(" input is not valid\n")
    sys.stderr.write("    -o: write to outfile instead of standard")
    sys.stderr.write(" output\n")
    sys.stderr.write("    palette: BMP or JASC PAL-compatible file")
    sys.stderr.write(" containing new color table\n")
    sys.stderr.write("    infile: BMP file to operate on\n\n")
    sys.exit(EXIT_INVALID_OPTIONS)

#initial values for options

outfilename = "standard output"
outfile = sys.stdout
force = False

# process command-line options

try:
    opts, args = getopt.getopt(sys.argv[1:], "o:F")
    for k, v in opts:
        if k == '-o':
            #output to file instead of stdout
            outfilename = v
            outfile = open(outfilename, 'wb')
        if k == '-F':
            force = True
    if len(args) > 2: raise IndexError
    palfilename = args[0]
    imgfilename = args[1]
except IndexError:
    usage()
except getopt.GetoptError:
    usage()
except:
    raise

# get the new color table from a text file or BMP file

p = get_color_table(palfilename)
# get number of colors in table
n = len(p) / COLOR_TABLE_ENTRY_LENGTH
eprint("Read new palette with %i colors" % n)

# try to open input image file

try:
    imgfile = open(imgfilename, 'rb')
except IOError, e:
    eprint(e);
    sys.exit(EXIT_IO_ERROR)
except:
    raise

if not force:
    #make sure input image is a BMP and has the correct palette length
    try:
        realn = get_num_colors(imgfile)
        eprint("Image actually has %i colors" % realn)
    except BMPValidError, e:
        eprint(e)
        sys.exit(EXIT_INVALID_BMP)
    except BMPTypeError, e:
        eprint("Input image is a 24-bit BMP. That doesn't make any sense. ")
        eprint("I give up.")
        sys.exit(EXIT_INVALID_BMP)
    except:
        raise

    if realn != n:
        eprint("New palette and image palette are different lengths!")
        sys.exit(EXIT_PALETTE_LENGTH_MISMATCH)

eprint("About to replace palette on file %s and write to %s" % (imgfilename,
        outfilename))

# actually replace palette, writing to output file
replace_palette(p, imgfile, outfile)

#close input file
imgfile.close()

#perform final sanity check and close output file, unless we wrote to stdout
if (outfile != sys.stdout):
    outfile.close()
    oldsize = os.stat(imgfilename)[stat.ST_SIZE]
    newsize = os.stat(outfilename)[stat.ST_SIZE]

    if oldsize == newsize:
        eprint("Done.")
    else:
        eprint("Done, but old file size was %i and new size is %i" % (oldsize,
                newsize)) 
else:
    eprint("Done.")