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

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, pnginfo=info)

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

        return new_filename, "", png_encode_coords

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


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 decode(image_name, txt=None):
    try:
        BGRimage = cv2.imread(f"{image_name}")
        image = cv2.cvtColor(BGRimage, cv2.COLOR_BGR2RGB)
        binary_data = ""
        for row in image:
            for pixel in row:
                r, g, b = to_bin(pixel)
                binary_data += r[-1]
                binary_data += g[-1]
                binary_data += b[-1]
        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
        this = decoded_data[:-5].split("#####", 1)[0]
        return this
    except Exception as e:
        return str(e)


def encode(image_name, secret_data, txt=None):
    msg = ""
    BGRimage = cv2.imread(f"{image_name}")
    image = cv2.cvtColor(BGRimage, cv2.COLOR_BGR2RGB)
    n_bytes = image.shape[0] * image.shape[1] * 3 // 8
    print("[*] Maximum bytes to encode:", n_bytes)
    secret_data1 = secret_data
    while True:
        if len(secret_data1) + 5 < (n_bytes):
            secret_data1 = f'{secret_data1}#####'
        elif len(secret_data1) + 5 >= (n_bytes):
            break
    secret_data = secret_data1
    if len(secret_data) > n_bytes:
        msg = "Watermark is too large for Image Size"
        return image_name, msg
    secret_data += "====="
    data_index = 0
    binary_secret_data = to_bin(secret_data)
    data_len = len(binary_secret_data)
    for row in image:
        for pixel in row:
            r, g, b = to_bin(pixel)
            if data_index < data_len:
                pixel[0] = int(r[:-1] + binary_secret_data[data_index], 2)
                data_index += 1
            if data_index < data_len:
                pixel[1] = int(g[:-1] + binary_secret_data[data_index], 2)
                data_index += 1
            if data_index < data_len:
                pixel[2] = int(b[:-1] + binary_secret_data[data_index], 2)
                data_index += 1
            if data_index >= data_len:
                break
    return image, msg