import argparse
import chess
import chess.pgn
from board_detector import find_board
from board_detector import find_pieces
from board_detector import init_show_cv
from board_detector import display_img
import color_analyzer
import move_translator
import cv2
import os

import time
from picamera2 import Picamera2, Preview

class ChessGame:
    def __init__(self, difficulty, show_cv, show_cam, calibrate, test_img = None):
        self.board = chess.Board()
        self.difficulty = difficulty
        self.show_cv = show_cv
        self.show_cam = show_cam
        self.test_img = test_img
        self.calibrate = calibrate
        self.img_idx = 0
        self.left_cut = 30
        self.right_cut = 290
        self.top_cut = 0
        self.bottom_cut = 290
        self.picam2 = Picamera2()

    def calibrate_cam():
        pass

    def start_game(self):

        if (self.calibrate):
            self.calibrate_cam()

        print(f"Starting chess game (difficulty: {self.difficulty})")

        # TODO - call initialize board in board_detector, initialize colors for color analysis,
        # then loop until checkmate. also handle illegal moves (writing to screen if we end up doing that or just LEDs)

        init_show_cv(self.show_cv)


        preview_config = self.picam2.create_preview_configuration(main={"size": (2464, 2464)})
        self.picam2.configure(preview_config)
        if (self.show_cam):
            self.picam2.start_preview(Preview.QTGL)
            self.picam2.start()
            input("Press enter to continue.")
        else:
            self.picam2.start()

        # self.get_img()

        if (self.test_img):
            img_path = os.path.join('test_images', self.test_img)
            orig_img = cv2.imread(img_path)
        else:
            # orig_img = cv2.imread('test_images/board1.jpg') # TODO - CHANGE TO MAKE IT RECEIVE INPUT FROM CAMERA
            orig_img = self.take_pic()
        h,w,c = orig_img.shape
        # w_cut = 30
        # h_cut = 100
        cropped_img = orig_img[self.top_cut:h-self.bottom_cut, self.left_cut:w-self.right_cut]
        # cropped_img = orig_img
        img = cv2.resize(cropped_img, (512, 512))
        # img = orig_img

        if (self.show_cv):
            display_img([orig_img, img])
            # display_img([img])

        warped_img, sorted_warped_points = find_board(img)

        if (warped_img is None):
            return
        
        color_grid = find_pieces(warped_img, sorted_warped_points)

        while(1): # game loop
            self.player_turn()
            # handle cheating
            if self.board.is_checkmate():
                break

            quit() # TODO - REMOVE

            self.ai_turn()
            if self.board.is_checkmate():
                break

        # game is over
            
    def player_turn(self):
        # TODO - wait for user button, then check for valid move. loop until a valid move has been made
        pass

    def ai_turn(self):
        # TODO
        # 1. use detect_pieces in color_analyzer
        # 2. update internal representation of board
        # 3. give to python-chess and get best move
        # 4. call move_translator with argument of best move
        pass

    def take_pic(self):
        self.img_idx += 1
        time.sleep(2)
        img_txt = 'board' + str(self.img_idx) + '.jpg'
        # print(img_txt)
        img_path = os.path.join('game_images', img_txt)
        metadata = self.picam2.capture_file(img_path)
        return cv2.imread(img_path)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="AI Chess Robot with Computer Vision")
    parser.add_argument("--difficulty", choices=["easy", "medium", "hard"], 
                        default="medium", help="Chess AI difficulty (how far it looks ahead)")
    parser.add_argument("--show_cv", action="store_true", help="Show opencv images as processing occurs during game")
    parser.add_argument("--show_cam", action="store_true", help="Show persistent camera view")
    parser.add_argument("--calibrate", action="store_true", help="Loop showing camera with cutoff lines first, then start game")
    parser.add_argument("--test_img", help="If specified, will use said image in test_images folder rather than camera input")
    args = parser.parse_args()

    game = ChessGame(args.difficulty, args.show_cv, args.show_cam, args.calibrate, args.test_img)
    game.start_game()