#!/usr/bin/env python3
import os
import time
import numpy as np
from Xlib import display, X
from PIL import Image


COLUMNS = 7
WINDOW_WIDTH = 1600
WINDOW_HEIGHT = 900
BLOCK_SIZE = 60
BOARD_X = 367
BOARD_Y = 129
BOARD_WIDTH = COLUMNS * BLOCK_SIZE
BOARD_HEIGHT = 638

MAX_COLUMN_HEIGHT = 546
DETECT_COLUMN_OFFSET_X = 8, 50
DETECT_COLUMN_OFFSET_Y = -11
MIN_COLUMN_SAT = 130
MIN_COLUMN_VAL = 120

COLUMN_VSHIFT = [-2, -2, -1, -1, 0, 0, 0]

RED, PINK, GREEN, BLUE, YELLOW, NONE = range(6)
BOMB_OFFSET = NONE + 1
BASIC_HUES = [248, 224, 118, 158, 26]
BOMB_HUES = [250, 219, 131, 174, 38]
HUE_TOLERANCE = 5
DETECT_BASIC_X = 9
DETECT_BASIC_Y = 15
DETECT_BOMB_X = 25
DETECT_BOMB_Y = 44
MIN_BASIC_SAT = 180

DETECT_EXA_X = 30
DETECT_EXA_Y = 547
EXA_HUE = 129
EXA_MIN_VAL = 240
DETECT_HELD_Y = 579
DETECT_EXA_LIGHT_X = 30
DETECT_EXA_LIGHT_Y = 609
EXA_LIGHT_HUE = 226
MIN_EXA_LIGHT_VAL = 180


def is_hue(h, hexpect):
    return abs(h - hexpect) <= HUE_TOLERANCE


def find_window(name):
    def traverse(window):
        if window.get_wm_name() == name:
            return window

        for child in window.query_tree().children:
            win = traverse(child)
            if win:
                return win

    return traverse(display.Display().screen().root)


def get_exapunks_window():
    win = find_window('EXAPUNKS')
    assert win, 'EXAPUNKS window not found'
    geo = win.get_geometry()
    assert geo.width == WINDOW_WIDTH
    assert geo.height == WINDOW_HEIGHT
    return win


def focus_window(window):
    window.set_input_focus(X.RevertToNone, X.CurrentTime)
    window.raise_window()
    display.Display().sync()


def screenshot_board(window):
    start = time.time()
    raw = window.get_image(BOARD_X, BOARD_Y, BOARD_WIDTH, BOARD_HEIGHT,
                           X.ZPixmap, 0xffffffff)
    dim = BOARD_WIDTH, BOARD_HEIGHT
    im = Image.frombytes('RGB', dim, raw.data, 'raw', 'BGRX')
    return im.convert('HSV')


def detect_columns(board):
    def saturated(x, y):
        h, s, v = board.getpixel((x, y))
        return s > MIN_COLUMN_SAT and v > MIN_COLUMN_VAL

    a, b = DETECT_COLUMN_OFFSET_X

    for y in range(MAX_COLUMN_HEIGHT, BLOCK_SIZE, -1):
        for col in range(COLUMNS):
            x = col * BLOCK_SIZE
            if saturated(x + a, y) and saturated(x + b, y):
                #print('found bottom in column', col)
                return y + 1 - DETECT_COLUMN_OFFSET_Y + COLUMN_VSHIFT[col]


def detect_block_type(board, x, y):
    h, s, v = board.getpixel((x + DETECT_BASIC_X,
                              y + DETECT_BASIC_Y))

    # check for basic blocks first, use saturation filter to avoid confusing
    # green blocks with background
    if s >= MIN_BASIC_SAT:
        for ty, hexpect in enumerate(BASIC_HUES):
            if is_hue(h, hexpect):
                return ty

    # if no basic block is detected, check another pixel for bomb contents
    h, s, v = board.getpixel((x + DETECT_BOMB_X,
                              y + DETECT_BOMB_Y))
    for ty, hexpect in enumerate(BOMB_HUES):
        if is_hue(h, hexpect):
            return ty + BOMB_OFFSET

    # no basic block or bomb -> empty slot
    return NONE


def detect_blocks(board):
    maxy = detect_columns(board) - BLOCK_SIZE
    for y in range(maxy, 0, -BLOCK_SIZE):
        for col in range(COLUMNS):
            x = col * BLOCK_SIZE
            yield detect_block_type(board, x, y + COLUMN_VSHIFT[col])


def detect_exa(board):
    for col in range(COLUMNS):
        x = col * BLOCK_SIZE + DETECT_EXA_X
        y = DETECT_EXA_Y + COLUMN_VSHIFT[col]
        h, s, v = board.getpixel((x, y))
        if is_hue(h, EXA_HUE) and v >= EXA_MIN_VAL:
            return col


def detect_held(board, exa):
    if exa is not None:
        x = exa * BLOCK_SIZE + DETECT_EXA_LIGHT_X
        y = DETECT_EXA_LIGHT_Y + COLUMN_VSHIFT[exa]
        h, s, v = board.getpixel((x, y))
        if not is_hue(h, EXA_LIGHT_HUE) or v < MIN_EXA_LIGHT_VAL:
            return detect_block_type(board, exa * BLOCK_SIZE, DETECT_HELD_Y)

    return NONE


def print_board(blocks, exa, held):
    rows = len(blocks) // COLUMNS
    for row in range(rows):
        row = rows - row
        row_blocks = blocks[(row - 1) * COLUMNS:row * COLUMNS]
        print(' ' + ''.join('rpgby.RPGBY'[ty] for ty in row_blocks))

    if exa is not None:
        print(' ' * exa + ' |')
        print('-' * exa, '(', 'rpgby RPGBY'[held], ')',
              '-' * (COLUMNS - exa - 1), sep='')


if __name__ == '__main__':
    win = get_exapunks_window()
    win.raise_window()

    while True:
        board = screenshot_board(win)
        blocks = list(detect_blocks(board))
        exa = detect_exa(board)
        held = detect_held(board, exa)
        print('\033c', end='')
        print_board(blocks, exa, held)
        time.sleep(0.05)