from __future__ import annotations

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


# Colors
PAGE_BACKGROUND_COLOR = "#d7bfa0"
NOT_SELECTED_COLOR = "#7f7260"
PAGE_BORDER_COLOR = "#c19d62"
BORDER_NOT_SELECTED_COLOR = "#76603c"
WORD_BORDER_COLOR = "#fffcf2"
WORD_SHADOW_COLOR = "#000000"

# Drawing constants
SCREEN_WIDTH = 5120
SCREEN_HEIGHT = 2880

PAGE_MARGIN = 30
PAGE_BORDER_RADIUS = 80
PAGE_BORDER_WIDTH = 20

WORD_WIDTH = 432
WORD_HEIGHT = 90
WORD_MARGIN = 20
WORD_BORDER_RADIUS = 60
WORD_BORDER_WIDTH = 8
WORD_FONT_SIZE = 44
WORD_SHADOW_SIZE = 8

GRID_WIDTH = 6
GRID_HEIGHT = 25

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

TAB_WIDTH = 120
TAB_HEIGHT = 160
TAB_MARGIN = 30
TAB_POS = Vector2(WORD_PAGE_WIDTH+PAGE_MARGIN*2,PAGE_MARGIN)

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 = 40
TEXT_TOP_MARGIN = 68
TEXT_WIDTH = TEXT_PAGE_WIDTH - PAGE_BORDER_WIDTH*2 - TEXT_MARGIN*2
TEXT_HEIGHT = TEXT_PAGE_HEIGHT - PAGE_BORDER_WIDTH*2 - TEXT_MARGIN*2

WORD_SLOT_SPACING = " "*16
WORD_SLOT_NUDGING = 25
WORD_SLOT_NUDGE_DOWN = -5
SLANT = 80

INDICATOR_SIZE = 80
INDICATOR_BORDER = 20
INDICATOR_MARGIN = 60
INDICATOR_POS = Vector2(TEXT_PAGE_WIDTH-(INDICATOR_SIZE+INDICATOR_MARGIN),TEXT_PAGE_HEIGHT-(INDICATOR_SIZE+INDICATOR_MARGIN))

class WordColor(Enum):
    Red = 0     # Names
    Green = 1   # Nouns
    Lime = 2    # Numbers and such
    Blue = 3    # Verbs
    Grey = 4    # Other
    Purple = 5  # Places
    Yellow = 6  # Titles
    Orange = 7  # Classes

    def color(self):
        return [
            "#b03334",
            "#6a5f31",
            "#62a032",
            "#297cb7",
            "#767978",
            "#623a75",
            "#bda627",
            "#b76c2d",
        ][self.value]

class SolvedState(Enum):
    NotFinished = 0
    Wrong = 1
    AlmostCorrect = 2
    Correct = 3

@dataclass
class Page():
    visible: bool
    text: None|str
    position: Vector2
    size: Vector2
    slots: list[tuple[list[str],WordSlot]]

    def solved(self):
        if any([s[1].word is None for s in self.slots]):
            return SolvedState.NotFinished
        elif all([s[1].word.word in s[0] for s in self.slots]):
            return SolvedState.Correct
        elif sum([s[1].word.word in s[0] for s in self.slots]) >= len(self.slots)-2:
            return SolvedState.AlmostCorrect
        else:
            return SolvedState.Wrong

@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():
    colors: None|list[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_BACKGROUND_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")
            pos = Vector2(
                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)
            )
            if i == 0:
                width,_ = text.get_size()
                text = pygame.transform.scale2x(text)
                pos.y = page.position.y + PAGE_BORDER_WIDTH + TEXT_MARGIN
                pos.x = page.position.x + PAGE_BORDER_WIDTH + TEXT_MARGIN + (TEXT_WIDTH-width*2)//2
            screen.blit(text,pos)

    if page.slots != []:
        pygame.draw.circle(
            screen,
            "#000000",
            page.position + INDICATOR_POS,
            INDICATOR_SIZE
        )
        status = page.solved()
        if status == SolvedState.Correct:
            color = WordColor.Lime.color()
        elif status == SolvedState.AlmostCorrect:
            color = WordColor.Yellow.color()
        elif status == SolvedState.Wrong:
            color = WordColor.Red.color()
        else:
            color = "#000000"

        pygame.draw.circle(
            screen,
            color,
            page.position + INDICATOR_POS,
            INDICATOR_SIZE-INDICATOR_BORDER
        )

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

    if word_slot.colors is None:
        color = "#00000000"
    else:
        color = word_slot.colors[0].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
    )

    if len(word_slot.colors) > 1:
        pygame.draw.rect(
            screen,
            word_slot.colors[1].color(),
            pygame.Rect(
                pos.x + WORD_BORDER_WIDTH + WORD_WIDTH//4,
                pos.y + WORD_BORDER_WIDTH,
                (WORD_WIDTH*3)//4 - WORD_BORDER_WIDTH*2,
                WORD_HEIGHT - WORD_BORDER_WIDTH*2
            ),
            border_radius=WORD_BORDER_RADIUS
        )

        pygame.draw.polygon(
            screen,
            color,
            [
                (
                    pos.x + WORD_BORDER_WIDTH + WORD_WIDTH//4,
                    pos.y + WORD_BORDER_WIDTH
                ),
                (
                    pos.x + WORD_BORDER_WIDTH + WORD_WIDTH//2 + SLANT//2,
                    pos.y + WORD_BORDER_WIDTH
                ),
                (
                    pos.x + WORD_BORDER_WIDTH + WORD_WIDTH//2 - SLANT//2,
                    pos.y + WORD_BORDER_WIDTH + WORD_HEIGHT - WORD_BORDER_WIDTH*2 - 1
                ),
                (
                    pos.x + WORD_BORDER_WIDTH + WORD_WIDTH//4,
                    pos.y + WORD_BORDER_WIDTH + WORD_HEIGHT - WORD_BORDER_WIDTH*2 - 1
                ),
            ]
        )

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 draw_tab(n: int, screen: pygame.surface.Surface, font: pygame.font.Font, selected: bool):
    pos = TAB_POS + Vector2((TAB_WIDTH+TAB_MARGIN)*n,0)

    pygame.draw.rect(
        screen,
        PAGE_BACKGROUND_COLOR if selected else NOT_SELECTED_COLOR,
        pygame.Rect(
            pos.x+PAGE_BORDER_WIDTH//2,
            pos.y+PAGE_BORDER_WIDTH//2,
            TAB_WIDTH-PAGE_BORDER_WIDTH,
            TAB_HEIGHT-PAGE_BORDER_WIDTH
        )
    )

    pygame.draw.rect(
        screen,
        PAGE_BORDER_COLOR if selected else BORDER_NOT_SELECTED_COLOR,
        pygame.Rect(
            pos.x,
            pos.y,
            TAB_WIDTH,
            TAB_HEIGHT
        ),
        PAGE_BORDER_WIDTH//2, PAGE_BORDER_RADIUS//2
    )

    text = font.render(str(n+1),False,"#000000")
    x,y = text.get_size()
    screen.blit(text,(
        pos.x+(TAB_WIDTH-x)//2,
        pos.y+(TAB_HEIGHT-y)//2+1
    ))


def main(data: dict):
    pygame.init()
    pygame.font.init()
    font = pygame.font.SysFont("Comic Code",WORD_FONT_SIZE,True)
    print(font.size(WORD_SLOT_SPACING))
    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,None,Vector2(PAGE_MARGIN,PAGE_MARGIN),Vector2(WORD_PAGE_WIDTH,WORD_PAGE_HEIGHT),[])
    pages = [word_page] + [
        Page(
            False,
            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:
            if text[-1] == "":
                space = ""
            text[-1] += space+word
            return text

    for p in pages[1:]:
        new_text = []
        for t in p.text:
            new_text.append("")
            text_words = t.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,slot[0]+WORD_SLOT_SPACING+slot[3],"")
                        height = WORD_FONT_SIZE*len(new_text)
                        width, _ = font.size(new_text[-1])
                        extra_width, _ = font.size(slot[3])
                        pos = Vector2(
                            TEXT_MARGIN + width - extra_width - WORD_WIDTH,
                            TEXT_MARGIN + height - WORD_HEIGHT//2 + TEXT_TOP_MARGIN*len(new_text) + WORD_SLOT_NUDGE_DOWN
                        )
                        # if new_text[-1] != WORD_SLOT_SPACING+slot[2]:
                        pos.x += WORD_SLOT_NUDGING
                        new_slot = WordSlot([WordColor(int(i)) for i in slot[2].split("+")],p,pos)
                        word_slots.append(new_slot)
                        p.slots.append((slot[1].replace("%"," ").split("+"),new_slot))

        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)

    all_words = [w.word for w in words]
    for p in pages:
        not_present = [word for s in p.slots for word in s[0] if word not in all_words]
        if not_present:
            print(not_present)

    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
            if not w.page.visible:
                continue

            pos = w.page.position+w.position
            rect = pygame.Rect(
                pos.x - WORD_MARGIN//2,
                pos.y - WORD_MARGIN//2,
                WORD_WIDTH + WORD_MARGIN,
                WORD_HEIGHT + WORD_MARGIN
            )
            if rect.collidepoint(scaled_xy):
                return w

        for i in range(len(pages)-1):
            pos = TAB_POS + Vector2((TAB_WIDTH+TAB_MARGIN)*i,0)
            rect = pygame.Rect(
                pos.x - TAB_MARGIN//2,
                pos.y - TAB_MARGIN//2,
                TAB_WIDTH + TAB_MARGIN,
                TAB_HEIGHT + TAB_MARGIN
            )
            if rect.collidepoint(scaled_xy):
                return i+1

        return None

    def click_slot():
        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:
            if not w.page.visible:
                continue

            pos = w.page.position+w.position
            rect = pygame.Rect(
                pos.x - WORD_MARGIN//2,
                pos.y - WORD_MARGIN//2,
                WORD_WIDTH + WORD_MARGIN,
                WORD_HEIGHT + WORD_MARGIN
            )
            if rect.collidepoint(scaled_xy):
                slot = w
                break

        if slot is None:
            if held_word is not None and held_word.slot.colors is not None:
                words.remove(held_word)
                held_word = None
            return

        if held_word is None and slot.word is not None:
            held_word = slot.word
            held_word.page = None
            slot.word = None
        elif held_word is not None and slot.word is None:
            if slot.colors is not None and held_word.color not in slot.colors:
                return

            slot.word = held_word
            if slot.colors is not None and held_word.slot.colors is 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)

            if slot.word.slot.colors is None or slot.colors is not None:
                slot.word.position = slot.position
                slot.word.page = slot.page
                slot.word.slot = slot
            else:
                words.remove(slot.word)
                slot.word = None

            held_word = None


    while running:
        for event in pygame.event.get():
            if event.type == pygame.MOUSEBUTTONDOWN:
                click_slot()
                i = focused()
                if isinstance(i,int):
                    pages[visible_page].visible = False
                    visible_page = i
                    pages[visible_page].visible = True
            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.colors 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)

        # drawing tabs
        for i in range(len(pages)-1):
            draw_tab(i,drawer,font,visible_page==i+1)

        # drawing held word
        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)