---

import os
import cv2
import numpy as np
import json
import time
import keyboard
from mss import mss

# --- НАСТРОЙКИ И ПУТИ ---
TEMPLATES_DIR = "extracted_pieces"
COLORS_FILE = "colors.json"
CONFIG_FILE = "config.json"

ALL_PIECES = ['K', 'Q', 'R', 'B', 'N', 'P', 'k', 'q', 'r', 'b', 'n', 'p']

PIECES_CONFIG = {
    "K": ("white_K.png", "white"), "Q": ("white_Q.png", "white"),
    "R": ("white_R.png", "white"), "B": ("white_B.png", "white"),
    "N": ("white_N.png", "white"), "P": ("white_P.png", "white"),
    "k": ("black_k.png", "black"), "q": ("black_q.png", "black"),
    "r": ("black_r.png", "black"), "b": ("black_b.png", "black"),
    "n": ("black_n.png", "black"), "p": ("black_p.png", "black")
}

is_active = False
board_roi = None
need_selection = False

def clear_console():
    os.system('cls' if os.name == 'nt' else 'clear')

def toggle_scanner():
    global is_active
    is_active = not is_active
    print(f"\n[СТАТУС] Сканирование: {'ВКЛ' if is_active else 'ПАУЗА'}")

def request_selection():
    global need_selection
    need_selection = True

def extract_dominant_piece_color(cell_bgr):
    """Извлекает медианный цвет центральной части фигуры."""
    h, w, _ = cell_bgr.shape
    center = cell_bgr[int(h*0.25):int(h*0.75), int(w*0.25):int(w*0.75)]
    median_color = np.median(center, axis=(0, 1))
    return median_color.tolist()

def generate_colors_database():
    """Автоматически генерирует файл colors.json на основе шаблонов."""
    print("[ИНФО] Файл colors.json не найден. Создаём базу цветов по шаблонам...")
    colors_db = {}
    for sym, (filename, _) in PIECES_CONFIG.items():
        img_path = os.path.join(TEMPLATES_DIR, filename)
        if os.path.exists(img_path):
            img = cv2.imread(img_path)
            colors_db[sym] = extract_dominant_piece_color(img)
        else:
            # Дефолтные значения, если файла шаблона не оказалось
            colors_db[sym] = [240, 240, 240] if sym.isupper() else [30, 30, 30]
            
    with open(COLORS_FILE, 'w', encoding='utf-8') as f:
        json.dump(colors_db, f, indent=4)
    print("[УСПЕХ] Файл colors.json успешно создан!\n")

def select_board_region(sct_instance):
    """Выбор области доски мышкой."""
    global board_roi, is_active
    was_active = is_active
    is_active = False

    print("\n[ВЫБОР] Выделите доску мышкой и нажмите ENTER (ESC — отмена)")
    monitor = sct_instance.monitors[1]
    frame = np.array(sct_instance.grab(monitor))[:, :, :3]

    win_name = "Выделите доску (ENTER - ок, ESC - отмена)"
    cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)
    roi = cv2.selectROI(win_name, frame, fromCenter=False, showCrosshair=True)
    cv2.waitKey(1)
    cv2.destroyWindow(win_name)
    cv2.waitKey(1)

    x, y, w, h = roi
    if w > 50 and h > 50:
        board_roi = {"top": int(y), "left": int(x), "width": int(w), "height": int(h)}
        with open(CONFIG_FILE, 'w', encoding='utf-8') as f:
            json.dump(board_roi, f, indent=4)
        print(f"[УСПЕХ] Область доски успешно сохранена в config.json!")
    else:
        print("[ОТМЕНА] Выделение отменено.")
    
    time.sleep(0.5)
    is_active = was_active

def print_visual_board(board_matrix, fen):
    clear_console()
    status = "[ВКЛ]" if is_active else "[ПАУЗА]"
    print("=" * 60)
    print(f"   {status} | F7: Старт/Пауза | F8: Выделить доску мышкой")
    print("=" * 60)
    print("    a   b   c   d   e   f   g   h")
    print("  +---------------------------------+")
    for idx, row in enumerate(board_matrix):
        rank = 8 - idx
        row_str = " ".join([f"[{p}]" if p != '.' else " . " for p in row])
        print(f"{rank} | {row_str} | {rank}")
    print("  +---------------------------------+")
    print("    a   b   c   d   e   f   g   h")
    print("=" * 60)
    print("FEN:", fen)

def scan_frame(cell_board, color_database, templates_edges):
    """Главная функция сканирования доски (Canny + Color Penalty)."""
    board = cv2.resize(cell_board, (400, 400))
    board_matrix = []
    
    for row in range(8):
        row_pieces = []
        for col in range(8):
            cell = board[row*50:(row+1)*50, col*50:(col+1)*50]
            gray_cell = cv2.cvtColor(cell, cv2.COLOR_BGR2GRAY)
            
            # Контуры Canny только по центру (7:43)
            cell_edge = cv2.Canny(gray_cell[7:43, 7:43], 50, 150)

            # 1. Проверка на пустую клетку
            if np.count_nonzero(cell_edge) < 25:
                row_pieces.append('.')
                continue

            # 2. Вычисление цвета и поиск максимального совпадения
            cell_color = np.array(extract_dominant_piece_color(cell))
            best_match, max_score = '.', -999.0

            for piece_symbol in ALL_PIECES:
                if piece_symbol not in templates_edges:
                    continue
                
                # Сравнение контуров
                res = cv2.matchTemplate(cell_edge, templates_edges[piece_symbol], cv2.TM_CCOEFF_NORMED)
                _, shape_score, _, _ = cv2.minMaxLoc(res)
                
                # Штраф за цвет
                color_dist = np.linalg.norm(cell_color - np.array(color_database[piece_symbol]))
                score = shape_score - (color_dist / 100.0)
                
                if score > max_score:
                    max_score, best_match = score, piece_symbol
                    
            row_pieces.append(best_match)
        board_matrix.append(row_pieces)

    # Генерация FEN
    fen_rows = []
    for r in board_matrix:
        empty, row_str = 0, ""
        for p in r:
            if p == '.': 
                empty += 1
            else:
                if empty > 0: 
                    row_str += str(empty)
                    empty = 0
                row_str += p
        if empty > 0: 
            row_str += str(empty)
        fen_rows.append(row_str)
        
    return board_matrix, "/".join(fen_rows) + " w - - 0 1"

def main():
    global is_active, board_roi, need_selection

    # Проверка наличия папки с шаблонами
    if not os.path.exists(TEMPLATES_DIR):
        print(f"[!] ОШИБКА: Папка '{TEMPLATES_DIR}' с изображениями фигур не найдена!")
        print("Создайте папку 'extracted_pieces' и положите туда шаблоны фигур.")
        return

    # Авто-генерация базы цветов
    if not os.path.exists(COLORS_FILE):
        generate_colors_database()

    # Загрузка ROI если есть
    if os.path.exists(CONFIG_FILE):
        try:
            with open(CONFIG_FILE, 'r', encoding='utf-8') as f:
                board_roi = json.load(f)
        except: 
            board_roi = None

    # Загружаем созданный/существующий colors.json
    with open(COLORS_FILE, 'r', encoding='utf-8') as f:
        color_database = json.load(f)

    # Подготовка Canny шаблонов
    templates_edges = {}
    for sym, (filename, _) in PIECES_CONFIG.items():
        t_path = os.path.join(TEMPLATES_DIR, filename)
        if os.path.exists(t_path):
            t_img = cv2.imread(t_path)
            t_img = cv2.resize(t_img, (50, 50))
            gray_t = cv2.cvtColor(t_img, cv2.COLOR_BGR2GRAY)
            templates_edges[sym] = cv2.Canny(gray_t[7:43, 7:43], 50, 150)

    sct = mss()
    keyboard.add_hotkey('f7', toggle_scanner)
    keyboard.add_hotkey('f8', request_selection)

    last_matrix, last_fen = [['.']*8 for _ in range(8)], "---"

    print("Программа успешно запущена!")
    print("Нажмите [F8] для выбора области доски мышкой.")

    while True:
        start_time = time.time()

        if need_selection:
            need_selection = False
            select_board_region(sct)

        if is_active:
            if board_roi is None:
                clear_console()
                print("⚠️ Область доски не задана! Нажмите [F8] для выбора мышкой.")
                time.sleep(1)
                continue
            try:
                frame = np.array(sct.grab(board_roi))[:, :, :3]
                last_matrix, last_fen = scan_frame(frame, color_database, templates_edges)
            except Exception as e:
                print(f"Ошибка захвата экрана: {e}")
                board_roi = None
            print_visual_board(last_matrix, last_fen)
        else:
            print_visual_board(last_matrix, last_fen)

        time.sleep(max(0.01, 0.4 - (time.time() - start_time)))

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"\n[КРИТИЧЕСКАЯ ОШИБКА]: {e}")
        input("\nНажмите Enter, чтобы закрыть программу...")