import pygame as pg
from pygame import Rect, Vector2 as v2
from dataclasses import dataclass, field
from enum import Enum
from typing import Set

import os
from websudoku import *

@dataclass
class Cursor:
    row: int = 5
    col: int = 5

@dataclass
class Cell:
    value: int = 0
    given: bool = False

@dataclass
class PencilMark:
    center: Set[int] = field(default_factory=set)
    border: Set[int] = field(default_factory=set)

class InputMethod(Enum):
    """
    This enum helps representing Snyder notation in the cells
    """
    FILL = 1,
    CENTER = 2,
    BORDER = 3


##################
# PYGAME INIT
##################
oswinx, oswiny = 1940,50
os.environ['SDL_VIDEO_WINDOW_POS'] = "%d,%d" % (oswinx,oswiny)
pg.init()

screen = pg.display.set_mode((900, 900))
pg.display.set_caption("Pydoku")
pg.key.set_repeat(200, 35)

clock = pg.time.Clock()
running = True

##################
# FONTS
##################
title_font = pg.font.Font(None, 64)
filled_num_font = pg.font.Font(None, 60)
btn_font = pg.font.Font(None, 28)
pm_border_font = pg.font.Font(None, 18)
pm_center_font = pg.font.Font(None, 32)
title = title_font.render("Pydoku", True, "black")

##################
# GAME
##################
CELL_SIZE = 60
CURSOR_SIZE = CELL_SIZE + 6
GRID_X = screen.get_width() / 2 - 9 * CELL_SIZE // 2
GRID_Y = screen.get_height() / 2 - 9 * CELL_SIZE // 2

cursorSurface = pg.Surface((CURSOR_SIZE, CURSOR_SIZE), pg.SRCALPHA)

cursor = Cursor()
input_method = InputMethod.FILL

# TODO: We have to make this async
sudoku = get_sudoku_puzzle(Difficulty.EVIL)
# sudoku = get_sudoku_puzzle_test()

board = []

for i in range(81):
    if sudoku.editmask[i] == 0:
        board.append(Cell(sudoku.solution[i], True))
    else:
        board.append(Cell())

pencil_marks = []
for i in range(81):
    pm = PencilMark() if not board[i].given else None
    pencil_marks.append(pm)


def draw_grid():
    for row in range(9):
        for col in range(9):
            rect = Rect(CELL_SIZE * col + GRID_X,
                        CELL_SIZE * row + GRID_Y,
                        CELL_SIZE, CELL_SIZE)
            rect = Rect(CELL_SIZE * col + GRID_X,
                        CELL_SIZE * row + GRID_Y,
                        CELL_SIZE, CELL_SIZE)

            pg.draw.rect(screen, "white", rect)
            pg.draw.rect(screen, "gray", rect, width=1)
            for l in range(4):
                pg.draw.line(screen, "black",
                             start_pos=(GRID_X, GRID_Y + CELL_SIZE * 3 * l),
                             end_pos  =(GRID_X + CELL_SIZE * 9, GRID_Y + CELL_SIZE * 3 * l),
                             width=2)
            for l in range(4):
                pg.draw.line(screen, "black",
                             start_pos=(GRID_X + CELL_SIZE * 3 * l, GRID_Y),
                             end_pos  =(GRID_X + CELL_SIZE * 3 * l, GRID_Y + CELL_SIZE * 9),
                             width=2)

def draw_cursor():
    pad_diff = (CURSOR_SIZE - CELL_SIZE) // 2
    pg.draw.rect(cursorSurface, (100, 0, 155, 100), [0, 0, CURSOR_SIZE, CURSOR_SIZE])
    rect = Rect(CELL_SIZE * (cursor.col - 1) + GRID_X - pad_diff,
                CELL_SIZE * (cursor.row - 1) + GRID_Y - pad_diff,
                CURSOR_SIZE, CURSOR_SIZE)
    screen.blit(cursorSurface, rect)

def draw_pm_border(row, col, idx):
    third = CELL_SIZE // 3
    for n in pencil_marks[idx].border:
        digit = pm_border_font.render(str(n), True, "red")
        padding = 6
        basepos = (GRID_X + CELL_SIZE * col + padding, GRID_Y + CELL_SIZE * row + padding)
        pos = basepos[0] + ((n - 1) % 3) * third , basepos[1] + ((n - 1) // 3) * third
        screen.blit(digit, pos)

def draw_pm_center(row, col, idx):
    third = CELL_SIZE // 3
    nums = "".join(map(str, sorted(pencil_marks[idx].center)))
    digits = pm_center_font.render(nums, True, "dark green")
    pos = (GRID_X + CELL_SIZE * col + CELL_SIZE // 2 - digits.get_width() // 2,
           GRID_Y + CELL_SIZE * row + CELL_SIZE // 2 - digits.get_height() // 2.3)
    screen.blit(digits, pos)

done_txt = btn_font.render("Check", True, "black")
done_rect = Rect(screen.get_width() - 10 - 100, 10, 100, 30)
def draw_buttons():
    pg.draw.rect(screen, "gray", done_rect)
    pg.draw.rect(screen, "black", done_rect, width=2)
    screen.blit(done_txt,
                (done_rect.x + done_rect.w // 2 - done_txt.get_width() // 2,
                 done_rect.y + done_rect.h // 2 - done_txt.get_height() // 2))

def draw_numbers():
    for row in range(9):
        for col in range(9):
            idx = row * 9 + col
            cell = board[idx]
            if cell.value > 0:
                color = "black" if cell.given else "dark blue"
                digit = filled_num_font.render(str(cell.value), True, color)

                pos = (GRID_X + CELL_SIZE * col + CELL_SIZE // 2 - digit.get_width() // 2,
                    # Here we divide the height by 2.3 because there seems to be some extra
                    # padding at the bottom of the surface
                    GRID_Y + CELL_SIZE * row + CELL_SIZE // 2 - digit.get_height() // 2.3)

                screen.blit(digit, pos)
                continue
            if pencil_marks[idx].center:
                draw_pm_center(row, col, idx)
                continue
            if pencil_marks[idx].border:
                draw_pm_border(row, col, idx)

def draw_hud():
    screen.blit(title, (screen.get_width() // 2 - title.get_width() // 2, 70))

def check_board():
    for i,cell in enumerate(board):
        if cell.value == 0:
            print("You still have to finish the puzzle!")
            return
        if cell.value != sudoku.solution[i]:
            print("You made a mistake!")
            return
    print("You completed it!")

while running:
    for event in pg.event.get():
        if event.type == pg.QUIT:
            running = False
        if event.type == pg.MOUSEBUTTONDOWN:
            if event.button == 1:
                grid_max_x = GRID_X + CELL_SIZE * 9
                grid_max_y = GRID_Y + CELL_SIZE * 9
                mx = event.pos[0]
                my = event.pos[1]
                if (mx >= GRID_X
                    and mx <= grid_max_x
                    and my >= GRID_Y
                    and my <= grid_max_y):
                    cursor.row = int((my - GRID_Y) // CELL_SIZE + 1)
                    cursor.col = int((mx - GRID_X) // CELL_SIZE + 1)

                if done_rect.collidepoint(event.pos):
                    check_board()


        if event.type == pg.KEYDOWN:
            # keys = pg.key.get_pressed()
            if (event.key == pg.K_h or event.key == pg.K_a) and cursor.col > 1:
                cursor.col -= 1
            if (event.key == pg.K_k or event.key == pg.K_w) and cursor.row > 1:
                cursor.row -= 1
            if (event.key == pg.K_j or event.key == pg.K_s) and cursor.row < 9:
                cursor.row += 1
            if (event.key == pg.K_l or event.key == pg.K_d) and cursor.col < 9:
                cursor.col += 1

            num = pg.key.name(event.key)
            if num.isdigit() and num != '0':
                num = int(num)
                idx = (cursor.row - 1) * 9 + cursor.col - 1
                if not board[idx].given:
                    match input_method:
                        case InputMethod.FILL:
                            board[idx].value = num
                        case InputMethod.CENTER:
                            if num in pencil_marks[idx].center:
                                pencil_marks[idx].center.remove(num)
                            # Max 3 center pencil marks
                            elif len(pencil_marks[idx].center) < 3:
                                pencil_marks[idx].center.add(num)
                        case InputMethod.BORDER:
                            if num in pencil_marks[idx].border:
                                pencil_marks[idx].border.remove(num)
                            else:
                                pencil_marks[idx].border.add(num)

            if event.key == pg.K_BACKSPACE:
                idx = (cursor.row - 1) * 9 + cursor.col - 1
                if not board[idx].given:
                    board[idx].value = 0

            match event.key:
                case pg.K_f:
                    input_method = InputMethod.FILL
                case pg.K_c:
                    input_method = InputMethod.CENTER
                case pg.K_b:
                    input_method = InputMethod.BORDER

            ##############
            # Debug stuff
            ##############
            if event.key == pg.K_F1:
                print(sudoku)
            if event.key == pg.K_F2:
                idx = (cursor.row - 1) * 9 + cursor.col - 1
                if not board[idx].given:
                    pencil_marks[idx].border = {i for i in range(1, 10)}
            if event.key == pg.K_F3:
                for i,cell in enumerate(board):
                    cell.value = sudoku.solution[i]


    screen.fill('cornflower blue')

    draw_grid()
    draw_numbers()
    draw_cursor()
    draw_hud()
    draw_buttons()


    pg.display.flip()

    dt = clock.tick(60) / 1000

pg.quit()