## keycard-class
import hashlib
from binascii import hexlify

## ui-class
from time import sleep, sleep_ms
import display
import buttons
import sys
import system
import identification

import json

KEYS = {
    b"t": {"name": "Techie", "key": b'793f4a34debf4ed18b27'},
    b"f": {"name": "Fixer", "key": b'e85d8ea54fa115a979c8'},
    b"r": {"name": "Rocker", "key": b'd244ea7b55b566935460'},
    b"n": {"name": "Netrunner", "key": b'6942fbc3ee60b2282638'},
    b"c": {"name": "Corporate", "key": b'822423ab2f58df256c82'}
}
types = [
    (b"t", "Techie"),
    (b"f", "Fixer"),
    (b"r", "Rocker"),
    (b"n", "Netrunner"),
    (b"c", "Corporate")
]


class badge_ui:
    def __init__(self, keycard=None, test=False):
        self.credentials = []

        if test:
            self.credentials = [
                {"host": "host0", "username": "user1", "password": "password1"},
                {"host": "host1", "username": "user2", "password": "password2"}
            ]

        self.keyfile_path = "storage.key"
        self.logfile_path = "badge_logfile"

        self.memory_index = 0
        self.scroll_index = 0
        self.multiline_window_height = 1
        self.font = "Roboto_Regular12"

        self.keycard = keycard

        try:
            self.load_credentials()
        except OSError:
            pass

    def word_wrap(self, text, font="Roboto_Regular12"):
        screen_width = display.width()
        text_width = display.getTextWidth(text, font)
        lines = []
        if screen_width > text_width:
            lines.append(text)
        else:
            tmp = ""
            buf = ""
            lines = []
            for char in text:
                tmp += char
                if screen_width < display.getTextWidth(tmp, font):
                    lines.append(buf)
                    buf = ""
                    tmp = "" + char
                if char == '\r':
                    lines.append(buf)
                    buf = ""
                    tmp = ""
                buf += char
            if len(buf) > 0:
                lines.append(buf)
        return lines

    def drawMultiline(self, lines, font="Roboto_Regular12"):
        # Use word_wrap() to split long texts into multiple lines

        self.scroll_index = 0
        self.scroll_lines = lines
        self.font = font

        y_diff = display.getTextHeight("I", font) + 1
        y_coord = 0

        self.multiline_window_height = int(display.height() / y_diff) - 1  # Number of lines for screen

        display.drawFill(0x000000)  # Fill the screen with black

        header = "KEY " + str(self.memory_index + 1) + "/" + str(len(self.credentials)) + "    " + "ROW " + str(self.scroll_index + 1) + "/" + str(len(self.scroll_lines))
        display.drawText(0, y_coord, header, 0xFFFFFF, font)
        y_coord += y_diff

        for line in self.scroll_lines[self.scroll_index:self.multiline_window_height + self.scroll_index]:
            display.drawText(0, y_coord, line, 0xFFFFFF, font)
            y_coord += y_diff

        display.flush()  # Write the contents of the buffer to the display

        buttons.attach(buttons.BTN_UP, self.multiline_down)
        buttons.attach(buttons.BTN_DOWN, self.multiline_up)

    def multiline_up(self, pressed):
        if pressed:

            if self.scroll_index < (len(self.scroll_lines) - self.multiline_window_height):
                self.scroll_index += 1

            y_diff = display.getTextHeight("I", self.font) + 1
            y_coord = 0

            display.drawFill(0x000000)  # Fill the screen with black

            header = "KEY " + str(self.memory_index + 1) + "/" + str(len(self.credentials)) + "    " + "ROW " + str(self.scroll_index + 1) + "/" + str(len(self.scroll_lines))
            display.drawText(0, y_coord, header, 0xFFFFFF, self.font)
            y_coord += y_diff

            for line in self.scroll_lines[self.scroll_index:self.multiline_window_height + self.scroll_index]:
                display.drawText(0, y_coord, line, 0xFFFFFF, self.font)
                y_coord += y_diff

            display.flush()  # Write the contents of the buffer to the display

    def multiline_down(self, pressed):
        if pressed:
            if self.scroll_index > 0:
                self.scroll_index -= 1

            y_diff = display.getTextHeight("I", self.font) + 1
            y_coord = 0

            display.drawFill(0x000000)  # Fill the screen with black

            header = "KEY " + str(self.memory_index + 1) + "/" + str(len(self.credentials)) + "    " + "ROW " + str(self.scroll_index + 1) + "/" + str(len(self.scroll_lines))
            display.drawText(0, y_coord, header, 0xFFFFFF, self.font)
            y_coord += y_diff

            for line in self.scroll_lines[self.scroll_index:self.multiline_window_height + self.scroll_index]:
                display.drawText(0, y_coord, line, 0xFFFFFF, self.font)
                y_coord += y_diff

            display.flush()  # Write the contents of the buffer to the display

    def memory_up(self, pressed):
        if pressed:
            if self.memory_index < len(self.credentials) - 1:
                self.memory_index += 1
            else:
                self.memory_index = 0
        self.refresh()

    def memory_down(self, pressed):
        if pressed:
            if self.memory_index > 0:
                self.memory_index -= 1
            else:
                self.memory_index = len(self.credentials) - 1
        self.refresh()

    def show_creds(self, credentials=None):
        display.drawFill(0x000000)  # Fill the screen with black
        if credentials is None:
            display.drawText(10, 2, "HELLO " + self.keycard.name, 0xFFFFFF, "Roboto_Regular12")
            display.drawLine(5, 16, 123, 16, 0xFFFFFF)  # Draw a white line from (10,10) to (20,20)
            display.drawText(10, 18, "EMPTY MEMORY", 0xFFFFFF, "Roboto_Regular12")
            display.drawText(10, 32, "A: TRANSMIT", 0xFFFFFF, "Roboto_Regular12")
            display.drawText(10, 44, "B: BACK", 0xFFFFFF, "Roboto_Regular12")
        elif "host" in credentials:
            host = credentials["host"]
            username = credentials["username"]
            password = credentials["password"]
            display.drawText(10, 5, host, 0xFFFFFF, "Roboto_Regular12")
            display.drawLine(5, 30, 123, 30, 0xFFFFFF)  # Draw a white line from (10,10) to (20,20)
            display.drawText(10, 35, username, 0xFFFFFF, "Roboto_Regular12")
            display.drawText(10, 50, password, 0xFFFFFF, "Roboto_Regular12")
        else:
            self.drawMultiline(self.word_wrap(credentials))
        display.flush()  # Write the contents of the buffer to the display

    def send(self, pressed):
        if pressed:
            logfile = open(self.logfile_path, 'w')
            rx = False
            retry_count = 5
            res = "WAITING RESPONSE"

            while rx is False and retry_count > 0:
                display.drawFill(0x000000)  # Fill the screen with black
                display.drawText(10, 5, "SENDING", 0xFFFFFF, "Roboto_Regular18")
                display.drawLine(5, 30, 123, 30, 0xFFFFFF)  # Draw a white line from (10,10) to (20,20)
                display.drawText(10, 35, str(retry_count), 0xFFFFFF, "Roboto_Regular12")
                display.drawText(10, 50, res, 0xFFFFFF, "Roboto_Regular12")
                display.flush()  # Write the contents of the buffer to the display
                c.send()
                sleep(0.1)
                c.send()
                sleep(0.1)
                rx = sys.stdin.readline()  # !!! This is blocking call, better solution needed.
                logfile.write(rx)
                logfile.write(b"\n")

                if rx[0] == "#":  # Potential credential data received
                    self.credentials.append(rx)
                    self.store_credentials()
                    break
                else:
                    res = rx
                    rx = False
                    retry_count -= 1

            if rx is False:
                display.drawFill(0x000000)  # Fill the screen with black
                display.drawText(10, 5, "TIMEOUT", 0xFFFFFF, "Roboto_Regular18")
                display.drawLine(5, 30, 123, 30, 0xFFFFFF)  # Draw a white line from (10,10) to (20,20)
                display.drawText(10, 35, "RECEIVED NO", 0xFFFFFF, "Roboto_Regular12")
                display.drawText(10, 50, "CREDENTIALS", 0xFFFFFF, "Roboto_Regular12")
                display.flush()  # Write the contents of the buffer to the display
                sleep_ms(5000)

            logfile.close()
            self.refresh()

    def store_credentials(self):
        with open(self.keyfile_path, 'w') as f:
            f.write(json.dumps({"credentials": self.credentials}))

    def load_credentials(self):
        with open(self.keyfile_path, 'r') as f:
            try:
                d = json.load(f)
                if "credentials" in d:
                    self.credentials = d["credentials"]
            except ValueError:
                pass

    def refresh(self, pressed=True):
        if pressed:
            if len(self.credentials) == 0:
                self.show_creds(None)
            else:
                self.show_creds(self.credentials[self.memory_index])

    def exit(self, pressed):
        if pressed:
            system.home()


class keycard:
    def __init__(self):
        ident = identification.detect_type()
        self.i = 0
        if "type" in ident and "name" in ident:
            self.type = ident["type"]
            self.name = ident["name"]
            self.key = KEYS[self.type]["key"]
        else:
            self.type = None
            self.name = "Unidentified"
            self.key = None

    def send(self):
        """
        Sends type and key via serial
        """
        m = hashlib.sha256()
        m.update(self.type + self.key)
        checksum = m.digest()
        packet = b"# " + self.type + self.key + checksum
        print(str(hexlify(packet), 'utf-8'))

    def impersonate(self, pressed=True):
        if pressed:
            self.i += 1
            if self.i >= len(types):
                self.i = 0
            self.type = types[self.i][0]
            self.name = types[self.i][1]
            self.key = KEYS[self.type]["key"]
            display.drawFill(0x000000)  # Fill the screen with black
            display.drawText(10, 10, "switched to", 0xFFFFFF, "Roboto_Regular12")
            display.drawText(10, 35, self.name, 0xFFFFFF, "Roboto_Regular12")
            display.drawText(10, 50, "(B to go back)", 0xFFFFFF, "Roboto_Regular12")
            display.flush()  # Write the contents of the buffer to the display


c = keycard()
b = badge_ui(keycard=c)
buttons.attach(buttons.BTN_B, b.refresh)
buttons.attach(buttons.BTN_SELECT, c.impersonate)
buttons.attach(buttons.BTN_START, b.exit)
buttons.attach(buttons.BTN_A, b.send)
buttons.attach(buttons.BTN_RIGHT, b.memory_up)
buttons.attach(buttons.BTN_LEFT, b.memory_down)
b.refresh()