import os
import cv2
import json
import time
import base64
import hashlib
import keyboard
import numpy as np
import tkinter as tk
import matplotlib.pyplot as plt

from tqdm import tqdm
from time import sleep
from PIL import Image, ImageTk
from qrcode import QRCode, constants

from qrtransfert.utils import max_bytes_per_version


def chunk_byte(data, length):
  return [data[i:i + length] for i in range(0, len(data), length)]

def generate_qr_code(data, sequence_number, total_sequences):

    error_correction_map = {
        "L": constants.ERROR_CORRECT_L,
        "M": constants.ERROR_CORRECT_M,
        "Q": constants.ERROR_CORRECT_Q,
        "H": constants.ERROR_CORRECT_H
    }

    data = int(sequence_number).to_bytes(2, 'big') + int(total_sequences).to_bytes(2, 'big') + data

    # Set max verion or auto when there is less data
    version_to_use=QR_MAX_VERSION
    if len(data) < MAX_QR_CODE_SIZE:
        version_to_use=None

    qr_code = QRCode(
        version = version_to_use,
        error_correction = error_correction_map[QR_CORRECTION]
    )
    qr_code.add_data(data, optimize=0)
    qr_code.make()
    qr_code_image = qr_code.make_image(fill_color="black", back_color="white")

    qr_code_image = qr_code_image.resize((QR_GEN_SIZE, QR_GEN_SIZE)).convert("L")
    
    return qr_code_image

def on_press_q(key):
    keyboard.unhook_all()
    stop_loop()

def stop_loop():
    global STOP_LOOP
    STOP_LOOP=True

def clear_console():
    os.system('cls' if os.name=='nt' else 'clear')

def display_qrcodes_in_terminal(qrcodes):

    # Listen for press 'q'
    keyboard.on_press_key("q", on_press_q)

    image_index = -1
    while True:
        if STOP_LOOP:
            clear_console()
            break
        clear_console()
        image_index = (image_index + 1) % len(qrcodes)
        qrcode = qrcodes[image_index]
        plt.imshow(qrcode)
        plt.axis('off')
        plt.show(block=False)
        print("Press Q to exit")
        sleep(LOOP_DELAY)
        plt.close()

def display_qrcodes(qrcodes):

    # Listen for press 'q'
    keyboard.on_press_key("q", on_press_q)

    def display_next_image():

        qrcode = qrcodes[image_index]
        image = qrcode
        image = image.resize((WINDOW_WIDTH, WINDOW_HEIGHT))
        photo = ImageTk.PhotoImage(image)

        label_index = tk.Label(window, text=f"Image {image_index + 1}/{len(qrcodes)}", font=("Arial", 12))
        label_index.place(relx=0.5, y=(WINDOW_HEIGHT+25), anchor=tk.CENTER)

        label_image.configure(image=photo)
        label_image.image = photo

    window = tk.Tk()
    window.geometry(WINDOW_SIZE)
    frame_image = tk.Frame(window)
    frame_image.pack()
    label_image = tk.Label(frame_image)
    label_image.pack()

    # Listen for press 'q'
    window.protocol("WM_DELETE_WINDOW", stop_loop)

    image_index = -1
    print("Press Q to exit")
    while True:
        if STOP_LOOP:
            window.destroy()
            break
        image_index = (image_index + 1) % len(qrcodes)
        display_next_image()
        window.update_idletasks()
        window.update()
        sleep(LOOP_DELAY)

def encode(file, qr_color, terminal, qr_max_version, qr_correction, loop_delay, window_width, window_height, out_to_file):

    # Update global variables
    global QR_MODE
    QR_MODE="byte"
    global BYTE_CONTROL
    BYTE_CONTROL=4
    global STOP_LOOP
    STOP_LOOP=False
    global LOOP_DELAY
    LOOP_DELAY=loop_delay
    global QR_CORRECTION
    QR_CORRECTION=qr_correction
    global QR_MAX_VERSION
    QR_MAX_VERSION=qr_max_version
    global MAX_QR_CODE_SIZE
    MAX_QR_CODE_SIZE=max_bytes_per_version(qr_max_version, qr_correction, QR_MODE)
    global MAX_SIZE
    MAX_SIZE=(MAX_QR_CODE_SIZE - BYTE_CONTROL)
    global WINDOW_WIDTH
    WINDOW_WIDTH=window_width
    global WINDOW_HEIGHT
    WINDOW_HEIGHT=window_height
    global WINDOW_SIZE
    WINDOW_SIZE=f"{WINDOW_WIDTH}x{(WINDOW_HEIGHT+50)}"
    global QR_GEN_SIZE
    QR_GEN_SIZE=1050

    # Def vars
    total_chunks = 2
    chunks = []
    qrcodes = []

    # Convert data to base64
    filename = os.path.basename(file.name)
    file_content = file.read()
    checksum = hashlib.md5(file_content).hexdigest()
    binary_data = base64.b64encode(file_content)

    if len(binary_data) > MAX_SIZE:
        chunks = chunk_byte(binary_data, MAX_SIZE)
    else:
        chunks = [binary_data]

    total_chunks = (len(chunks) + 1)

    # Estimated minimum transfer time
    ett=round(loop_delay*(total_chunks/3))
    m, s = divmod(ett, 60)
    h, m = divmod(m, 60)
    print(f'Minimum transfer time: {h:d}h {m:02d}m {s:02d}s')

    # Generate QRCodes
    with tqdm(len(chunks)) as pbar:
        for i, chunk in enumerate(chunks):
            pbar.set_description(f"Generating QRCode {i+1}/{total_chunks}")
            qrcode = generate_qr_code(chunk, (i + 1), total_chunks)
            qrcodes.append(qrcode)

        # Generate last QRCode
        pbar.set_description(f"Generating QRCode {total_chunks}/{total_chunks}")
        json_data = {
            "filename": f"{filename}",
            "checksum": f"{checksum}"
        }
        json_string = json.dumps(json_data)
        last_data = json_string.encode("utf-8")
        last_data_base64 = base64.b64encode(last_data)
        qrcode = generate_qr_code(last_data_base64, total_chunks, total_chunks)
        qrcodes.append(qrcode)

    qr_to_display = qrcodes

    if qr_color:
        # Merge to RGB
        nb_bands = 3
        splited_qrcodes = [qrcodes[i:i+nb_bands] for i in range(0, len(qrcodes), nb_bands)]
        empty_image = Image.new('RGB', (1, 1), (255, 255, 255)).resize((QR_GEN_SIZE, QR_GEN_SIZE)).convert("L")
        chroma_qrcodes = []
        for i in splited_qrcodes:
            while len(i) != nb_bands:
                i.append(empty_image)
            chroma_qrcodes.append(Image.merge('RGB', i))

        qr_to_display = chroma_qrcodes

    # Display QRCodes
    if out_to_file:
        fps = 10
        fourcc = cv2.VideoWriter.fourcc(*'MJPG')
        video = cv2.VideoWriter(out_to_file, fourcc, fps, (QR_GEN_SIZE,QR_GEN_SIZE))
        for image in qr_to_display:
            cv_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
            for _ in range(fps*2):
                video.write(np.array(cv_img))
        cv2.destroyAllWindows()
        video.release()
    else:
        if terminal:
            display_qrcodes_in_terminal(qr_to_display)
        else:
            display_qrcodes(qr_to_display)
