import sys, os
import numpy as np
from PyQt5.QtWidgets import QApplication
from PyQt5.QtCore import Qt, QUrl, QTimer
from PyQt5.QtWebEngineWidgets import QWebEngineView, QWebEngineSettings
from PIL import Image
import urllib.parse

def getRandomFEN():
    fen_chars = list('1KQRBNPkqrbnp')
    pieces = np.random.choice(fen_chars, 64)
    fen = '/'.join([''.join(pieces[i*8:(i+1)*8]) for i in range(8)])
    return fen

class Screenshot(QWebEngineView):

    def __init__(self):
        super().__init__()
        self.current_index = 0
        self.total_screenshots = 0
        self.output_filenames = []
        self.theme_index = 0
        self.piece_index = 0
        self.background_index = 0
        self.themes = []
        self.pieceSets = []
        self.backgrounds = []
        self.out_folder = ''

    def capture(self, url, output_file, total_screenshots, output_filenames, themes, pieceSets, backgrounds, out_folder):
        self.output_file = output_file
        self.total_screenshots = total_screenshots
        self.output_filenames = output_filenames
        self.themes = themes
        self.pieceSets = pieceSets
        self.backgrounds = backgrounds
        self.out_folder = out_folder
        self.load(QUrl(url))
        self.loadFinished.connect(self.on_loaded)
        self.setAttribute(Qt.WA_DontShowOnScreen)
        self.page().settings().setAttribute(QWebEngineSettings.ShowScrollBars, False)
        self.show()

    def on_loaded(self):
        self.resize(1920, 1080)
        QTimer.singleShot(2000, self.take_screenshot)

    def take_screenshot(self):
        self.grab().save(self.output_file, b'PNG')
        QTimer.singleShot(1000, self.process_next)

    def process_next(self):
        if os.path.exists(self.output_file):
            im = Image.open(self.output_file).crop([488, 167, 1118, 797])
            im.save(self.output_file)
            self.current_index += 1

            # Update indices for theme, piece set, and background before generating the next screenshot
            if self.current_index >= self.total_screenshots:
                self.current_index = 0
                self.theme_index += 1  # Increment the theme first

                # After all themes are used, move to the next piece set
                if self.theme_index >= len(self.themes):
                    self.theme_index = 0  # Reset the theme index
                    self.piece_index += 1  # Increment the piece set

                    # After all piece sets are used, move to the next background
                    if self.piece_index >= len(self.pieceSets):
                        self.piece_index = 0  # Reset the piece set index
                        self.background_index += 1  # Increment the background

                        # If all backgrounds have been used, stop the app
                        if self.background_index >= len(self.backgrounds):
                            self.app.quit()  # Quit the app when all combinations are done
                            return
                        
            print(f"Theme: {self.themes[self.theme_index]}, Piece: {self.pieceSets[self.piece_index]}, Background: {self.backgrounds[self.background_index]}")

            fen = getRandomFEN().replace('/', '-')
            next_output_file = f"{self.out_folder}/{fen}.png"
            theme = self.themes[self.theme_index]
            piece = self.pieceSets[self.piece_index]
            background = self.backgrounds[self.background_index]
            # URL Encoding for theme
            encoded_theme = urllib.parse.quote(theme)
            url = f"http://en.lichess.org/editor/{fen.replace('-', '/')}?theme={encoded_theme}&pieceSet={piece}&bg={background}"
            print(theme,piece)
            print(f"#{self.current_index} : {fen}")
            print(f"\tSaved screenshot to '{next_output_file}'")
            self.capture(url, next_output_file, self.total_screenshots, self.output_filenames, self.themes, self.pieceSets, self.backgrounds, self.out_folder)
            print("\t...Success")
        else:
            print(f"File not found: {self.output_file}")
            QTimer.singleShot(2000, self.process_next)

app = QApplication(sys.argv)
s = Screenshot()
s.app = app

themes = ["brown","wood","wood2","wood3","wood4","maple","maple2","horsey","leather","blue","blue2","blue3","canvas","blue-marble","ic","green","marble","green-plastic","olive","grey","metal","newspaper","purple","purple-diag","pink"]
pieceSets = ["cburnett","merida","alpha","pirouetti","chessnut","chess7","reillycraig","companion","riohacha","kosal","leipzig","fantasy","spatial","celtic","california","caliente","pixel","maestro","fresca","cardinal","gioco","tatiana","staunty","cooke","monarchy","governor","dubrovny","icpieces","mpchess","kiwen-suwi","letter","shapes"]
backgrounds = ["light"]

# Number of random screenshots to generate per combination
N = 3

out_folder = 'train_images'
if not os.path.exists(out_folder):
    os.makedirs(out_folder, exist_ok=True)

fen = getRandomFEN().replace('/', '-')
s.output_filenames = [f"{out_folder}/{fen}.png" for i in range(N)]
print(f"#0 : {fen}")
print(f"\tSaved screenshot to '{s.output_filenames[0]}'")
# URL Encoding for theme
encoded_theme = urllib.parse.quote(themes[0])
url = f"http://en.lichess.org/editor/{fen.replace('-', '/')}?theme={encoded_theme}&piece={pieceSets[0]}&bg={backgrounds[0]}"
s.capture(url, s.output_filenames[0], N, s.output_filenames, themes, pieceSets, backgrounds, out_folder)
print("\t...Success")

sys.exit(app.exec_()) : File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png
File not found: train_images/QPBkPKrP-PNRbnPR1-nkbq1rRK-R1bPB1pk-11nNbnNp-nbrpnnnR-NpNPPppn-nnRNkknn.png