from PySide6.QtCore import QSize, QRectF
from app.ui.key import KeyItem


class KeyboardLayout():
    def __init__(self, scene):
        self.scene = scene
        self.keys = {}
        self.letter_rows = [
            "QWERTZUIOP",
            "ASDFGHJKL",
            "YXCVBNM",
        ]
        self.space_key = KeyItem("Space")
        self.space_key.clicked.connect(lambda: self.on_key_clicked("Space"))
        self.scene.addItem(self.space_key)
        self.create_keys()

        # callback to be set by scene
        self.key_pressed_callback = None

    def create_keys(self):
        for row in self.letter_rows:
            for char in row:
                key = KeyItem(char)
                key.clicked.connect(lambda c=char: self.on_key_clicked(c))
                self.scene.addItem(key)
                self.keys[char] = key

    def layout_keys(self, view_size: QSize, top_offset: float = 0.0):
        view_width = view_size.width()
        view_height = view_size.height()

        padding = 10.0
        spacing = 6.0
        rows = self.letter_rows + [" "]
        row_count = len(rows)
        max_keys_in_row = max(len(row) for row in rows)

        total_spacing_x = (max_keys_in_row - 1) * spacing
        total_spacing_y = (row_count - 1) * spacing

        available_width = view_width - 2 * padding - total_spacing_x
        available_height = view_height - top_offset - 2 * padding - total_spacing_y

        key_width = available_width / max_keys_in_row
        key_height = available_height / row_count

        y = top_offset + padding
        for row in self.letter_rows:
            x = padding
            for char in row:
                key = self.keys[char]
                key.set_geometry(x, y, key_width, key_height)
                x += key_width + spacing
            y += key_height + spacing

        space_width = key_width * 7 + spacing * 6
        self.space_key.set_geometry(padding, y, space_width, key_height)

    def set_scale_factors(self, predictions: dict):
        for key in self.keys.values():
            key.set_scale_factor(1.0)
        for key, probability in predictions.items():
            if key.upper() in self.keys:
                self.keys[key.upper()].set_scale_factor(1 + probability)

    def on_key_clicked(self, label: str):
        if self.key_pressed_callback:
            self.key_pressed_callback(label)