from PIL import Image
import numpy as np

def string_to_binary(message):
    return ''.join(format(ord(char), '08b') for char in message)

def embed_message(image_path, message, exifraw):
    # Open the image
    img = Image.open(image_path)
    # Convert image to numpy array
    img_array = np.array(img)

    # Flatten the array
    flat_array = img_array.flatten()

    # Convert message to binary
    binary_message = string_to_binary(message)
    
    # Check if the message can fit in the image
    if len(binary_message) > len(flat_array):
        raise ValueError("Message is too long to be embedded in this image")

    # Embed the message
    for i, bit in enumerate(binary_message):
        flat_array[i] = (flat_array[i] & 0xFE) | int(bit)

    # Reshape the array back to the original image shape
    stego_array = flat_array.reshape(img_array.shape)

    # Create a new image from the modified array
    stego_img = Image.fromarray(stego_array.astype('uint8'), img.mode)
    
    # Save the image
    stego_img.save(image_path, exif=exifraw)

def extract_message(image_path, message_length):
    # Open the image
    img = Image.open(image_path)
    # Convert image to numpy array
    img_array = np.array(img)

    # Flatten the array
    flat_array = img_array.flatten()

    # Extract the binary message
    binary_message = ''.join([str(pixel & 1) for pixel in flat_array[:message_length * 8]])

    # Convert binary to string
    message = ''.join([chr(int(binary_message[i:i+8], 2)) for i in range(0, len(binary_message), 8)])

    return message