#!/usr/bin/env python

"""
corrupt
usage: ./corrupt.py --source=input.jpg [--output=corrupted.jpg] [--iterations=10]

originally: http://www.recyclism.com/corrupt.php
ported: @mattdennewitz
"""

import argparse
import os
from random import randint

from PIL import Image


DEFAULT_ITERATIONS = 10
DEFAULT_OUTPUT = "corrupt.jpg"


class InvalidImageException(Exception):
    pass


def corrupt(source, output=DEFAULT_OUTPUT, iterations=DEFAULT_ITERATIONS):
    """Reconsider byte arrangement from ``input``, render to ``output.
    
    @param source Path to source image
    @param output Path to output image
    @param iterations How many times to rethink and rearrange.

    """

    def read_chunks(fh):
        "Read a file in chunks. Useful for larger images."
        while True:
            chunk = fh.read(4096)
            if not chunk:
                break
            yield chunk

    try:
        Image.open(source)
    except:
        raise InvalidImageException("%s appears to be broken." % source)

    # read the image, copy into a buffer for manipulation
    source_image = open(source, "r")
    output_image = open(output, "w+b")
    for chunk in read_chunks(source_image):
        output_image.write(chunk)
    source_image.close()

    # herald the destroyer
    iters = 0
    steps = 0
    block_start = 100
    block_end = os.path.getsize(source) - 400
    replacements = randint(1, 30)

    while iters <= iterations:
        while steps <= replacements:
            pos_a = randint(block_start, block_end)
            pos_b = randint(block_start, block_end)

            output_image.seek(pos_a)
            content_from_pos_a = output_image.read(1)
            output_image.seek(0)

            output_image.seek(pos_b)
            content_from_pos_b = output_image.read(1)
            output_image.seek(0)

            # overwrite A with B
            output_image.seek(pos_a)
            output_image.write(content_from_pos_b)
            output_image.seek(0)

            # overwrite B with A
            output_image.seek(pos_b)
            output_image.write(content_from_pos_a)
            output_image.seek(0)
    
            steps += 1
        iters += 1

    output_image.close()


if __name__ == "__main__":
    arg_parser = argparse.ArgumentParser()

    arg_parser.add_argument("--source", help="Path to source image",
                            dest="source", required=True)
    arg_parser.add_argument("--output", help="Path to glitched mayhem output",
                            dest="output", required=False,
                            default="glitched.jpg")
    arg_parser.add_argument("--iterations", type=int, help="# Iterations", 
                            default=DEFAULT_ITERATIONS, dest="iterations")

    app_args = arg_parser.parse_args()

    corrupt(app_args.source, output=app_args.output,
            iterations=app_args.iterations)