from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
import json
import pygame
from pygame import Vector2
import re


# Colors
BACKGROUND_COLOR = "#d7bfa0"
PAGE_BORDER_COLOR = "#c19d62"
WORD_BORDER_COLOR = "#fffcf2"
WORD_SHADOW_COLOR = "#000000"

# Drawing constants
SCREEN_WIDTH = 2560
SCREEN_HEIGHT = 1440
PAGE_MARGIN = 20

PAGE_BORDER_RADIUS = 40
PAGE_BORDER_WIDTH = 20

TAB_WIDTH = 80
TAB_HEIGHT = 80

WORD_WIDTH = 228
WORD_HEIGHT = 55
WORD_MARGIN = 16
WORD_BORDER_RADIUS = 30
WORD_BORDER_WIDTH = 6
WORD_FONT_SIZE = 30
WORD_SHADOW_SIZE = 5

GRID_WIDTH = 5
GRID_HEIGHT = 18

WORD_PAGE_WIDTH = PAGE_BORDER_WIDTH*2+(WORD_WIDTH+WORD_MARGIN)*GRID_WIDTH+WORD_MARGIN
WORD_PAGE_HEIGHT = SCREEN_HEIGHT-PAGE_MARGIN*2

TEXT_PAGE_POS = Vector2(WORD_PAGE_WIDTH+PAGE_MARGIN*2,TAB_HEIGHT+PAGE_MARGIN*2)
TEXT_PAGE_WIDTH = SCREEN_WIDTH-(WORD_PAGE_WIDTH+PAGE_MARGIN*3)
TEXT_PAGE_HEIGHT = SCREEN_HEIGHT-(TAB_HEIGHT+PAGE_MARGIN*3)

TEXT_MARGIN = 20
TEXT_TOP_MARGIN = 35
TEXT_WIDTH = TEXT_PAGE_WIDTH - TEXT_MARGIN*2
TEXT_HEIGHT = TEXT_PAGE_HEIGHT - TEXT_MARGIN*2

WORD_SLOT_SPACING = " "*12
WORD_SLOT_NUDGING = 22

class WordColor(Enum):
    Red = 0     # Names
    Green = 1   # Nouns
    Lime = 2    # Numbers and such
    Blue = 3    # Verbs
    Grey = 4    # Other

    def color(self):
        return [
            "#b03334",
            "#6a5f31",
            "#4b7a28",
            "#297cb7",
            "#767978"
        ][self.value]

@dataclass
class Page():
    visible: bool
    color: str
    text: None|str
    position: Vector2
    size: Vector2

@dataclass
class Word():
    word: str
    color: WordColor
    page: None|Page
    position: Vector2
    slot: None|WordSlot = None

    def copy(self):
        return Word(self.word,self.color,self.page,self.position,self.slot)

@dataclass
class WordSlot():
    color: None|WordColor
    page: Page
    position: Vector2
    word: None|Word = None

def draw_page(page: Page, screen: pygame.surface.Surface, font: pygame.font.Font):
    pygame.draw.rect(
        screen,
        page.color,
        pygame.Rect(
            page.position.x + PAGE_BORDER_WIDTH,
            page.position.y + PAGE_BORDER_WIDTH,
            page.size.x - PAGE_BORDER_WIDTH*2,
            page.size.y - PAGE_BORDER_WIDTH*2
        )
    )

    pygame.draw.rect(
        screen,
        PAGE_BORDER_COLOR,
        pygame.Rect(page.position.x,page.position.y,page.size.x,page.size.y),
        PAGE_BORDER_WIDTH, PAGE_BORDER_RADIUS
    )

    if page.text is not None:
        for i, t in enumerate(page.text):
            text = font.render(t,False,"#000000")
            # x,y = text.get_size()
            screen.blit(text,(
                page.position.x + PAGE_BORDER_WIDTH + TEXT_MARGIN,
                page.position.y + PAGE_BORDER_WIDTH + TEXT_MARGIN + TEXT_TOP_MARGIN + i*(WORD_FONT_SIZE+TEXT_TOP_MARGIN)
            ))

def draw_word_slot(word_slot: WordSlot, screen: pygame.surface.Surface):
    pos = word_slot.page.position+word_slot.position

    if word_slot.color is None:
        color = "#00000000"
    else:
        color = word_slot.color.color()

    pygame.draw.rect(
        screen,
        color,
        pygame.Rect(
            pos.x + WORD_BORDER_WIDTH,
            pos.y + WORD_BORDER_WIDTH,
            WORD_WIDTH - WORD_BORDER_WIDTH*2,
            WORD_HEIGHT - WORD_BORDER_WIDTH*2
        ),
        border_radius=WORD_BORDER_RADIUS
    )

def draw_word(word: Word, screen: pygame.surface.Surface, font: pygame.font.Font):
    if word.page is not None:
        pos = word.page.position+word.position
    else:
        pos = word.position

    pygame.draw.rect(
        screen,
        WORD_SHADOW_COLOR,
        pygame.Rect(
            pos.x,
            pos.y + WORD_SHADOW_SIZE,
            WORD_WIDTH,
            WORD_HEIGHT
        ),
        border_radius=WORD_BORDER_RADIUS
    )

    pygame.draw.rect(
        screen,
        word.color.color(),
        pygame.Rect(
            pos.x + WORD_BORDER_WIDTH//2,
            pos.y + WORD_BORDER_WIDTH//2,
            WORD_WIDTH - WORD_BORDER_WIDTH,
            WORD_HEIGHT - WORD_BORDER_WIDTH
        ),
        border_radius=WORD_BORDER_RADIUS
    )

    pygame.draw.rect(
        screen,
        WORD_BORDER_COLOR,
        pygame.Rect(
            pos.x,
            pos.y,
            WORD_WIDTH,WORD_HEIGHT
        ),
        WORD_BORDER_WIDTH, WORD_BORDER_RADIUS
    )

    text = font.render(word.word,False,"#ffffff")
    x,y = text.get_size()
    screen.blit(text,(
        pos.x+(WORD_WIDTH-x)//2,
        pos.y+(WORD_HEIGHT-y)//2+1
    ))

def main(data: dict):
    pygame.init()
    pygame.font.init()
    font = pygame.font.SysFont("Comic Code",WORD_FONT_SIZE,True)
    screen = pygame.display.set_mode((1280, 720),pygame.RESIZABLE)
    drawer = pygame.surface.Surface((SCREEN_WIDTH, SCREEN_HEIGHT))
    clock = pygame.time.Clock()
    running = True

    word_page = Page(True,BACKGROUND_COLOR,None,Vector2(PAGE_MARGIN,PAGE_MARGIN),Vector2(WORD_PAGE_WIDTH,WORD_PAGE_HEIGHT))
    pages = [word_page] + [
        Page(
            False,
            BACKGROUND_COLOR,
            [p],
            TEXT_PAGE_POS,
            Vector2(TEXT_PAGE_WIDTH,TEXT_PAGE_HEIGHT)
        ) for p in data["pages"]
    ]
    if len(pages) > 1:
        pages[1].visible = True
        visible_page = 1
    else:
        visible_page = 0

    held_word = None

    word_slots: list[WordSlot] = []
    for y in range(GRID_HEIGHT):
        for x in range(GRID_WIDTH):
            word_slots.append(WordSlot(
                None,
                word_page,
                Vector2(
                    PAGE_BORDER_WIDTH+WORD_MARGIN+(WORD_WIDTH+WORD_MARGIN)*x,
                    PAGE_BORDER_WIDTH+WORD_MARGIN+(WORD_HEIGHT+WORD_MARGIN)*y
                )
            ))

    def add_to_text(text,word,space=" "):
        width, _ = font.size(text[-1]+space+word)
        if width > TEXT_WIDTH:
            return text+[word]
        else:
            text[-1] += space+word
            return text

    for p in pages[1:]:
        new_text = [""]
        text_words = p.text[0].split(" ")
        for w in text_words:
            res = re.findall(r"\[(.*?)\/(\d)\](.)?",w)
            if res == []:
                new_text = add_to_text(new_text,w)
            else:
                if new_text[0] != "":
                    new_text = add_to_text(new_text,"")

                for slot in res:
                    new_text = add_to_text(new_text,WORD_SLOT_SPACING,"")
                    height = WORD_FONT_SIZE*len(new_text)
                    width, _ = font.size(new_text[-1])
                    new_text = add_to_text(new_text,slot[2],"")
                    pos = Vector2(
                        TEXT_MARGIN + width - WORD_WIDTH,
                        TEXT_MARGIN + height - WORD_HEIGHT//2 + TEXT_TOP_MARGIN*len(new_text)
                    )
                    # if new_text[-1] != WORD_SLOT_SPACING+slot[2]:
                    pos.x += WORD_SLOT_NUDGING
                    word_slots.append(WordSlot(WordColor(int(slot[1])),p,pos))

        p.text = new_text

    words: list[Word] = []
    for i,w in enumerate(data["words"]):
        slot = word_slots[i]
        word = Word(w[0],WordColor(w[1]),word_page,Vector2(slot.position.x,slot.position.y))
        slot.word = word
        word.slot = slot
        words.append(word)

    def focused():
        x, y = pygame.mouse.get_pos()
        ratio_x = (screen.get_width() / drawer.get_width())
        ratio_y = (screen.get_height() / drawer.get_height())
        scaled_xy = (x/ratio_x,y/ratio_y)
        for w in words:
            if w.page is None:
                return w
            pos = w.page.position+w.position
            rect = pygame.Rect(
                pos.x - WORD_MARGIN,
                pos.y - WORD_MARGIN,
                WORD_WIDTH + WORD_MARGIN*2,
                WORD_HEIGHT + WORD_MARGIN*2
            )
            if rect.collidepoint(scaled_xy):
                return w

        return None

    def click():
        nonlocal held_word
        slot = None
        x, y = pygame.mouse.get_pos()
        ratio_x = (screen.get_width() / drawer.get_width())
        ratio_y = (screen.get_height() / drawer.get_height())
        scaled_xy = (x/ratio_x,y/ratio_y)
        for w in word_slots:
            pos = w.page.position+w.position
            rect = pygame.Rect(
                pos.x - WORD_MARGIN,
                pos.y - WORD_MARGIN,
                WORD_WIDTH + WORD_MARGIN*2,
                WORD_HEIGHT + WORD_MARGIN*2
            )
            if rect.collidepoint(scaled_xy):
                slot = w
                break

        if slot is None:
            return

        if held_word is None and slot.word is not None:
            if slot.color is None:
                held_word = slot.word
                held_word.page = None
            else:
                words.remove(slot.word)
            slot.word = None
        elif held_word is not None and slot.word is None:
            if slot.color is not None and slot.color != held_word.color:
                return

            slot.word = held_word
            if slot.color is not None:
                new_word = slot.word.copy()
                new_word.page = slot.word.slot.page
                new_word.position = slot.word.slot.position
                new_word.slot = slot.word.slot
                new_word.slot.word = new_word
                words.append(new_word)
            slot.word.position = slot.position
            slot.word.page = slot.page
            slot.word.slot = slot
            held_word = None


    while running:
        for event in pygame.event.get():
            if event.type == pygame.MOUSEBUTTONUP:
                click()
            elif event.type == pygame.QUIT:
                running = False

        # fill the screen with a color to wipe away anything from last frame
        drawer.fill("#4d2905")

        # Held word
        if held_word is not None:
            x, y = pygame.mouse.get_pos()
            ratio_x = (screen.get_width() / drawer.get_width())
            ratio_y = (screen.get_height() / drawer.get_height())
            scaled_xy = (x/ratio_x,y/ratio_y)
            held_word.position = Vector2(scaled_xy) - Vector2(WORD_WIDTH//2,WORD_HEIGHT//2)

        # RENDER YOUR GAME HERE
        for p in pages:
            if not p.visible:
                continue

            draw_page(p, drawer, font)

        # drawing word slots
        for w in word_slots:
            if not w.page.visible or w.color is None:
                continue

            draw_word_slot(w, drawer)

        # drawing words
        for w in words:
            if w.page is not None and not w.page.visible and w != held_word:
                continue

            draw_word(w, drawer, font)

        if held_word is not None:
            draw_word(held_word, drawer, font)

        # Draw everything to the screen
        frame = pygame.transform.scale(drawer,screen.get_size())
        screen.blit(frame,frame.get_rect())

        if focused() is not None:
            pygame.mouse.set_cursor(*pygame.cursors.tri_left)
        else:
            pygame.mouse.set_cursor(*pygame.cursors.arrow)

        # flip() the display to put your work on screen
        pygame.display.flip()

        clock.tick(60)  # limits FPS to 60

    pygame.quit()

if __name__ == "__main__":
    with open("data.json","r") as file_pointer:
        data = json.load(file_pointer)

    main(data)