import numpy as np
from PIL import Image


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