import os
import numpy as np
from imageio.v2 import imread, imwrite
import tkinter as tk
from tkinter import filedialog
from PIL import Image, ImageTk

max_value = 255  # max uint value per pixel per channel
header_len = 4 * 8  # uint32 bit length

def read_image(img_path):
    img = np.array(imread(img_path), dtype=np.uint8)
    orig_shape = img.shape
    return img.flatten(), orig_shape

def write_image(img_path, img_data, shape):
    img_data = np.reshape(img_data, shape)
    imwrite(img_path, img_data)

def bytes2array(byte_data):
    byte_array = np.frombuffer(byte_data, dtype=np.uint8)
    return np.unpackbits(byte_array)

def array2bytes(bit_array):
    byte_array = np.packbits(bit_array)
    return byte_array.tobytes()

def read_file(file_path):
    file_bytes = open(file_path, "rb").read()
    return bytes2array(file_bytes)

def write_file(file_path, file_bit_array):
    bytes_data = array2bytes(file_bit_array)
    with open(file_path, 'wb') as f:
        f.write(bytes_data)

def encode_data(image, file_data):
    or_mask = file_data
    and_mask = np.zeros_like(or_mask)
    and_mask = (and_mask + max_value - 1) + or_mask 
    res = np.bitwise_or(image, or_mask)
    res = np.bitwise_and(res, and_mask)
    return res

def decode_data(encoded_data):
    out_mask = np.ones_like(encoded_data)
    output = np.bitwise_and(encoded_data, out_mask)
    return output

def update_zoom(event=None):
    if 'original_image' not in globals():
        print("No image loaded to zoom")
        return
    
    zoom_value = zoom_slider.get() / 100  # Scale the value down to a fraction
    resized_image = original_image.resize((int(original_image.width * zoom_value), int(original_image.height * zoom_value)))
    zoomed_photo = ImageTk.PhotoImage(resized_image)

    image_label.config(image=zoomed_photo)
    image_label.image = zoomed_photo
    canvas.config(scrollregion=canvas.bbox(tk.ALL))

def hide_images():
    global original_image  # Ensure this is a global variable
    img_path = original_entry_hide.get()
    file_path = hide_file_entry.get()
    output_path = save_file_entry_hide.get()

    if not os.path.isfile(img_path):
        print("Original image file does not exist")
        return
    if not os.path.isfile(file_path):
        print("File to hide does not exist")
        return

    image, shape_orig = read_image(img_path)
    file = read_file(file_path)
    file_len = file.shape[0]
    len_array = np.array([file_len], dtype=np.uint32).view(np.uint8)
    len_array = np.unpackbits(len_array)
    img_len = image.shape[0]

    if file_len >= img_len - header_len:
        print("File too big, error")
        return
    else:
        tmp = file
        file = np.random.randint(2, size=img_len, dtype=np.uint8)
        file[header_len:header_len+file_len] = tmp

    file[:header_len] = len_array
    encoded_data = encode_data(image, file)

    write_image(output_path, encoded_data, shape_orig)
    print("Image encoded")

    # Update the preview image
    original_image = Image.fromarray(encoded_data.reshape(shape_orig))
    resized_image = original_image.resize((200, 400))  # First-time view as 200x400
    photo = ImageTk.PhotoImage(resized_image)
    
    # Update label and configure scroll region
    image_label.config(image=photo)
    image_label.image = photo
    canvas.config(scrollregion=canvas.bbox(tk.ALL))
    zoom_slider.set(30)  # Reset the zoom slider to default

# Function to refresh the GUI
def refresh_gui():
    # Clear the entry fields
    original_entry_hide.delete(0, tk.END)
    hide_file_entry.delete(0, tk.END)
    save_file_entry_hide.delete(0, tk.END)
    original_entry_unhide.delete(0, tk.END)
    save_file_entry_unhide.delete(0, tk.END)

    # Reset the zoom slider
    zoom_slider.set(50)

    # Clear the image displayed on the canvas
    image_label.config(image='')
    canvas.config(scrollregion=canvas.bbox(tk.ALL))

    print("GUI refreshed")


# Create the root window
root = tk.Tk()
root.geometry("1000x700")
root.title("Najeeb AI Generator Image Tool")
root.configure(bg="#282c34")

# Grid configuration
root.grid_rowconfigure(2, weight=1)
root.grid_columnconfigure(0, weight=1)
root.grid_columnconfigure(1, weight=1)
root.grid_columnconfigure(2, weight=4)  # Make column 2 wider for the canvas and zoom slider

# Create a frame for the left side (fields and buttons)
left_frame = tk.Frame(root, bg="#282c34")
left_frame.grid(row=0, column=0, padx=5, pady=5, sticky='nsew')

# Styling for labels and buttons
label_font = ('Arial', 12, 'bold')
entry_font = ('Arial', 11)
button_font = ('Arial', 12, 'bold')
button_color = "#61afef"
button_fg = "white"

# Input for Original Image to hide file in
tk.Label(left_frame, text="Select Original Image:", font=label_font, bg="#282c34", fg="white").grid(row=1, column=0, sticky="e", padx=10)
original_entry_hide = tk.Entry(left_frame, width=20, font=entry_font)
original_entry_hide.grid(row=1, column=1, padx=10, pady=5)
browse_button_hide_image = tk.Button(left_frame, text="Browse", font=button_font, bg=button_color, fg=button_fg, command=lambda: original_entry_hide.insert(0, filedialog.askopenfilename()))
browse_button_hide_image.grid(row=1, column=2, padx=10, pady=5)

# Input for File to hide
tk.Label(left_frame, text="Select Convert AI Image", font=label_font, bg="#282c34", fg="white").grid(row=1, column=3, sticky="e", padx=10)
hide_file_entry = tk.Entry(left_frame, width=20, font=entry_font)
hide_file_entry.grid(row=1, column=4, padx=10, pady=5)
browse_button_hide_file = tk.Button(left_frame, text="Browse", font=button_font, bg=button_color, fg=button_fg, command=lambda: hide_file_entry.insert(0, filedialog.askopenfilename()))
browse_button_hide_file.grid(row=1, column=5, padx=10, pady=5)

# Input for Save Location of encoded image
tk.Label(left_frame, text="Save Bind AI Image:", font=label_font, bg="#282c34", fg="white").grid(row=2, column=0, sticky="e", padx=10)
save_file_entry_hide = tk.Entry(left_frame, width=20, font=entry_font)
save_file_entry_hide.grid(row=2, column=1, padx=10, pady=5)
save_button_hide_file = tk.Button(left_frame, text="Save As", font=button_font, bg=button_color, fg=button_fg, command=lambda: save_file_entry_hide.insert(0, filedialog.asksaveasfilename(defaultextension=".png")))
save_button_hide_file.grid(row=2, column=2, padx=10, pady=5)

# Button for encoding
hide_button = tk.Button(left_frame, text="AI Bind Image", font=button_font, bg="#98c379", fg="white", command=hide_images)
hide_button.grid(row=2, column=3, columnspan=2, padx=10, pady=5)

# Add the refresh button to the left frame
refresh_button = tk.Button(left_frame, text="Refresh", font=button_font, bg="#e06c75", fg="white", command=refresh_gui)
refresh_button.grid(row=2, column=5, columnspan=2, padx=10, pady=5)

# Create a canvas for image display with scrollbar
canvas = tk.Canvas(root, bg="#000000")
canvas.grid(row=2, column=0, columnspan=3, sticky="nsew")

# Add vertical and horizontal scrollbars
vertical_scroll = tk.Scrollbar(root, orient="vertical", command=canvas.yview)
vertical_scroll.grid(row=2, column=3, sticky="ns")
canvas.configure(yscrollcommand=vertical_scroll.set)

horizontal_scroll = tk.Scrollbar(root, orient="horizontal", command=canvas.xview)
horizontal_scroll.grid(row=3, column=0, columnspan=3, sticky="ew")
canvas.configure(xscrollcommand=horizontal_scroll.set)

# Create an image label on the canvas
image_label = tk.Label(canvas, bg="#000000")
canvas.create_window((0, 0), window=image_label, anchor='nw')

# Zoom slider
zoom_slider = tk.Scale(root, from_=10, to=200, bg="lightgreen", orient='horizontal', command=update_zoom)
zoom_slider.set(100)  # Set default zoom level to 100%
zoom_slider.grid(row=1, column=0, columnspan=3, sticky='ew', padx=10)

# Set initial scroll region
canvas.config(scrollregion=canvas.bbox(tk.ALL))

# Run the application
root.mainloop()