import cv2
import numpy as np
from PIL import Image, PngImagePlugin
import os
import uuid

def to_bin(data):
    if isinstance(data, str):
        return ''.join([format(ord(i), "08b") for i in data])
    elif isinstance(data, bytes):
        return ''.join([format(i, "08b") for i in data])
    elif isinstance(data, np.ndarray):
        return [format(i, "08b") for i in data]
    elif isinstance(data, int) or isinstance(data, np.uint8):
        return format(data, "08b")
    else:
        raise TypeError("Type not supported.")

def encode(image_path, secret_data):
    try:
        # Read the image
        image = cv2.imread(image_path)
        if image is None:
            return None, "Failed to read image"
            
        # Calculate maximum bytes for encoding
        n_bytes = image.shape[0] * image.shape[1] * 3 // 8
        
        # Check if the secret data is too large
        secret_data_with_delimiter = f'{secret_data}#####'
        if len(secret_data_with_delimiter) + 5 >= n_bytes:
            return None, "Watermark is too large for Image Size"
            
        # Add ending delimiter
        secret_data_with_delimiter += "====="
        
        # Convert data to binary
        binary_secret_data = to_bin(secret_data_with_delimiter)
        data_len = len(binary_secret_data)
        data_index = 0
        
        # Encode the data into the image
        for i in range(image.shape[0]):
            for j in range(image.shape[1]):
                pixel = image[i, j]
                for color_channel in range(3):
                    if data_index < data_len:
                        # Get the binary value of the pixel
                        binary_value = to_bin(pixel[color_channel])
                        # Replace the least significant bit
                        new_binary = binary_value[:-1] + binary_secret_data[data_index]
                        # Update the pixel with the new value
                        pixel[color_channel] = int(new_binary, 2)
                        data_index += 1
                image[i, j] = pixel
                if data_index >= data_len:
                    break
            if data_index >= data_len:
                break
                
        return image, None
    except Exception as e:
        return None, f"Error during encoding: {str(e)}"

def decode(image_path):
    try:
        # Read the image
        image = cv2.imread(image_path)
        if image is None:
            return "Failed to read image"
            
        binary_data = ""
        for row in image:
            for pixel in row:
                for color_channel in pixel:
                    binary_data += to_bin(color_channel)[-1]
        
        # Convert binary to text
        all_bytes = [binary_data[i: i + 8] for i in range(0, len(binary_data), 8)]
        decoded_data = ""
        for byte in all_bytes:
            decoded_data += chr(int(byte, 2))
            if decoded_data[-5:] == "=====":
                break
        
        # Remove delimiters
        return decoded_data[:-5].split("#####")[0]
    except Exception as e:
        return f"Error during decoding: {str(e)}"

def png_encode(im_name, extra):
    try:
        im = Image.open(im_name)
        info = PngImagePlugin.PngInfo()
        info.add_text("TXT", extra)

        unique_id = str(uuid.uuid4())[:8]
        filename, ext = os.path.splitext(im_name)
        new_filename = f"{filename}_{unique_id}{ext}"
        im.save(new_filename, "PNG", pnginfo=info)

        width, height = im.size
        rect_width, rect_height = 200, 50
        x = width - rect_width - 10
        y = height - rect_height - 10
        global png_encode_coords
        png_encode_coords = (x, y, rect_width, rect_height)

        return new_filename, None

    except Exception as e:
        return None, str(e)