from typing import NamedTuple

from PySide6.QtCore import Qt
from PySide6.QtWidgets import (
    QCheckBox,
    QComboBox,
    QFormLayout,
    QGridLayout,
    QGroupBox,
    QLabel,
    QLineEdit,
    QMenu,
)

from ...models.animal import AnimalConfig
from ...models.session import RewardEnum, SessionConfig
from .animal import AnimalCard


class ExperimentConfigGroup(QGroupBox):
    def __init__(self, parent, experimenters: list[str]):
        super().__init__("Config", parent)

        layout = QGridLayout(self)
        self.setLayout(layout)
        layout.setColumnStretch(1, 1)
        layout.setColumnStretch(3, 1)

        # Row 0: experimenter | hide cursor
        layout.addWidget(QLabel("experimenter", self), 0, 0)
        self.combo_experimenter = QComboBox(self)
        self.combo_experimenter.addItems(experimenters)
        layout.addWidget(self.combo_experimenter, 0, 1)

        layout.addWidget(QLabel("hide cursor", self), 0, 2)
        self.check_hide_cursor = QCheckBox(self)
        layout.addWidget(self.check_hide_cursor, 0, 3)

        # Row 1: reward type | fullscreen
        layout.addWidget(QLabel("reward type", self), 1, 0)
        self.combo_reward_type = QComboBox(self)
        self.combo_reward_type.addItems(list(RewardEnum))
        layout.addWidget(self.combo_reward_type, 1, 1)

        layout.addWidget(QLabel("fullscreen", self), 1, 2)
        self.check_fullscreen = QCheckBox(self)
        layout.addWidget(self.check_fullscreen, 1, 3)

        # Row 2: notes (span all columns)
        layout.addWidget(QLabel("notes", self), 2, 0)
        self.line_notes = QLineEdit(self)
        self.line_notes.setPlaceholderText("Notes")
        layout.addWidget(self.line_notes, 2, 1, 1, 3)

    def load_config(self, config: SessionConfig):
        self.combo_experimenter.setEditText(config.experimenter)
        self.combo_reward_type.setEditText(config.reward_type)
        self.check_hide_cursor.setChecked(config.hide_cursor)
        self.check_fullscreen.setChecked(config.fullscreen)

    def result(self) -> SessionConfig:
        return SessionConfig(
            experimenter=self.combo_experimenter.currentText(),
            reward_type=RewardEnum(self.combo_reward_type.currentText()),
            note=self.line_notes.text(),
            hide_cursor=self.check_hide_cursor.isChecked(),
            fullscreen=self.check_fullscreen.isChecked(),
        )


class SceneSelection(NamedTuple):
    default_scene: str
    unknown_animal_fallback: str
    fault_fallback: str


class ExperimentSceneGroup(QGroupBox):
    def __init__(self, parent=None, *, stages: list[str]):
        super().__init__("Scene", parent)

        layout = QFormLayout(self)
        self.setLayout(layout)

        label_default = QLabel("default scene", self)
        self.combo_default_scene = QComboBox(self)
        self.combo_default_scene.addItems(stages)
        layout.addRow(label_default, self.combo_default_scene)

        label_unknown = QLabel("unknown animal fallback", self)
        self.combo_unknown_animal_fallback = QComboBox(self)
        self.combo_unknown_animal_fallback.addItems(stages)
        layout.addRow(label_unknown, self.combo_unknown_animal_fallback)

        label_fault = QLabel("fault fallback", self)
        self.combo_fault_fallback = QComboBox(self)
        self.combo_fault_fallback.addItems(stages)
        layout.addRow(label_fault, self.combo_fault_fallback)

    def load_config(self, config: SessionConfig):
        if config.default_scene:
            self.combo_default_scene.setCurrentText(config.default_scene)
        if config.unknown_animal_fallback:
            self.combo_unknown_animal_fallback.setCurrentText(
                config.unknown_animal_fallback
            )
        if config.fault_fallback:
            self.combo_fault_fallback.setCurrentText(config.fault_fallback)

    def result(self) -> SceneSelection:
        return SceneSelection(
            default_scene=self.combo_default_scene.currentText(),
            unknown_animal_fallback=self.combo_unknown_animal_fallback.currentText(),
            fault_fallback=self.combo_fault_fallback.currentText(),
        )


class ExperimentAnimalsGroup(QGroupBox):
    def __init__(
        self,
        parent=None,
        *,
        animals: dict[str, str],
        stage_level_tables: dict[str, dict[str, list[int]]],
    ):
        super().__init__("Animals", parent)

        self._animals = animals
        self._stage_level_tables = stage_level_tables
        self.layout_animals = QGridLayout(self)
        self.setLayout(self.layout_animals)

        self.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
        self.customContextMenuRequested.connect(self._on_context_menu)

    def _on_context_menu(self, pos):
        menu = QMenu(self)
        action = menu.addAction("Add animal")
        action.triggered.connect(lambda _checked=False: self._on_add_animal())
        menu.exec(self.mapToGlobal(pos))

    def _on_add_animal(self):
        animal_card = AnimalCard(self, self._animals, self._stage_level_tables)
        animal_card.remove_requested.connect(
            lambda _card=animal_card: self._on_remove_animal(_card)
        )
        self._add_animal_card(animal_card)

    def _add_animal_card(self, animal_card: AnimalCard) -> None:
        index = self.layout_animals.count()
        row = index // 4
        col = index % 4
        self.layout_animals.addWidget(animal_card, row, col)

    def _on_remove_animal(self, animal_card: AnimalCard) -> None:
        self.layout_animals.removeWidget(animal_card)
        animal_card.setParent(None)
        animal_card.deleteLater()
        self._reflow_animal_cards()

    def _reflow_animal_cards(self) -> None:
        cards: list[AnimalCard] = []
        for i in range(self.layout_animals.count()):
            item = self.layout_animals.itemAt(i)
            if item:
                widget = item.widget()
                if isinstance(widget, AnimalCard):
                    cards.append(widget)

        while self.layout_animals.count() > 0:
            self.layout_animals.takeAt(0)

        for card in cards:
            self._add_animal_card(card)

    def load_config(self, config: SessionConfig):
        for animal_configs in config.animals:
            animal_card = AnimalCard(self, self._animals, self._stage_level_tables)
            animal_card.load_config(animal_configs)
            animal_card.remove_requested.connect(
                lambda _card=animal_card: self._on_remove_animal(_card)
            )
            self._add_animal_card(animal_card)

    def result(self) -> list[AnimalConfig]:
        results: list[AnimalConfig] = []
        for i in range(self.layout_animals.count()):
            item = self.layout_animals.itemAt(i)
            if not item:
                continue
            widget = item.widget()
            if isinstance(widget, AnimalCard):
                results.append(widget.result)
        return results
