#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Author: @abarnert, @willwade, and @morinted
# Code taken and modified from
# <https://github.com/willwade/PyUserInput/blob/master/pykeyboard/mac_keycode.py>
# <https://stackoverflow.com/questions/1918841/how-to-convert-ascii-character-to-cgkeycode>

from threading import Thread
import ctypes
import ctypes.util
import re
import struct
import unicodedata

from PyObjCTools import AppHelper
import AppKit
import Foundation

from plover import log
from plover.key_combo import CHAR_TO_KEYNAME
from plover.misc import popcount_8


carbon_path = ctypes.util.find_library('Carbon')
carbon = ctypes.cdll.LoadLibrary(carbon_path)

CFIndex = ctypes.c_int64


class CFRange(ctypes.Structure):
    _fields_ = [('loc', CFIndex),
                ('len', CFIndex)]


carbon.TISCopyCurrentKeyboardInputSource.argtypes = []
carbon.TISCopyCurrentKeyboardInputSource.restype = ctypes.c_void_p
carbon.TISCopyCurrentASCIICapableKeyboardLayoutInputSource.argtypes = []
carbon.TISCopyCurrentASCIICapableKeyboardLayoutInputSource.restype = ctypes.c_void_p

carbon.TISGetInputSourceProperty.argtypes = [ctypes.c_void_p, ctypes.c_void_p]
carbon.TISGetInputSourceProperty.restype = ctypes.c_void_p
carbon.LMGetKbdType.argtypes = []
carbon.LMGetKbdType.restype = ctypes.c_uint32
carbon.CFDataGetLength.argtypes = [ctypes.c_void_p]
carbon.CFDataGetLength.restype = ctypes.c_uint64
carbon.CFDataGetBytes.argtypes = [ctypes.c_void_p, CFRange, ctypes.c_void_p]
carbon.CFDataGetBytes.restype = None
carbon.CFRelease.argtypes = [ctypes.c_void_p]
carbon.CFRelease.restype = None

kTISPropertyUnicodeKeyLayoutData = ctypes.c_void_p.in_dll(
    carbon, 'kTISPropertyUnicodeKeyLayoutData')
kTISPropertyInputSourceID = ctypes.c_void_p.in_dll(
    carbon, 'kTISPropertyInputSourceID')
kTISPropertyInputSourceIsASCIICapable = ctypes.c_void_p.in_dll(
    carbon, 'kTISPropertyInputSourceIsASCIICapable')
COMMAND = '⌘'
SHIFT = '⇧'
OPTION = '⌥'
CONTROL = '⌃'
CAPS = '⇪'
UNKNOWN = '?'
UNKNOWN_L = '?_L'

KEY_CODE_VISUALIZATION = """
┌───────────────────Layout of Apple Extended Keyboard II───────────────────────┐
│ 53  122 120 99 118  96 97 98 100  101 109 103 111   105 107 113              │
│                                                                              │
│ 50─── 18 19 20 21 23 22 26 28 25 29 27 24 ─────51   114 115 116  71 81 75 67 │
│ 48──── 12 13 14 15 17 16 32 34 31 35 33 30 ────42   117 119 121  89 91 92 78 │
│ 57───── 0  1  2  3  5  4  38 40 37 41 39 ──────36                86 87 88 69 │
│ 56────── 6  7  8  9  11 45 46  43 47 44 ───────56       126      83 84 85 76 │
│ 59── 58─ 55─ ──────────49───────── ─55 ──58 ───59   123 125 124  82─── 65 76 │
└──────────────────────────────────────────────────────────────────────────────┘
Valid keycodes not visible on layout:
  10, 52, 54, 60-64, 66, 68, 70, 72-74, 77, 79, 80,
  90, 93-95, 102, 104, 106, 108, 110, 112, 127
"""

SPECIAL_KEY_NAMES = {
    '\x1b': 'Esc',   # Will map both Esc and Clear (key codes 53 and 71)
    '\xa0': 'nbsp',
    '\x08': 'Bksp',
    '\x05': 'Help',
    '\x01': 'Home',
    '\x7f': 'Del',
    '\x04': 'End',
    '\x0c': 'PgDn',
    '\x0b': 'PgUp',  # \x0b is also the clear character signal
}

DEFAULT_SEQUENCE = (None, 0),


def is_printable(string):
    for character in string:
        category = unicodedata.category(character)
        if category[0] in 'C':
            # Exception: the "Apple" character that most Mac layouts have
            return False if string != "" else True
        elif category == 'Zs' and character != ' ':
            return False
        elif category in 'Zl, Zp':
            return False
    return True


def get_printable_string(s):
    if s is None:
        return 'None'
    return s if is_printable(s) else SPECIAL_KEY_NAMES.setdefault(
        s, s.encode('unicode_escape').decode("utf-8")
    )


class KeyboardLayout:
    def __init__(self, watch_layout=True):
        self._char_to_key_sequence = None
        self._key_sequence_to_char = None
        self._modifier_masks = None
        self._deadkey_symbol_to_key_sequence = None

        # Spawn a thread that responds to system keyboard layout changes.
        if watch_layout:
            self._watcher = Thread(
                target=self._layout_watcher,
                name="LayoutWatcher")
            self._watcher.start()

        self._update_layout()

    def _layout_watcher(self):
        layout = self

        class LayoutWatchingCallback(AppKit.NSObject):
            def layoutChanged_(self, event):
                log.info('Mac keyboard layout changed, updating')
                layout._update_layout()

        center = Foundation.NSDistributedNotificationCenter.defaultCenter()
        watcher_callback = LayoutWatchingCallback.new()
        center.addObserver_selector_name_object_suspensionBehavior_(
            watcher_callback,
            'layoutChanged:',
            'com.apple.Carbon.TISNotifySelectedKeyboardInputSourceChanged',
            None,
            Foundation.NSNotificationSuspensionBehaviorDeliverImmediately
        )
        AppHelper.runConsoleEventLoop(installInterrupt=True)

    def _update_layout(self):
        char_to_key_sequence, key_sequence_to_char, modifier_masks = KeyboardLayout._get_layout()
        self._char_to_key_sequence = char_to_key_sequence
        self._key_sequence_to_char = key_sequence_to_char
        self._modifier_masks = modifier_masks
        self._deadkey_symbol_to_key_sequence = self._deadkeys_by_symbols()

    def deadkey_symbol_to_key_sequence(self, symbol):
        return self._deadkey_symbol_to_key_sequence.get(symbol, DEFAULT_SEQUENCE)

    def key_code_to_char(self, key_code, modifier=0):
        """Provide a key code and a modifier and it provides the character"""
        return get_printable_string(
            self._key_sequence_to_char[key_code, modifier]
        )

    def char_to_key_sequence(self, char):
        """Finds the key code and modifier sequence for the character"""
        key_sequence = self._char_to_key_sequence.get(char, DEFAULT_SEQUENCE)
        if not isinstance(key_sequence[0], tuple):
            key_sequence = key_sequence,
        return key_sequence

    def _deadkeys_by_symbols(self):
        # We store deadkeys as "characters"; dkX, where X is the symbol.
        symbols = {
            '`': '`',
            '´': '´',
            '^': ('^', 'ˆ'),
            '~': ('~', '˜'),
            '¨': '¨',
        }
        deadkeys_by_symbol = {}
        for symbol, equivalent_symbols in symbols.items():
            for equivalent_symbol in equivalent_symbols:
                sequence = self.char_to_key_sequence('dk%s' % equivalent_symbol)
                if sequence[0][0] is not None:
                    deadkeys_by_symbol[symbol] = sequence
        return deadkeys_by_symbol

    def format_modifier_header(self):
        modifiers = (
            '| {}\t'.format(KeyboardLayout._modifier_string(mod)).expandtabs(8)
            for mod in sorted(self._modifier_masks.values())
        )
        header = 'Keycode\t{}'.format(''.join(modifiers))
        return '%s\n%s' % (header, re.sub(r'[^|]', '-', header))

    def format_keycode_keys(self, keycode):
        """Returns all the variations of the keycode with modifiers"""
        keys = ('| {}\t'.format(get_printable_string(
            self._key_sequence_to_char[keycode, mod])).expandtabs(8)
                for mod in sorted(self._modifier_masks.values()))

        return '{}\t{}'.format(keycode, ''.join(keys)).expandtabs(8)

    @staticmethod
    def _modifier_dictionary(modifier_mask):
        """Provide a dictionary of active modifier keys from mod mask"""
        modifiers = {
            SHIFT: False,
            COMMAND: False,
            CONTROL: False,
            OPTION: False,
            CAPS: False,
            UNKNOWN: False,
            UNKNOWN_L: False
        }
        if modifier_mask & 16:
            modifiers[CONTROL] = True
        if modifier_mask & 8:
            modifiers[OPTION] = True
        if modifier_mask & 4:
            modifiers[CAPS] = True
        if modifier_mask & 2:
            modifiers[SHIFT] = True
        if modifier_mask & 1:
            modifiers[COMMAND] = True
        return modifiers

    @staticmethod
    def _modifier_string(modifier_mask):
        """Turn modifier mask into string representing modifiers"""
        s = ''
        modifiers = KeyboardLayout._modifier_dictionary(modifier_mask)
        for key in modifiers:
            s += key if modifiers[key] else ''
        return s

    @staticmethod
    def _get_layout():
        keyboard_input_source = carbon.TISCopyCurrentKeyboardInputSource()

        layout_source = carbon.TISGetInputSourceProperty(
            keyboard_input_source, kTISPropertyUnicodeKeyLayoutData
        )
        # Some keyboard layouts don't return UnicodeKeyLayoutData so we turn to a different source.
        if layout_source is None:
            carbon.CFRelease(keyboard_input_source)
            keyboard_input_source = carbon.TISCopyCurrentASCIICapableKeyboardLayoutInputSource()
            layout_source = carbon.TISGetInputSourceProperty(
                keyboard_input_source, kTISPropertyUnicodeKeyLayoutData
            )
        layout_size = carbon.CFDataGetLength(layout_source)
        layout_buffer = ctypes.create_string_buffer(b'\000' * layout_size)
        carbon.CFDataGetBytes(
            layout_source, CFRange(0, layout_size), ctypes.byref(layout_buffer)
        )
        keyboard_type = carbon.LMGetKbdType()
        parsed_layout = KeyboardLayout._parse_layout(
            layout_buffer, keyboard_type
        )
        carbon.CFRelease(keyboard_input_source)
        return parsed_layout

    @staticmethod
    def _parse_layout(buf, ktype):
        hf, dv, featureinfo, ktcount = struct.unpack_from('HHII', buf)
        offset = struct.calcsize('HHII')
        ktsize = struct.calcsize('IIIIIII')
        kts = [struct.unpack_from('IIIIIII', buf, offset + ktsize * i)
               for i in range(ktcount)]
        for i, kt in enumerate(kts):
            if kt[0] <= ktype <= kt[1]:
                kentry = i
                break
        else:
            kentry = 0
        ktf, ktl, modoff, charoff, sroff, stoff, seqoff = kts[kentry]

        # Modifiers
        mf, deftable, mcount = struct.unpack_from('HHI', buf, modoff)
        modtableoff = modoff + struct.calcsize('HHI')
        modtables = struct.unpack_from('B' * mcount, buf, modtableoff)
        modifier_masks = {}
        for i, table in enumerate(modtables):
            modifier_masks.setdefault(table, i)

        # Sequences
        sequences = []
        if seqoff:
            sf, scount = struct.unpack_from('HH', buf, seqoff)
            seqtableoff = seqoff + struct.calcsize('HH')
            lastoff = -1
            for soff in struct.unpack_from('H' * scount, buf, seqtableoff):
                if lastoff >= 0:
                    sequences.append(
                        buf[seqoff + lastoff:seqoff + soff].decode('utf-16'))
                lastoff = soff

        def lookupseq(key):
            if key >= 0xFFFE:
                return None
            if key & 0xC000:
                seq = key & ~0xC000
                if seq < len(sequences):
                    return sequences[seq]
            return chr(key)

        # Dead keys
        deadkeys = []
        if sroff:
            srf, srcount = struct.unpack_from('HH', buf, sroff)
            srtableoff = sroff + struct.calcsize('HH')
            for recoff in struct.unpack_from('I' * srcount, buf, srtableoff):
                cdata, nextstate, ecount, eformat = struct.unpack_from('HHHH',
                                                                       buf,
                                                                       recoff)
                recdataoff = recoff + struct.calcsize('HHHH')
                edata = buf[recdataoff:recdataoff + 4 * ecount]
                deadkeys.append((cdata, nextstate, ecount, eformat, edata))

        if stoff:
            stf, stcount = struct.unpack_from('HH', buf, stoff)
            sttableoff = stoff + struct.calcsize('HH')
            dkterms = struct.unpack_from('H' * stcount, buf, sttableoff)
        else:
            dkterms = []

        # Get char tables
        char_to_key_sequence = {}
        key_sequence_to_char = {}

        deadkey_state_to_key_sequence = {}
        key_sequence_to_deadkey_state = {}

        def favored_modifiers(modifier_a, modifier_b):
            """0 if they are equal, 1 if a is better, -1 if b is better"""
            a_favored_over_b = 0
            modifiers_a = KeyboardLayout._modifier_dictionary(modifier_a)
            modifiers_b = KeyboardLayout._modifier_dictionary(modifier_b)
            count_a = popcount_8(modifier_a)
            count_b = popcount_8(modifier_b)

            # Favor no CAPS modifier
            if modifiers_a[CAPS] and not modifiers_b[CAPS]:
                a_favored_over_b = -1
            elif modifiers_b[CAPS] and not modifiers_a[CAPS]:
                a_favored_over_b = 1
            # Favor no command modifier
            elif modifiers_a[COMMAND] and not modifiers_b[COMMAND]:
                a_favored_over_b = -1
            elif modifiers_b[COMMAND] and not modifiers_a[COMMAND]:
                a_favored_over_b = 1
            # Favor no control modifier
            elif modifiers_a[CONTROL] and not modifiers_b[CONTROL]:
                a_favored_over_b = -1
            elif modifiers_b[CONTROL] and not modifiers_a[CONTROL]:
                a_favored_over_b = 1
            # Finally, favor fewer modifiers
            elif count_a > count_b:
                a_favored_over_b = -1
            elif count_b > count_a:
                a_favored_over_b = 1
            return a_favored_over_b

        def save_shortest_key_sequence(character, new_sequence):
            current_sequence = char_to_key_sequence.setdefault(character,
                                                               new_sequence)
            if current_sequence != new_sequence:
                # Convert responses to tuples...
                if not isinstance(current_sequence[0], tuple):
                    current_sequence = current_sequence,
                if not isinstance(new_sequence[0], tuple):
                    new_sequence = new_sequence,

                first_current_better = favored_modifiers(
                    current_sequence[-1][1], new_sequence[-1][1]
                )
                last_current_better = favored_modifiers(
                    current_sequence[0][1], new_sequence[0][1]
                )

                # Favor key sequence with best modifiers (avoids a short sequence with awful modifiers)
                # e.g. ABC Extended wants ¯ from (⌘⌥a,) instead of the saner (⌥a, space)
                if first_current_better == last_current_better != 0:
                    if first_current_better < 0:
                        char_to_key_sequence[character] = new_sequence
                # Favor shortest sequence (fewer separate key presses)
                elif len(current_sequence) < len(new_sequence):
                    pass
                elif len(new_sequence) < len(current_sequence):
                    char_to_key_sequence[character] = new_sequence[0]
                # Favor fewer modifiers on last item
                elif last_current_better < 0:
                    char_to_key_sequence[character] = new_sequence
                # Favor lower modifiers on first item if last item is the same
                elif last_current_better == 0 and first_current_better < 0:
                    char_to_key_sequence[character] = new_sequence

        def lookup_and_add(key, j, mod):
            ch = lookupseq(key)
            save_shortest_key_sequence(ch, (j, mod))
            key_sequence_to_char[j, mod] = ch

        cf, csize, ccount = struct.unpack_from('HHI', buf, charoff)
        chartableoff = charoff + struct.calcsize('HHI')

        for i, table_offset in enumerate(
                struct.unpack_from('I' * ccount, buf, chartableoff)):
            mod = modifier_masks[i]
            for j, key in enumerate(
                    struct.unpack_from('H' * csize, buf, table_offset)):
                if key == 65535:
                    key_sequence_to_char[j, mod] = 'mod'
                elif key >= 0xFFFE:
                    key_sequence_to_char[j, mod] = '<{}>'.format(key)
                elif key & 0x0C000 == 0x4000:
                    dead = key & ~0xC000
                    if dead < len(deadkeys):
                        cdata, nextstate, ecount, eformat, edata = deadkeys[dead]
                        if eformat == 0 and nextstate:
                            # Initial: symbols, e.g. `, ~, ^
                            new_deadkey = (j, mod)
                            current_deadkey = deadkey_state_to_key_sequence.setdefault(
                                nextstate, new_deadkey
                            )
                            if new_deadkey != current_deadkey:
                                if favored_modifiers(current_deadkey[1],
                                                     new_deadkey[1]) < 0:
                                    deadkey_state_to_key_sequence[nextstate] = new_deadkey
                            if nextstate - 1 < len(dkterms):
                                base_key = lookupseq(dkterms[nextstate - 1])
                                dead_key_name = 'dk{}'.format(base_key)
                            else:
                                dead_key_name = 'dk#{}'.format(nextstate)
                            key_sequence_to_char[j, mod] = dead_key_name
                            save_shortest_key_sequence(dead_key_name, (j, mod))
                        elif eformat == 1 or (eformat == 0 and not nextstate):
                            # Terminal: letters, e.g. a, e, o, A, E, O
                            key_sequence_to_deadkey_state[j, mod] = deadkeys[dead]

                            lookup_and_add(cdata, j, mod)
                        elif eformat == 2:  # range
                            # TODO!
                            pass
                else:
                    lookup_and_add(key, j, mod)

        for key, dead in key_sequence_to_deadkey_state.items():
            j, mod = key
            cdata, nextstate, ecount, eformat, edata = dead
            entries = [struct.unpack_from('HH', edata, i * 4) for i in
                       range(ecount)]
            for state, key in entries:
                # Ignore if unknown state...
                if state in deadkey_state_to_key_sequence:
                    dj, dmod = deadkey_state_to_key_sequence[state]
                    ch = lookupseq(key)
                    save_shortest_key_sequence(ch, ((dj, dmod), (j, mod)))
                    key_sequence_to_char[(dj, dmod), (j, mod)] = ch

        char_to_key_sequence['\n'] = (36, 0)
        char_to_key_sequence['\r'] = (36, 0)
        char_to_key_sequence['\t'] = (48, 0)

        return char_to_key_sequence, key_sequence_to_char, modifier_masks


if __name__ == '__main__':
    layout = KeyboardLayout(False)
    print(KEY_CODE_VISUALIZATION)
    print()
    print(layout.format_modifier_header())
    for keycode in range(127):
        print(layout.format_keycode_keys(keycode))
    print()
    unmapped_characters = []
    for char, keyname in sorted(CHAR_TO_KEYNAME.items()):
        sequence = []
        for code, mod in layout.char_to_key_sequence(char):
            if code is not None:
                sequence.append(
                    (code, '{}{}'.format(
                        layout._modifier_string(mod),
                        layout.key_code_to_char(code, 0)
                    ))
                )
            else:
                unmapped_characters.append(char)
        if sequence:
            print('Name:\t\t{}\nCharacter:\t{}\nSequence:\t‘{}’\nBase codes:\t‘{}’\n'.format(
                keyname, char, '’, ‘'.join(combo[1] for combo in sequence),
                '’, ‘'.join(str(combo[0]) for combo in sequence)
            ))
    print('No mapping on this layout for characters: ‘{}’'.format(
        '’, ‘'.join(unmapped_characters)
    ))
