dcs-retribution/qt_ui/windows/basemenu/NewUnitTransferDialog.py

from __future__ import annotations

import logging
from collections import defaultdict
from dataclasses import dataclass
from typing import Callable, Dict, Type

from PySide2.QtCore import Qt
from PySide2.QtWidgets import (
    QCheckBox,
    QComboBox,
    QDialog,
    QFrame,
    QGridLayout,
    QGroupBox,
    QHBoxLayout,
    QLabel,
    QMessageBox,
    QPushButton,
    QScrollArea,
    QSizePolicy,
    QSpacerItem,
    QVBoxLayout,
    QWidget,
)
from dcs.task import PinpointStrike
from dcs.unittype import UnitType

from game import Game, db
from game.theater import ControlPoint, SupplyRoute
from game.transfers import AirliftPlanner, RoadTransferOrder
from qt_ui.models import GameModel
from qt_ui.widgets.QLabeledWidget import QLabeledWidget


class TransferDestinationComboBox(QComboBox):
    def __init__(self, game: Game, origin: ControlPoint) -> None:
        super().__init__()
        self.game = game
        self.origin = origin

        for cp in self.game.theater.controlpoints:
            if (
                cp != self.origin
                and cp.is_friendly(to_player=True)
                and cp.can_deploy_ground_units
            ):
                self.addItem(cp.name, cp)
        self.model().sort(0)
        self.setCurrentIndex(0)


class UnitTransferList(QFrame):
    def __init__(self, cp: ControlPoint, game_model: GameModel):
        super().__init__(self)
        self.cp = cp
        self.game_model = game_model

        self.bought_amount_labels = {}
        self.existing_units_labels = {}

        main_layout = QVBoxLayout()
        self.setLayout(main_layout)

        scroll_content = QWidget()
        task_box_layout = QGridLayout()
        scroll_content.setLayout(task_box_layout)

        units_column = sorted(
            cp.base.armor,
            key=lambda u: db.unit_get_expanded_info(
                self.game_model.game.player_country, u, "name"
            ),
        )

        count = 0
        for count, unit_type in enumerate(units_column):
            self.add_purchase_row(unit_type, task_box_layout, count)
        stretch = QVBoxLayout()
        stretch.addStretch()
        task_box_layout.addLayout(stretch, count, 0)

        scroll_content.setLayout(task_box_layout)
        scroll = QScrollArea()
        scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        scroll.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOn)
        scroll.setWidgetResizable(True)
        scroll.setWidget(scroll_content)
        main_layout.addWidget(scroll)


@dataclass(frozen=True)
class AirliftCapacity:
    helicopter: int
    cargo_plane: int

    @property
    def total(self) -> int:
        return self.helicopter + self.cargo_plane

    @classmethod
    def to_control_point(cls, game: Game) -> AirliftCapacity:
        helo_capacity = 0
        plane_capacity = 0
        for cp in game.theater.player_points():
            inventory = game.aircraft_inventory.for_control_point(cp)
            for unit_type, count in inventory.all_aircraft:
                if unit_type.helicopter:
                    helo_capacity += count
        return AirliftCapacity(helicopter=helo_capacity, cargo_plane=plane_capacity)


class TransferOptionsPanel(QVBoxLayout):
    def __init__(
        self,
        game: Game,
        origin: ControlPoint,
        airlift_capacity: AirliftCapacity,
        airlift_required: bool,
    ) -> None:
        super().__init__()

        self.source_combo_box = TransferDestinationComboBox(game, origin)
        self.addLayout(QLabeledWidget("Destination:", self.source_combo_box))
        self.airlift = QCheckBox()
        self.airlift.setChecked(airlift_required)
        self.airlift.setDisabled(airlift_required)
        self.addLayout(QLabeledWidget("Airlift:", self.airlift))
        self.addWidget(
            QLabel(
                f"{airlift_capacity.total} airlift capacity "
                f"({airlift_capacity.cargo_plane} from cargo planes, "
                f"{airlift_capacity.helicopter} from helicopters)"
            )
        )

    @property
    def changed(self):
        return self.source_combo_box.currentIndexChanged

    @property
    def current(self) -> ControlPoint:
        return self.source_combo_box.currentData()


class TransferControls(QGroupBox):
    def __init__(
        self,
        increase_text: str,
        on_increase: Callable[[TransferControls], None],
        decrease_text: str,
        on_decrease: Callable[[TransferControls], None],
        initial_amount: int = 0,
        disabled: bool = False,
    ) -> None:
        super().__init__()

        self.quantity = initial_amount

        self.setProperty("style", "buy-box")
        self.setMaximumHeight(36)
        self.setMinimumHeight(36)
        layout = QHBoxLayout()
        self.setLayout(layout)

        decrease = QPushButton(decrease_text)
        decrease.setProperty("style", "btn-sell")
        decrease.setDisabled(disabled)
        decrease.setMinimumSize(16, 16)
        decrease.setMaximumSize(16, 16)
        decrease.setSizePolicy(QSizePolicy(QSizePolicy.Fixed, QSizePolicy.Fixed))
        decrease.clicked.connect(lambda: on_decrease(self))
        layout.addWidget(decrease)

        self.count_label = QLabel()
        self.count_label.setSizePolicy(
            QSizePolicy(QSizePolicy.Fixed, QSizePolicy.Fixed)
        )
        self.set_quantity(initial_amount)
        layout.addWidget(self.count_label)

        increase = QPushButton(increase_text)
        increase.setProperty("style", "btn-buy")
        increase.setDisabled(disabled)
        increase.setMinimumSize(16, 16)
        increase.setMaximumSize(16, 16)
        increase.clicked.connect(lambda: on_increase(self))
        increase.setSizePolicy(QSizePolicy(QSizePolicy.Fixed, QSizePolicy.Fixed))
        layout.addWidget(increase)

    def set_quantity(self, quantity: int) -> None:
        self.quantity = quantity
        self.count_label.setText(f"<b>{self.quantity}</b>")


class ScrollingUnitTransferGrid(QFrame):
    def __init__(
        self,
        cp: ControlPoint,
        airlift: bool,
        airlift_capacity: AirliftCapacity,
        game_model: GameModel,
    ) -> None:
        super().__init__()
        self.cp = cp
        self.airlift = airlift
        self.remaining_capacity = airlift_capacity.total
        self.game_model = game_model
        self.transfers: Dict[Type[UnitType, int]] = defaultdict(int)

        main_layout = QVBoxLayout()

        scroll_content = QWidget()
        task_box_layout = QGridLayout()

        unit_types = set(
            db.find_unittype(PinpointStrike, self.game_model.game.player_name)
        )
        sorted_units = sorted(
            {u for u in unit_types if self.cp.base.total_units_of_type(u)},
            key=lambda u: db.unit_get_expanded_info(
                self.game_model.game.player_country, u, "name"
            ),
        )
        for row, unit_type in enumerate(sorted_units):
            self.add_unit_row(unit_type, task_box_layout, row)
        stretch = QVBoxLayout()
        stretch.addStretch()
        task_box_layout.addLayout(stretch, task_box_layout.count(), 0)

        scroll_content.setLayout(task_box_layout)
        scroll = QScrollArea()
        scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        scroll.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOn)
        scroll.setWidgetResizable(True)
        scroll.setWidget(scroll_content)
        main_layout.addWidget(scroll)
        self.setLayout(main_layout)

    def add_unit_row(
        self,
        unit_type: Type[UnitType],
        layout: QGridLayout,
        row: int,
    ) -> None:
        exist = QGroupBox()
        exist.setProperty("style", "buy-box")
        exist.setMaximumHeight(36)
        exist.setMinimumHeight(36)
        origin_inventory_layout = QHBoxLayout()
        exist.setLayout(origin_inventory_layout)

        origin_inventory = self.cp.base.total_units_of_type(unit_type)

        unit_name = QLabel(
            "<b>"
            + db.unit_get_expanded_info(
                self.game_model.game.player_country, unit_type, "name"
            )
            + "</b>"
        )
        unit_name.setSizePolicy(
            QSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        )

        origin_inventory_label = QLabel(str(origin_inventory))
        origin_inventory_label.setSizePolicy(
            QSizePolicy(QSizePolicy.Fixed, QSizePolicy.Fixed)
        )

        def increase(controls: TransferControls):
            nonlocal origin_inventory
            nonlocal origin_inventory_label
            if not origin_inventory:
                return

            if self.airlift:
                if not self.remaining_capacity:
                    return
                self.remaining_capacity -= 1

            self.transfers[unit_type] += 1
            origin_inventory -= 1
            controls.set_quantity(self.transfers[unit_type])
            origin_inventory_label.setText(str(origin_inventory))

        def decrease(controls: TransferControls):
            nonlocal origin_inventory
            nonlocal origin_inventory_label
            if not controls.quantity:
                return

            if self.airlift:
                self.remaining_capacity += 1

            self.transfers[unit_type] -= 1
            origin_inventory += 1
            controls.set_quantity(self.transfers[unit_type])
            origin_inventory_label.setText(str(origin_inventory))

        transfer_controls = TransferControls("->", increase, "<-", decrease)

        origin_inventory_layout.addWidget(unit_name)
        origin_inventory_layout.addItem(
            QSpacerItem(20, 0, QSizePolicy.Minimum, QSizePolicy.Minimum)
        )
        origin_inventory_layout.addWidget(origin_inventory_label)
        origin_inventory_layout.addItem(
            QSpacerItem(20, 0, QSizePolicy.Minimum, QSizePolicy.Minimum)
        )

        layout.addWidget(exist, row, 1)
        layout.addWidget(transfer_controls, row, 2)


class NewUnitTransferDialog(QDialog):
    def __init__(
        self,
        game_model: GameModel,
        origin: ControlPoint,
        parent=None,
    ) -> None:
        super().__init__(parent)
        self.origin = origin
        self.setWindowTitle(f"New unit transfer from {origin.name}")

        self.game_model = game_model

        layout = QVBoxLayout()
        self.setLayout(layout)

        self.airlift_capacity = AirliftCapacity.to_control_point(game_model.game)
        airlift_required = len(SupplyRoute.for_control_point(origin)) == 1
        self.dest_panel = TransferOptionsPanel(
            game_model.game, origin, self.airlift_capacity, airlift_required
        )
        self.dest_panel.changed.connect(self.rebuild_transfers)
        layout.addLayout(self.dest_panel)

        self.transfer_panel = ScrollingUnitTransferGrid(
            origin,
            airlift=airlift_required,
            airlift_capacity=self.airlift_capacity,
            game_model=game_model,
        )
        self.dest_panel.airlift.toggled.connect(self.rebuild_transfers)
        layout.addWidget(self.transfer_panel)

        self.submit_button = QPushButton("Create Transfer Order", parent=self)
        self.submit_button.clicked.connect(self.on_submit)
        self.submit_button.setProperty("style", "start-button")
        layout.addWidget(self.submit_button)

    def rebuild_transfers(self) -> None:
        # Rebuild the transfer panel to reset everything. It's easier to recreate the
        # panel itself than to clear the grid layout in the panel.
        self.layout().removeWidget(self.transfer_panel)
        self.layout().removeWidget(self.submit_button)
        self.transfer_panel = ScrollingUnitTransferGrid(
            self.origin,
            airlift=self.dest_panel.airlift.isChecked(),
            airlift_capacity=self.airlift_capacity,
            game_model=self.game_model,
        )
        self.layout().addWidget(self.transfer_panel)
        self.layout().addWidget(self.submit_button)

    def on_submit(self) -> None:
        destination = self.dest_panel.current
        supply_route = SupplyRoute.for_control_point(self.origin)
        if not self.dest_panel.airlift.isChecked() and destination not in supply_route:
            QMessageBox.critical(
                self,
                "Could not create transfer",
                f"Transfers from {self.origin} to {destination} require airlift.",
                QMessageBox.Ok,
            )
            return
        transfers = {}
        for unit_type, count in self.transfer_panel.transfers.items():
            if not count:
                continue

            logging.info(
                f"Transferring {count} {unit_type.id} from {self.origin} to "
                f"{destination}"
            )
            transfers[unit_type] = count

        if self.dest_panel.airlift.isChecked():
            planner = AirliftPlanner(
                self.game_model.game,
                self.origin,
                destination,
                transfers,
            )
            planner.create_package_for_airlift()
        else:
            transfer = RoadTransferOrder(
                player=True,
                origin=self.origin,
                destination=destination,
                units=transfers,
            )
            self.game_model.transfer_model.new_transfer(transfer)
        self.close()