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 time

import threading
import sys

import os
from websudoku import *

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

@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)

@dataclass
class Button:
    id: str
    txt: pg.Surface
    rect: Rect

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

class Status(Enum):
    FETCHING = 1,
    SOLVING = 2,
    COMPLETED = 3


@dataclass
class GameState:
    start_time: float = 0
    status: Status = Status.FETCHING
    status_msg: str = ""

##################
# 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()

##################
# FONTS
##################
title_font = pg.font.Font(None, 64)
status_font = pg.font.Font(None, 44)
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()

board = []
pencil_marks = []
sudoku = None

start_time = [0]
game_state = GameState()

data_lock = threading.Lock()

def data_callback(sudoku_data):
    # Handle the data here, e.g., store it in a global variable
    for i in range(81):
        if sudoku_data.editmask[i] == 0:
            board.append(Cell(sudoku_data.solution[i], True))
        else:
            board.append(Cell())

    for i in range(81):
        pm = PencilMark() if not board[i].given else None
        pencil_marks.append(pm)
    global sudoku
    with data_lock:
        game_state.start_time = time.time()
        game_state.status = Status.SOLVING
        game_state.status_msg = ""
        sudoku = sudoku_data


def http_request_thread():
    if len(sys.argv) > 1:
        difficulty = Difficulty[sys.argv[1].upper()]
    else:
        difficulty = Difficulty.MEDIUM
    data = get_sudoku_puzzle(difficulty)
    data_callback(data)

http_thread = threading.Thread(target=http_request_thread)
http_thread.start()


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 + GRID_X - pad_diff,
                CELL_SIZE * cursor.row + 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)

def new_button(id, label, color, x, y, w, h):
    return Button(id,
                  btn_font.render(label, True, color),
                  Rect(x, y, w, h))
buttons = [
    new_button("check", "Check", "black", screen.get_width() - 10 - 100, 10, 100, 30),
    new_button("load", "Load new", "black", screen.get_width() - 10 - 150, 50, 150, 30),
]

def draw_buttons():
    for btn in buttons:
        pg.draw.rect(screen, "gray", btn.rect)
        pg.draw.rect(screen, "black", btn.rect, width=2)
        screen.blit(btn.txt,
                    (btn.rect.x + btn.rect.w // 2 - btn.txt.get_width() // 2,
                    btn.rect.y + btn.rect.h // 2 - btn.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))
    status_txt = status_font.render(game_state.status_msg, True, "black")
    screen.blit(status_txt, (screen.get_width() // 2 - status_txt.get_width() // 2,
                         screen.get_height() - status_txt.get_height() - 70))

def check_board():
    for i,cell in enumerate(board):
        if cell.value != 0 and cell.value != sudoku.solution[i]:
            game_state.status_msg = "You made a mistake!"
            return
    for i,cell in enumerate(board):
        if cell.value == 0:
            game_state.status_msg = "You still have to finish the puzzle!"
            return
    completion_time = time.time() - game_state.start_time
    mins = int(completion_time // 60)
    secs = int(completion_time % 60)
    game_state.status_msg = f"You finished it in {mins} min {secs} sec!"
    game_state.status = Status.COMPLETED

def main():
    running = True
    input_method = InputMethod.FILL
    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)
                        cursor.col = int((mx - GRID_X) // CELL_SIZE)

                    for btn in buttons:
                        if btn.rect.collidepoint(event.pos):
                            match btn.id:
                                case "check":
                                    if game_state.status == Status.SOLVING:
                                        check_board()
                                case "load":
                                    pass


            if event.type == pg.KEYDOWN:
                if (event.key == pg.K_h or event.key == pg.K_a):
                    cursor.col = (cursor.col - 1) % 9
                if (event.key == pg.K_k or event.key == pg.K_w):
                    cursor.row = (cursor.row - 1) % 9
                if (event.key == pg.K_j or event.key == pg.K_s):
                    cursor.row = (cursor.row + 1) % 9
                if (event.key == pg.K_l or event.key == pg.K_d):
                    cursor.col = (cursor.col + 1) % 9

                num = pg.key.name(event.key)
                if num.isdigit() and num != '0':
                    num = int(num)
                    idx = cursor.row * 9 + cursor.col
                    if board and 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 * 9 + cursor.col
                    if not board[idx].given:
                        if board[idx].value > 0:
                            board[idx].value = 0
                        elif pencil_marks[idx].center:
                            pencil_marks[idx].center.pop()

                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 * 9 + cursor.col
                    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()
        with data_lock:
            if sudoku is not None:
                draw_numbers()
        draw_cursor()
        draw_hud()
        draw_buttons()


        pg.display.flip()

        dt = clock.tick(60) / 1000

    pg.quit()

if __name__ == "__main__":
    main()