from PIL import Image
import numpy as np
from scipy.fftpack import dct, idct
import io

QUANTIZATION_FACTOR = 10

def convert_and_resize_image(image):
    gray_image = image.convert("L")
    w, h = gray_image.size
    w_crop = w - (w % 8)
    h_crop = h - (h % 8)
    if w_crop == 0 or h_crop == 0:
        return None
    return gray_image.crop((0, 0, w_crop, h_crop))

def message_to_binary(message):
    message += "_END_"
    return ''.join(format(ord(c), '08b') for c in message)

def binary_to_message(binary_message):
    delimiter = ''.join(format(ord(c), '08b') for c in "_END_")
    if delimiter not in binary_message:
        return None
    try:
        end_index = binary_message.index(delimiter)
        clean_binary = binary_message[:end_index]
        if len(clean_binary) % 8 != 0:
            return None
        return ''.join(chr(int(clean_binary[i:i+8], 2)) for i in range(0, len(clean_binary), 8))
    except:
        return None

def apply_dct(image_array):
    h, w = image_array.shape
    dct_blocks = np.zeros_like(image_array, dtype=float)
    for i in range(0, h, 8):
        for j in range(0, w, 8):
            block = image_array[i:i+8, j:j+8].astype(float)
            dct_blocks[i:i+8, j:j+8] = dct(dct(block.T, norm='ortho').T, norm='ortho')
    return dct_blocks

def apply_idct(dct_blocks):
    h, w = dct_blocks.shape
    image_blocks = np.zeros_like(dct_blocks, dtype=float)
    for i in range(0, h, 8):
        for j in range(0, w, 8):
            block = dct_blocks[i:i+8, j:j+8]
            image_blocks[i:i+8, j:j+8] = idct(idct(block.T, norm='ortho').T, norm='ortho')
    return np.clip(image_blocks, 0, 255).astype(np.uint8)

def embed_message(image, secret_message):
    gray_cropped = convert_and_resize_image(image)
    if gray_cropped is None:
        return None
    img_array = np.array(gray_cropped)
    h, w = img_array.shape

    binary_message = message_to_binary(secret_message)
    msg_len = len(binary_message)

    dct_coeffs = apply_dct(img_array)
    bit_index = 0
    for i in range(0, h, 8):
        for j in range(0, w, 8):
            if bit_index >= msg_len:
                break
            u, v = 4, 1
            coeff = dct_coeffs[i + u, j + v]
            bit = int(binary_message[bit_index])
            quantized_coeff = int(round(coeff / QUANTIZATION_FACTOR))
            if quantized_coeff % 2 != bit:
                quantized_coeff -= 1
            dct_coeffs[i + u, j + v] = quantized_coeff * QUANTIZATION_FACTOR
            bit_index += 1
        if bit_index >= msg_len:
            break
    stego_array = apply_idct(dct_coeffs)
    return Image.fromarray(stego_array)

def extract_message(image):
    gray_cropped = convert_and_resize_image(image)
    if gray_cropped is None:
        return None
    img_array = np.array(gray_cropped)
    h, w = img_array.shape
    dct_coeffs = apply_dct(img_array)
    binary_message = ""
    for i in range(0, h, 8):
        for j in range(0, w, 8):
            u, v = 4, 1
            coeff = dct_coeffs[i + u, j + v]
            quantized_coeff = int(round(coeff / QUANTIZATION_FACTOR))
            binary_message += str(quantized_coeff % 2)
    return binary_to_message(binary_message)