import collections
import cv2
import gettext
from itertools import count
import os
from send2trash import send2trash
import signal
import sys
import time
import tomllib
import tkinter as tk
from tkinter import filedialog, messagebox
import numpy as np


running_from_folder = os.path.realpath(__file__)
alphabet = ['A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z']
next_letter = 'A'

# l10n
LOCALE = os.getenv('LANG', 'en_EN')
_ = gettext.translation('template', localedir='locales', languages=[LOCALE]).gettext

# Config 
# defaults
project_settings_defaults = {
    # gphoto2 = 0, picam = 1, webcam = 2
    'camera_type': 0,
    'file_extension':'JPG',
    'trigger_mode': 'key',
    'projects_folder': '',
    # ~ 'project_letter': 'A'
    'onion_skin_onstartup' : False,
    'onionskin_alpha_default' : 0.4,
    'onionskin_fx' : False,
    'fullscreen_bool' : True,
    'screen_w' : 1920,
    'screen_h' : 1080,
    'framerate' : 16,
    'vflip' : False,
    'hflip' : False,
    'export_options' : 'scale=1920:-1,crop=1920:1080:0:102',
    'cache_images' : False,
    'liveview' : False,
}
# Load from file
config_locations = ["./", "~/.", "~/.config/"]
config_found_msg = _("No configuration file found, using defaults.")
project_settings = project_settings_defaults
for location in config_locations:
    # Optional config files, ~ is expanded to $HOME on *nix, %USERPROFILE% on windows
    if os.path.exists( os.path.expanduser(os.path.join(location, 'config.toml'))):
        with open(os.path.expanduser(location + 'config.toml'), 'rb') as config_file:
            project_settings = tomllib.load(config_file)
            if 'CHECK' in project_settings:
                camera_status = project_settings['CHECK']
            if 'CAMERA' in project_settings:
                camera_settings = project_settings['CAMERA']
            if 'DEFAULT' in project_settings:
                project_settings = project_settings['DEFAULT']
            config_found_msg = _("Found configuration file in {}").format(os.path.expanduser(location))
print(config_found_msg)

# Create blank image with 0s
blank_image = np.zeros((project_settings['screen_h'], project_settings['screen_w'],3), np.uint8)
# Set all pixels to light grey
blank_image[:,0:] = 200
font = cv2.FONT_HERSHEY_SIMPLEX
cv2.putText(blank_image,_("Pas d'image"),(int(project_settings['screen_w']/3),int(project_settings['screen_h']/2)), font, 4, (100, 100, 100), 8, cv2.LINE_AA)

def find_letter_after(letter:str):
    if letter in alphabet and alphabet.index(letter) < len(alphabet) - 1:
        return alphabet[alphabet.index(letter) + 1]
    else:
        return False


def get_projects_folder():
    if len(projects_folder):
        project_folder = projects_folder
    else:
        # Get user folder
        project_folder = os.path.expanduser('~')
        # If a project folder is defined in settings, use it
        if project_settings['projects_folder'] != '':
            subfolder = project_settings['projects_folder']
        else:
            # If it doesn't exist, use a default name
            subfolder = 'Stopmotion Projects'
        project_folder = os.path.join(project_folder, subfolder)
    # Create folder if it doesn't exist
    if os.path.exists(project_folder) == False:
        os.mkdir(project_folder)
    else:
        if not os.path.isdir(project_folder):
            # If file exists but is not a folder, can't create it, abort
            return False
    return project_folder


def get_session_folder():
    project_folder = get_projects_folder()
    if project_folder:
        sessions_list = []
        dir_list = os.listdir(project_folder)
        # Filter folders with name only one char long
        for dir in dir_list:
            if len(dir) == 1 and dir in alphabet:
                sessions_list.append(dir)
        # If folders exist, find last folder in alphabetical order
        if len(sessions_list):
            sessions_list.sort()
            last_letter = sessions_list[-1]
            # By default, find next letter for a new session
            next_letter = find_letter_after(last_letter)
            if next_letter is False:
                return False
            # A previous session folder was found; ask the user if they wish to resume session
            resume_session = tk.messagebox.askyesno(_("Resume session?"), _("A previous session was found in {}, resume shooting ?").format(os.path.join(project_folder, last_letter)))
            if resume_session:
                next_letter = last_letter
        else:
            next_letter = 'A'
        if os.path.exists(os.path.join(project_folder, next_letter)) is False:
            os.mkdir(os.path.join(project_folder, next_letter))
        print(_("Using {} as session folder.").format(os.path.join(project_folder, next_letter)))
        return os.path.join(project_folder, next_letter)
    return False


def get_frames_list(folder:str):
    # Get JPG files list in current directory
    # ~ existing_animation_files = []
    # ~ if not len(img_list):
    existing_animation_files = img_list
    # ~ existing_animation_files =
    file_list = os.listdir(folder)
    for file in file_list:
        if (file.startswith(project_letter) and file.endswith(project_settings['file_extension'])):
            if file not in existing_animation_files:
                existing_animation_files.append(file)
                # ~ existing_animation_files[file] = None
    if len(existing_animation_files) == 0:
        # If no images were found, return fake name set to -001 to init file count to 000
        return ["{}.{:04d}.{}".format(next_letter, -1, project_settings['file_extension'])]
        # ~ return {"{}.{:04d}.{}".format(project_letter, -1, project_settings['file_extension']):None}
    # ~ else:
        # Remove fake file name as soon as we have real pics
        # ~ if 'A.-001.JPG' in existing_animation_files:
            # ~ existing_animation_files.pop('A.-001.JPG')
    existing_animation_files.sort()
    # ~ existing_animation_files = collections.OrderedDict(sorted(existing_animation_files.items()))
    return existing_animation_files


def get_last_frame(folder:str):
    # Refresh file list
    existing_animation_files = get_frames_list(folder)
    # Get last file
    # Filename pattern is A.0001.JPG
    return existing_animation_files[-1].split('.')
    # ~ print(next(reversed(existing_animation_files.keys())))
    # ~ return next(reversed(existing_animation_files.keys())).split('.')


def get_onionskin_frame(folder:str, index=None):
    prev_image = get_last_frame(folder)
    prev_image = '.'.join(prev_image)
    if os.path.exists( os.path.expanduser(os.path.join(savepath, prev_image))):
        frm = cv2.imread(os.path.join(savepath, prev_image))
        frm = cv2.resize(frm, (project_settings['screen_w'], project_settings['screen_h']))
    # Imge does not exist, load blank image
    else:
        frm = blank_image
    return frm


def return_next_frame_number(last_frame_name):
    prefix, filecount, ext = last_frame_name
    filename = '.{:04d}.'.format(int(filecount)+1)
    # ~ filename = (".%04i." % x for x in count(int(filecount) + 1))
    # ~ return prefix + next(filename) + ext
    return prefix + filename + ext


def update_image(img_list, img_index):
    if len(img_list) == 0:
        return 0
    img_filename = img_list[img_index]
    if os.path.exists( os.path.expanduser(os.path.join(savepath, img_filename))):
        img = cv2.imread(os.path.join(savepath, img_filename))
        img = cv2.resize(img, (project_settings['screen_w'], project_settings['screen_h']))
    else:
        img = blank_image
    return img

def next_frame(img_index, loop=True):
    img_index = check_range(img_index+1, loop)
    return img_index, update_image(img_list, img_index)
    

def previous_frame(img_index):
    img_index = check_range(img_index-1)
    return img_index, update_image(img_list, img_index)


def clean_img_list(folder_path):
    # Check file in dict exists, else remove it
    file_list = os.listdir(folder_path)
    # Iterate over copy of dict to avoid OOR error
    img_list_copy = img_list
    for file in img_list_copy:
        if file not in file_list:
            img_list.remove(file)


def check_range(x, loop=True):
    if x < 0:
        if loop:
            return len(img_list)-1
        else:
            return 0
    elif x >= len(img_list)-1:
        if loop:
            return 0
        else:
            return len(img_list)-1
    else:
        return x


def batch_rename(folder:str):        
    # initialize counter to 0
    frame_list = get_frames_list(folder)
    counter = (".%04i." % x for x in count(0))
    # ~ for i in range(len(frame_list)):
    for i in frame_list:
        # ~ if os.path.exists(os.path.realpath(frame_list[i])):
        if os.path.exists(os.path.join(folder, i)):
            # ~ os.rename(os.path.realpath(frame_list[i]), os.path.realpath("{}{}{}".format(project_settings['project_letter'], next(counter), project_settings['file_extension'])))
            os.rename(os.path.join(folder, i), os.path.join(folder, "{}{}{}".format(project_letter, next(counter), project_settings['file_extension'])))
            # ~ print(os.path.join(folder, "{}{}{}".format(project_letter, next(counter), project_settings['file_extension'])))
        else:
            print(_("{} does not exist").format(str(i)))
    return get_frames_list(folder)


def offset_dictvalues(from_index=0):
    dict_copy = dict(img_list)
    for i in range(from_index, len(dict_copy)):
        if i < len(img_list)-1:
            img_list[list(img_list.keys())[i]] = list(img_list.values())[i+1]
        else:
            img_list[list(img_list.keys())[i]] = None


def remove_frame(img_list, img_index):
    if len(img_list):
        folder_path = os.path.realpath(savepath)
        frame_name = img_list[img_index]
        # ~ frame_path = os.path.realpath(frame_name)
        frame_path = os.path.join(folder_path, frame_name)
        if not os.path.exists(frame_path):
            return img_index, blank_image
        print(_("Removing {}").format(frame_path))
        # trash file
        send2trash(frame_path)
        # remove entry from dict
        img_list.remove(frame_name)
        # offset cached images
        # ~ offset_dictvalues(img_index)
        # rename files and get new list
        img_list = batch_rename(folder_path)
        clean_img_list(folder_path)
        # update index if possible
        img_index = check_range(img_index, False)
        # update display
        return img_index, update_image(img_list, img_index)
    else:
        return 0, blank_image

def playback_animation(img_list, img_index):
    # save onionskin state
    if onionskin:
        onionskin = False
        onionskin_was_on = True
        
    # ~ index = 0
    # Play all frames
    # ~ while index < len(img_list):
        # ~ print(img_list[index])
        # ~ img = update_image(img_list, index)
        # ~ cv2.imshow("StopiCV", img)
        # ~ time.sleep(.5)
        # ~ index += 1
    for file in img_list:
        img = update_image(img_list, img_list.index(file))
        print(str(img_list.index(file)) + " : " + file )
        cv2.imshow("StopiCV", img)
        time.sleep(.5)
    # Restore previous frame
    print(img_index)
    frame_before_playback = update_image(img_list, img_index)
    # ~ cv2.imshow("StopiCV", img)
    # Restore onionskin
    if 'onionskin_was_on' in locals():
        onionskin = True
    # Restore index
    return img_index, frame_before_playback

def testDevice(source):
    cap = cv2.VideoCapture(source) 
    if cap is None or not cap.isOpened():
        print(_("Warning: unable to open video source: {}").format(source))
        return False
    cap.release()
    return True


def signal_handler(sig, frame):
    global ctrlc_pressed
    ctrlc_pressed = True


def main(args):

    global onionskin, playback, loop_playback, playhead, index
    
    if not testDevice(0):
        print(_("No camera device found. Exiting..."))
        return 1
    cam = cv2.VideoCapture(0)
    cam.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
    cam.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
    
    cv2.namedWindow("StopiCV", cv2.WINDOW_GUI_NORMAL)
    cv2.setWindowProperty("StopiCV", cv2.WND_PROP_OPENGL, cv2.WINDOW_OPENGL)
    cv2.setWindowProperty("StopiCV", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
    cv2.setWindowProperty("StopiCV", cv2.WND_PROP_ASPECT_RATIO, cv2.WINDOW_KEEPRATIO)

    frame = get_onionskin_frame(savepath, index)
    while True:
        if playback:
            if onionskin:
                onionskin = False
                onionskin_was_on = True
            # Play all frames
            if playhead < len(img_list)-1:
                playhead, img = next_frame(playhead, loop_playback)
                cv2.imshow("StopiCV", img)
                time.sleep(1.0/project_settings['framerate'])
            else:
                playhead = index
                img = update_image(img_list, index)
                playback = False
                # Restore onionskin
                if 'onionskin_was_on' in locals():
                    onionskin = True
            loop_playback = False
            
        if onionskin:
            ret, overlay = cam.read()
            og_frame = overlay
            # Resize preview
            overlay = cv2.resize(overlay, (project_settings['screen_w'], project_settings['screen_h']))
            # ~ if overlay is not None:
            overlay = cv2.addWeighted(frame, 1.0, overlay, project_settings['onionskin_alpha_default'], 0)
            # Flip preview (0 = vert; 1 = hor)
            if project_settings['vflip'] or project_settings['hflip']:
                if project_settings['vflip'] and project_settings['hflip']:
                    flip_dir = -1
                elif project_settings['vflip']:
                    flip_dir = 0
                elif project_settings['hflip']:
                    flip_dir = 1
                frame = cv2.flip(overlay, flip_dir)
            if not ret:
                print(_("Failed to grab frame."))
                break
            cv2.imshow("StopiCV", overlay)
        
        if not playback and not onionskin:
            cv2.imshow("StopiCV", frame)

        k = cv2.waitKey(1)
        # Key o
        if (k%256 == 111) or (k%256 == 47):
            # Toggle onionskin
            onionskin = not onionskin
        # Key left
        elif (k%256 == 81) or (k%256 == 52):
            # Displau previous frame
            print("Left")
            if len(img_list):
                if playback:
                    playback = False
                index, frame = previous_frame(index)
            print("L: {}".format(index))
        elif (k%256 == 83) or (k%256 == 54):
            # Displau next frame
            print("Right")
            if len(img_list):
                if playback:
                    playback = False
                index, frame = next_frame(index)
            print("R: {}".format(index))
        # Key e / keypad *
        elif (k%256 == 101) or (k%256 == 42):
            # TODO : Export
            print("Export")
        # Key Return
        elif (k%256 == 13):
            # TODO : Playback
            print("Playback")
            playhead = index
            loop_playback = True
            playback = not playback
        # Key remove frame - backspace, del, numpad_minus
        elif (k%256 == 8) or (k%256 == 45) or (k == 255):
            # Remove frame
            print("Remove frame")   
            index, frame = remove_frame(img_list, index)
        # Quit app
        elif k%256 == 27:
            # ESC pressed
            print(_("Escape hit, exiting..."))
            break
        elif ctrlc_pressed:
            print(_("Ctrl-C hit, exiting..."))
            break
        elif cv2.getWindowProperty("StopiCV", cv2.WND_PROP_VISIBLE) < 1:
            print(_("Window was closed, exiting..."))
            break
        # Take pic
        elif (k%256 == 32) or (k%256 == 48):
            # SPACE pressed
            img_name = return_next_frame_number(get_last_frame(savepath))
            img_path = os.path.join(savepath, img_name)
            cv2.imwrite(img_path, og_frame)
            print(_("File {} written.").format(img_path))
            if len(img_list) and (img_list[index] == 'A.-001.png'):
                img_list[index] = img_name
            else:
                index += 1
            frame = get_onionskin_frame(savepath, index)
            # TODO :  go back to last frame ?
        # REMOVE : Debug print keycode
        elif k==-1:  # normally -1 returned,so don't print it
            continue
        else:
            print(k) # else print its value
    cam.release()
    cv2.destroyAllWindows()


ctrlc_pressed = False
projects_folder = project_settings['projects_folder']
img_list = []
savepath = get_session_folder()
onionskin = project_settings['onion_skin_onstartup']
playback = False
playhead = 0
loop_playback = True

if len(savepath):
    project_letter = savepath.split(os.sep)[-1]
else:
    project_letter = 'A'
img_list = get_frames_list(savepath)
index = len(img_list)-1
print(index)

if __name__ == '__main__':
    signal.signal(signal.SIGINT, signal_handler)
    sys.exit(main(sys.argv[1:]))