import os
import curses
import math

from contextlib import contextmanager

BIT_PAD_WIDTH = 40

RED = curses.COLOR_RED
BLACK = curses.COLOR_BLACK
WHITE = curses.COLOR_WHITE
CYAN = curses.COLOR_CYAN
BLUE = curses.COLOR_BLUE
GREEN = curses.COLOR_GREEN
MAGENTA = curses.COLOR_MAGENTA
YELLOW = curses.COLOR_YELLOW

HIGHLIGHTED_BIT_COLOR = 11
NO_BIT_COLOR = 10
HIGHLIGHTED_CELL_COLOR = 9
HEADER_COLOR = 8
BASIC_TEXT_COLOR = 7

HL_1, HL_2, HL_3, HL_4, HL_5, HL_6 = [20 + i for i in range(6)]
HL_COLORS = [HL_1,HL_2,HL_3,HL_4,HL_5,HL_6]

CMD_ABD = 'abd'
CMD_AB = 'ab'
CMD_C = 'c'
CMD_EN = 'en'
CMD_ENX = 'enx'
CMD_HIGHLIGHT = 'hl'
CMD_FILTER = 'filt'
CMD_CALC = 'calc'

BITS_TO_SHOW = 20

@contextmanager
def justdoit():
    try:
        yield
    except:
        pass

# Funcions for calc
def isprime(num):
    for i in range(2, isqrt(num)):
        if num % i == 0:
            return False
    return True

def isqrt(num):
    return int(math.sqrt(num))

def factors(num):
    factors = []
    cur_num = num
    i = 2
    irt = isqrt(num)
    while i < irt:
        if num % i == 0:
            factors.append(i)
            cur_num = cur_num // i
        else:
            i += 1
    return factors

EVAL_FUNCS = [isprime, isqrt, factors]


class RecordList:

    def __init__(self, height, width):
        self.height = height
        self.width = width
        self.indices = ['e', 'n', 'd', 'x', 'a', 'b', 'c', 'f', 'd+n', 'x+n', 'xx+e']
        self.filters = []
        self.highlights = []
        self.history = []
        self.history_index = len(self.history) - 1
        self.set_enx(1, 1, 1)

    def gen_records(self):
        h2 = self.height // 2
        bottom = h2
        top = h2
        if self.height % 2:
            top += 1

        cur_x = self.x - 1
        records = []

        found = 0
        while(found < bottom):
            r = self.get_record(self.e, self.n, cur_x)
            if r and self._passes_filter(r):
                records.insert(0, r)
                found += 1
            cur_x -= 1

        found = 0
        cur_x = self.x
        while(found < top):
            r = self.get_record(self.e, self.n, cur_x)
            if r and self._passes_filter(r):
                records.append(r)
                found += 1
            cur_x += 1
        self.records = records
        self.calc_params()

    def add_filter(self, filter):
        for i in self.indices:
            locals()[i] = 1
        for func in EVAL_FUNCS:
            locals()[func.__name__] = func
        try:
            eval(filter)
            self.filters.append(filter)
            return True
        except:
            return False

    def add_highlight(self, highlight):
        for i in self.indices:
            locals()[i] = 1
        for func in EVAL_FUNCS:
            locals()[func.__name__] = func
        try:
            eval(highlight)
            self.highlights.append(highlight)
            return True
        except:
            return False

    def _passes_filter(self, record):
        for i in range(len(self.indices)):
            locals()[self.indices[i]] = record[i]
        for func in EVAL_FUNCS:
            locals()[func.__name__] = func
        for filt in self.filters:
            if not eval(filt):
                return False
        return True

    def is_row_highlighted(self, row):
        record = self.records[row]
        for i in range(len(self.indices)):
            locals()[self.indices[i]] = record[i]
        for func in EVAL_FUNCS:
            locals()[func.__name__] = func

        out = -1
        for i in range(len(self.highlights)):
            hl = self.highlights[i]
            if eval(hl):
                out = i
        return out


    def shift_index(self, amt):
        cur_dex = (self.height // 2)
        new_dex = cur_dex + amt
        x = self.records[new_dex][3]
        self.set_enx(self.e, self.n, x)

    def rec_for_row(self, screen_row):
        return self.records[screen_row]

    def __format_text(self, text, index):
        w = self.params[index]

        if len(text) < w:
            text = ' '*(w-len(text)) + text
        else:
            text = text[-w:]
        return text

    def text_header(self, index):
        content = str(self.indices[index])
        return self.__format_text(content, index)

    def text_for_row(self, screen_row, index):
        content = str(self.rec_for_row(screen_row)[index])
        return self.__format_text(content, index)

    def set_enx(self, e, n, x, add_history=True):
        self.e = e
        self.n = n
        self.x = x

        if add_history:
            self.history = self.history[:self.history_index + 1]
            self.history.append([e, n, x])
            self.history_index = len(self.history) - 1

        self.gen_records()

    def go_back(self):
        nind = self.history_index - 1
        if nind > 0:
            e, n, x = self.history[nind]
            self.set_enx(e, n, x, add_history=False)
            self.history_index -= 1

    def go_forward(self):
        nind = self.history_index + 1
        if nind < len(self.history):
            e, n, x = self.history[nind]
            self.set_enx(e, n, x, add_history=False)
            self.history_index += 1

    def calc_params(self):
        max_widths = []
        num_elements = len(self.indices)

        MIN_WIDTH = 3

        for i in range(num_elements):
            max_len = MIN_WIDTH
            for rec in self.records:
                val = rec[i]
                strlen = len(str(val))
                if strlen > max_len:
                    max_len = strlen + 1
            max_widths.append(max_len)

        FILL = 1000
        if FILL:
            total = sum(max_widths)
            if total < self.width:
                for i in range(min(self.width - total, FILL * len(max_widths))):
                    max_widths[i%len(max_widths)] += 1
        self.params = max_widths

    def get_record(self, e, n, x):
        if e%2 != x%2:
            return

        na2 = x * x + e
        if na2 % 2 == 1:
            return

        na = (x * x + e) // 2
        if na % n > 0:
            return

        a = na // n
        # e, n, x, a
        d = a + x
        b = a + 2 * (x + n)

        f = 2 * d + 1 - e
        dn = d + n
        xn = x + n
        xxe = x * x + e
        c = a * b
        # self.indices = ['e', 'n', 'd', 'x', 'a', 'b', 'c', 'f', 'd+n', 'x+n', 'xx+e']
        return [e, n, d, x, a, b, c, f, dn, xn, xxe]




class Console:

    def __init__(self):
        x = os.get_terminal_size()
        self.width = x[0]
        self.height = x[1]
        self.set_up_curses()
        self.bit_index = 4 # Position of 'a' in (e, n, d, x, a, b)
        self.filters = []

    @property
    def current_record(self):
        return self.record_list_view.records[len(self.record_list_view.records) // 2]

    def set_up_curses(self):
        self.base_screen = curses.initscr()
        curses.start_color()
        curses.use_default_colors()

        curses.curs_set(0)
        curses.noecho()
        curses.cbreak()

        self.base_screen.keypad(1)
        self.set_up_colors()

        list_view_height = self.height - 2
        list_view_width = self.width - BIT_PAD_WIDTH

        self.record_list_view = RecordList(list_view_height, list_view_width)
        self.bit_pad = curses.newpad(list_view_height, BIT_PAD_WIDTH)
        self.record_pad = curses.newpad(list_view_height, list_view_width)

    def end_screen(self):
        curses.nocbreak()
        curses.echo()
        self.base_screen.keypad(0)
        curses.endwin()

    def set_up_colors(self):
        curses.init_pair(HIGHLIGHTED_BIT_COLOR, WHITE, WHITE)
        curses.init_pair(NO_BIT_COLOR, BLACK, BLACK)
        curses.init_pair(HIGHLIGHTED_CELL_COLOR, WHITE, CYAN)
        curses.init_pair(HEADER_COLOR, BLACK, MAGENTA)
        curses.init_pair(BASIC_TEXT_COLOR, WHITE, BLACK)

        curses.init_pair(HL_1, WHITE, RED)
        curses.init_pair(HL_2, WHITE, BLUE)
        curses.init_pair(HL_3, WHITE, GREEN)
        curses.init_pair(HL_4, WHITE, YELLOW)
        curses.init_pair(HL_5, WHITE, MAGENTA)
        curses.init_pair(HL_6, WHITE, CYAN)

    def main_loop(self):

        big_shift_size = 5

        while True:
            ch = self.base_screen.getch()

            if ch == ord('j'):
                self.record_list_view.shift_index(1)
            if ch == ord('k'):
                self.record_list_view.shift_index(-1)
            if ch == ord('J'):
                self.record_list_view.shift_index(big_shift_size)
            if ch == ord('K'):
                self.record_list_view.shift_index(-big_shift_size)
            if ch == ord('h'):
                self.bit_index = (self.bit_index - 1) % len(self.record_list_view.indices)
            if ch == ord('l'):
                self.bit_index = (self.bit_index + 1) % len(self.record_list_view.indices)
            if ch == ord('p'):
                self.record_list_view.go_back()
            if ch == ord('n'):
                self.record_list_view.go_forward()
            if ch == ord('D'):
                self.change_d(1)
            if ch == ord('d'):
                self.change_d(-1)
            if ch == ord(':'):
                self.enter_input()
            if ch == ord('Q'):
                self.end_screen()
                return
            self.refresh_screen()

    def clear_input_row(self):
        with justdoit():
            self.base_screen.addstr(self.height-1, 0, ':' + ' ' * (self.width - 1))
            self.refresh_screen()

    def write_output_row(self, text):
        with justdoit():
            self.base_screen.addstr(self.height-2, 0, text + ' '*(self.width - len(text)))
            self.refresh_screen()

    def change_d(self, amt):
        e, n, d, x, a, b = self.current_record[:6]
        c = a * b
        new_d = d + amt
        e = c - new_d**2
        x = new_d - a
        n = (x * x + e) // 2
        self.record_list_view.set_enx(e, n, x)

    def cmd_success(self):
        self.write_output_row("Success")

    def cmd_fail(self):
        self.write_output_row("Failure")

    def enter_input(self):
        row = self.height - 1
        self.clear_input_row()
        cur_text = ''
        ch = self.base_screen.getch()
        while(ch != curses.KEY_ENTER or ch != 10):
            # For some reason backspace wasn't working for me so I use the raw value
            # if it isn't working for you uncomment the line with !!! and then
            # # start the app and type ":" (just the colon). Then BEFORE YOU PRESS ANYTHING
            # ELSE type the key you want to find the value for (probably BACKSPACE or ENTER)
            # then look for "Exception: N" and N should be the number that you put where 127 is
            #
            # raise Exception(ch) # !!!

            if ch == 127 or ch == curses.KEY_BACKSPACE:
                cur_text = cur_text[:-1]
            elif ch == 10 or ch == curses.KEY_ENTER:
                try:
                    message = self.accept_command(cur_text)
                except Exception as e:
                    self.write_output_row(str(e))
                # Print message at this point
                return
            else:
                cur_text += chr(ch)

            with justdoit():
                self.base_screen.addstr(row, 0, ':' + cur_text + ' ') # ' ' added for backspace function

            ch = self.base_screen.getch()

    def accept_command(self, text):
        args = text.strip().split()

        command = args[0]
        rest = args[1:]

        if command in [CMD_C, CMD_AB, CMD_ABD]:
            if command == CMD_C:
                a, b = 1, int(rest[0])
            else:
                a, b = [int(x) for x in rest[:2]]

            if command == CMD_ABD:
                d = int(rest[2])
            else:
                d = int(math.sqrt(a*b))
            e = (a * b) - (d * d)
            x = d - a
            n = ((b - a) // 2) - x
            self.record_list_view.set_enx(e, n, x)
            return

        if command in [CMD_EN, CMD_ENX]:
            e, n = [int(x) for x in rest[:2]]
            if command == CMD_ENX:
                x = int(rest[2])
            else:
                x = self.record_list_view.x
                while self.record_list_view.get_record(e, n, x) is None:
                    x += 1
            self.record_list_view.set_enx(e, n, x)

        if command == CMD_FILTER:
            text = ''.join(rest)
            if text.strip() == '-':
                self.record_list_view.filters.pop()
                return
            if text.strip() == '--':
                self.record_list_view.filters = []
                return
            success = self.record_list_view.add_filter(text)
            if success:
                self.cmd_success()
            else:
                self.cmd_fail()

        if command == CMD_HIGHLIGHT:
            text = ''.join(rest)
            if text.strip() == '-':
                self.record_list_view.highlights.pop()
                return
            if text.strip() == '--':
                self.record_list_view.highlights = []
                return
            success = self.record_list_view.add_highlight(text)
            if success:
                self.cmd_success()
            else:
                self.cmd_fail()

        if command == CMD_CALC:
            text = ''.join(rest)
            output = 'Error'
            for func in EVAL_FUNCS:
                locals()[func.__name__] = func
            output = eval(text)
            self.write_output_row(str(output))


    def refresh_screen(self):
        self.refresh_list_view()
        self.refresh_bit_pad()

    def refresh_list_view(self):
        self.record_pad.clear()
        for row in range(self.record_list_view.height):
            row_str = ''

            base_color = BASIC_TEXT_COLOR
            mods = 0
            if row == 0:
                base_color = HEADER_COLOR
            if row == self.record_list_view.height // 2:
                mods = curses.A_UNDERLINE

            hl = self.record_list_view.is_row_highlighted(row)
            if hl != -1:
                base_color = HL_COLORS[hl % len(HL_COLORS)]

            def get_val(index):
                if row == 0:
                    return self.record_list_view.text_header(index)
                else:
                    return self.record_list_view.text_for_row(row, index)

            cur_col = 0
            for i in range(len(self.record_list_view.indices)):
                color = HIGHLIGHTED_CELL_COLOR if i == self.bit_index else base_color
                string_val = get_val(i) + '|'
                with justdoit():
                    self.record_pad.addstr(row, cur_col, string_val,
                                           curses.color_pair(color) | mods)
                    cur_col += len(string_val)

        refresh_args = (0, 0, 0, 0, self.record_list_view.height, self.record_list_view.width)
        self.record_pad.refresh(*refresh_args)

    def refresh_bit_pad(self):
        self.bit_pad.clear()
        start_val = 1 << (BIT_PAD_WIDTH - 1)

        for row in range(self.record_list_view.height):
            if row == 0:
                continue
            checker = start_val
            row_val = self.record_list_view.rec_for_row(row)[self.bit_index]
            index = 0

            hl = self.record_list_view.is_row_highlighted(row)
            base_color = NO_BIT_COLOR
            if hl != -1:
                base_color = HL_COLORS[hl % len(HL_COLORS)]

            while checker > 0:
                with justdoit():
                    if row_val & checker > 0:
                        self.bit_pad.addch(row, index, ord('X'), curses.color_pair(HIGHLIGHTED_BIT_COLOR))
                    else:
                        self.bit_pad.addch(row, index, ord(' '), curses.color_pair(base_color))
                index += 1
                checker = checker >> 1

        refresh_args = (0, 0, 0, self.record_list_view.width, \
                        self.record_list_view.height, \
                        self.record_list_view.width + BIT_PAD_WIDTH)
        self.bit_pad.refresh(*refresh_args)

def main():
    c = Console()
    c.main_loop()


if __name__ == '__main__':
    main()