"""
A library for the steganography program
Copyright (C) 2022  Valentin Moguérou

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""

from typing import BinaryIO
from PIL import Image

# ================= BINARY OPERATIONS =================

def write_lsb_in_bin(byte, bit):
    "Write the specified bit in the LSB of the specified byte"
    return byte & 254 | bit if bit==1 else byte & 254

def read_lsb_in_bin(byte):
    return byte & 1

def combine_bits_to_byte(bits):
    """
    Turns a sequence of bits into an integer.
    Example: turns [0, 1, 1, 0, 1, 1, 0, 1] into 0b01101101 or 109
    """
    return sum(bits[-1-i]<<i for i in range(len(bits)))

def split_byte_into_bits(integer):
    """
    Turns an integer into an array of bits
    Example: turns 109 or 0b01101101 into [0, 1, 1, 0, 1, 1, 0, 1]
    """
    return [integer>>(7-i)&1 for i in range(8)]

def hex_print(byte_list, margin=0, line_width=16):
    "Prints a byte list in hexadecimal"

    for line in range(0, len(byte_list), line_width):
        print(' '*margin + ' '.join(f'{byte:02X}' for byte in byte_list[line:line+line_width]))

# ================= STRING OPERATIONS =================

def encode_string(string: str) -> bytearray:
    "Turns the given string into a byte array"
    return bytearray(ord(ch) for ch in string+'\0')

def decode_string(bytelist: bytearray) -> str:
    "Turns the given byte array into a string"
    return ''.join(chr(b) for b in bytelist)

# ================= WRITE OPERATIONS =================

def write_band(band: bytearray, byte_list: bytearray, starting_pos=0) -> bool:
    "Writes a byte sequence in the given band"

    index = starting_pos
    for index in range(starting_pos, len(band), 8):
        band[index:index+8] = (write_lsb_in_bin(band[index+i], bit) for i, bit in enumerate(split_byte_into_bits(byte_list[index//8])))
        if byte_list[index//8] == 0:
            return True, index+8 # finished writing
    return False, index+8 # didn't finish writing, return last index

def write_image(img: Image.Image, byte_list, verbose=False):
    data = [bytearray(band.tobytes()) for band in img.split()]

    ok, position = write_band(data[0], byte_list)
    if not ok:
        ok, position = write_band(data[1], byte_list, position)
    if not ok:
        ok, position = write_band(data[2], byte_list, position)
    if not ok:
        raise ValueError("The byte sequence is too long.")

    if verbose:
        print(f"{position} bits, {position//8} bytes successfully written.")

    return Image.merge(img.mode, [Image.frombytes('L', img.size, bytes(band)) for band in data])

def write_file(from_file: BinaryIO, to_file: BinaryIO, byte_list, verbose=False):
    write_image(Image.open(from_file), byte_list, verbose=verbose).save(to_file)

# ================= READ OPERATIONS =================

def read_band(band: bytes, byte_list: bytearray) -> bool:
    "Turns a sequence of pixels into a byte list"

    for i in range(0, len(band), 8):
        bits = [read_lsb_in_bin(b) for b in band[i:i+8]]
        if all(b==0 for b in bits):
            return True # finished reading
        byte_list.append(combine_bits_to_byte(bits))
    return False # didn't finish reading

def read_image(img: Image.Image, verbose=False):
    "Read the whole image"

    data = [band.tobytes() for band in img.split()]
    byte_list = bytearray()

    if read_band(data[0], byte_list):
        if verbose:
            print("Successfully finished reading.")
    elif read_band(data[1], byte_list):
        if verbose:
            print("Successfully finished reading.")
            print("Read the whole red band.")
    elif read_band(data[1], byte_list):
        if verbose:
            print("Successfully finished reading.")
            print("Read the whole green band.")
    else:
        if verbose:
            print("Read the whole blue band.")
        raise ValueError("Invalid image, did not find end control sequence.")

    return byte_list

def read_file(from_file: BinaryIO, verbose=False):
    return read_image(Image.open(from_file), verbose=verbose)

# ================= EXAMPLE PROGRAM =================

def main():
    s = "This is a string"
    oldimg = Image.open('default_image.jpg')
    oldimg.load()
    print(encode_string(s))
    newimg = write_image(oldimg, encode_string(s))
    newimg.save('dissimulated.png', 'PNG')
    print('Written...')

    print(f"String: --> {read_image(Image.open('default_image.jpg'))} (default_image.jpg)")
    print(f"String: --> {decode_string(read_image(Image.open('dissimulated.png')))} (dissimulated.png)")

if __name__ == '__main__':
    main()