dcs-retribution/qt_ui/windows/AirWingDialog.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterator, Optional

from PySide6.QtCore import QItemSelectionModel, QModelIndex, QSize
from PySide6.QtWidgets import (
    QAbstractItemView,
    QCheckBox,
    QDialog,
    QHBoxLayout,
    QListView,
    QTabWidget,
    QTableWidget,
    QTableWidgetItem,
    QVBoxLayout,
    QWidget,
    QPushButton,
)

from game.armedforces.armedforces import ArmedForces
from game.ato.flight import Flight
from game.factions import Faction
from game.server import EventStream
from game.sim import GameUpdateEvents
from game.squadrons import Squadron
from game.theater import ConflictTheater
from qt_ui.delegates import TwoColumnRowDelegate
from qt_ui.models import AirWingModel, AtoModel, GameModel, SquadronModel
from qt_ui.simcontroller import SimController
from qt_ui.windows.AirWingConfigurationDialog import AirWingConfigurationDialog
from qt_ui.windows.SquadronDialog import SquadronDialog
from qt_ui.windows.newgame.WizardPages.QFactionSelection import QFactionUnits


class SquadronDelegate(TwoColumnRowDelegate):
    def __init__(self, air_wing_model: AirWingModel) -> None:
        super().__init__(rows=2, columns=2, font_size=12)
        self.air_wing_model = air_wing_model

    @staticmethod
    def squadron(index: QModelIndex) -> Squadron:
        return index.data(AirWingModel.SquadronRole)

    def text_for(self, index: QModelIndex, row: int, column: int) -> str:
        squadron = self.squadron(index)
        if (row, column) == (0, 0):
            if squadron.nickname:
                nickname = f' "{squadron.nickname}"'
            else:
                nickname = ""
            return f"{squadron.name}{nickname}"
        elif (row, column) == (0, 1):
            return squadron.aircraft.display_name
        elif (row, column) == (1, 0):
            return squadron.location.name
        elif (row, column) == (1, 1):
            squadron = self.squadron(index)
            active = len(squadron.active_pilots)
            available = len(squadron.available_pilots)
            on_leave = len(squadron.pilots_on_leave)
            return f"{active} active, {available} unassigned, {on_leave} on leave"
        return ""


class SquadronList(QListView):
    """List view for displaying the air wing's squadrons."""

    def __init__(
        self,
        ato_model: AtoModel,
        air_wing_model: AirWingModel,
        theater: ConflictTheater,
        sim_controller: SimController,
    ) -> None:
        super().__init__()
        self.ato_model = ato_model
        self.air_wing_model = air_wing_model
        self.theater = theater
        self.sim_controller = sim_controller
        self.dialog: Optional[SquadronDialog] = None

        self.setIconSize(QSize(91, 24))
        self.setItemDelegate(SquadronDelegate(self.air_wing_model))
        self.setModel(self.air_wing_model)
        self.selectionModel().setCurrentIndex(
            self.air_wing_model.index(0, 0, QModelIndex()),
            QItemSelectionModel.SelectionFlag.Select,
        )

        # self.setIconSize(QSize(91, 24))
        self.setSelectionBehavior(QAbstractItemView.SelectionBehavior.SelectItems)
        self.doubleClicked.connect(self.on_double_click)

    def on_double_click(self, index: QModelIndex) -> None:
        if not index.isValid():
            return
        self.dialog = SquadronDialog(
            self.ato_model,
            SquadronModel(self.air_wing_model.squadron_at_index(index)),
            self.theater,
            self.sim_controller,
            self,
        )
        self.dialog.show()


@dataclass(frozen=True)
class AircraftInventoryData:
    location: str
    unit_type: str
    task: str
    target: str
    pilot: str
    player: str

    @classmethod
    def headers(cls) -> list[str]:
        return ["Base", "Type", "Flight Type", "Target", "Pilot", "Player"]

    @property
    def columns(self) -> Iterator[str]:
        yield self.location
        yield self.unit_type
        yield self.task
        yield self.target
        yield self.pilot
        yield self.player

    @classmethod
    def from_flight(cls, flight: Flight) -> Iterator[AircraftInventoryData]:
        num_units = flight.count
        flight_type = flight.flight_type.value
        target = flight.package.target.name
        for idx in range(0, num_units):
            pilot = flight.roster.pilot_at(idx)
            if pilot is None:
                pilot_name = "Unassigned"
                player = ""
            else:
                pilot_name = pilot.name
                player = "Player" if pilot.player else "AI"
            yield AircraftInventoryData(
                flight.departure.name,
                flight.unit_type.display_name,
                flight_type,
                target,
                pilot_name,
                player,
            )

    @classmethod
    def each_untasked_from_squadron(
        cls, squadron: Squadron
    ) -> Iterator[AircraftInventoryData]:
        for _ in range(0, squadron.untasked_aircraft):
            yield AircraftInventoryData(
                squadron.name,
                squadron.aircraft.display_name,
                "Idle",
                "N/A",
                "N/A",
                "N/A",
            )


class AirInventoryView(QWidget):
    def __init__(self, game_model: GameModel) -> None:
        super().__init__()
        self.game_model = game_model

        self.only_unallocated = False
        self.enemy_info = False

        layout = QVBoxLayout()
        self.setLayout(layout)

        checkbox_row = QHBoxLayout()
        layout.addLayout(checkbox_row)

        self.only_unallocated_cb = QCheckBox("Unallocated only")
        self.only_unallocated_cb.toggled.connect(self.set_only_unallocated)
        checkbox_row.addWidget(self.only_unallocated_cb)

        self.enemy_info_cb = QCheckBox("Show enemy info")
        self.enemy_info_cb.toggled.connect(self.set_enemy_info)
        checkbox_row.addWidget(self.enemy_info_cb)

        checkbox_row.addStretch()

        self.table = QTableWidget()
        layout.addWidget(self.table)

        self.table.setEditTriggers(QAbstractItemView.EditTrigger.NoEditTriggers)
        self.table.verticalHeader().setVisible(False)
        self.set_only_unallocated(False)

    def set_only_unallocated(self, value: bool) -> None:
        self.only_unallocated = value
        self.update_table()

    def set_enemy_info(self, value: bool) -> None:
        self.enemy_info = value
        self.update_table()

    def update_table(self) -> None:
        self.table.setSortingEnabled(False)
        self.table.clear()

        inventory_rows = list(self.get_data())
        self.table.setRowCount(len(inventory_rows))
        headers = AircraftInventoryData.headers()
        self.table.setColumnCount(len(headers))
        self.table.setHorizontalHeaderLabels(headers)

        for row, data in enumerate(inventory_rows):
            for column, value in enumerate(data.columns):
                self.table.setItem(row, column, QTableWidgetItem(value))

        self.table.resizeColumnsToContents()
        self.table.setSortingEnabled(True)

    def iter_allocated_aircraft(self) -> Iterator[AircraftInventoryData]:
        coalition = self.game_model.game.coalition_for(not self.enemy_info)
        for package in coalition.ato.packages:
            for flight in package.flights:
                yield from AircraftInventoryData.from_flight(flight)

    def iter_unallocated_aircraft(self) -> Iterator[AircraftInventoryData]:
        coalition = self.game_model.game.coalition_for(not self.enemy_info)
        for squadron in coalition.air_wing.iter_squadrons():
            yield from AircraftInventoryData.each_untasked_from_squadron(squadron)

    def get_data(self) -> Iterator[AircraftInventoryData]:
        yield from self.iter_unallocated_aircraft()
        if not self.only_unallocated:
            yield from self.iter_allocated_aircraft()


class AirWingTabs(QTabWidget):
    def __init__(self, game_model: GameModel) -> None:
        super().__init__()

        self.game_model = game_model

        self.addTab(
            SquadronList(
                game_model.ato_model,
                game_model.blue_air_wing_model,
                game_model.game.theater,
                game_model.sim_controller,
            ),
            "Squadrons OWNFOR",
        )
        self.addTab(
            SquadronList(
                game_model.red_ato_model,
                game_model.red_air_wing_model,
                game_model.game.theater,
                game_model.sim_controller,
            ),
            "Squadrons OPFOR",
        )
        self.addTab(AirInventoryView(game_model), "Inventory")

        if game_model.game.settings.enable_air_wing_adjustments:
            pb = QPushButton("Open Air Wing Config Dialog")
            pb.clicked.connect(self.open_awcd)
            pb.setMaximumWidth(300)
            layout = QHBoxLayout()
            layout.addWidget(pb)
            w = QWidget(layout=layout)
            self.addTab(w, "Cheats")

        qfu_ownfor = QFactionUnits(
            game_model.game.coalition_for(True).faction,
            self,
            show_jtac=True,
        )
        qfu_ownfor.preset_groups_changed.connect(self.preset_group_updated_ownfor)
        self.addTab(
            qfu_ownfor,
            "Faction OWNFOR",
        )
        qfu_opfor = QFactionUnits(
            game_model.game.coalition_for(False).faction,
            self,
            show_jtac=True,
        )
        qfu_opfor.preset_groups_changed.connect(self.preset_group_updated_opfor)
        self.addTab(
            qfu_opfor,
            "Faction OPFOR",
        )

    def open_awcd(self):
        AirWingConfigurationDialog(self.game_model.game, True, self).exec_()
        events = GameUpdateEvents().begin_new_turn()
        EventStream.put_nowait(events)
        self.game_model.ato_model.on_sim_update(events)

    def preset_group_updated_ownfor(self, f: Faction) -> None:
        self.preset_group_updated(f, player=True)

    def preset_group_updated_opfor(self, f: Faction) -> None:
        self.preset_group_updated(f, player=False)

    def preset_group_updated(self, f: Faction, player: bool) -> None:
        self.game_model.game.coalition_for(player).armed_forces = ArmedForces(f)


class AirWingDialog(QDialog):
    """Dialog window showing the player's air wing."""

    def __init__(self, game_model: GameModel, parent) -> None:
        super().__init__(parent)
        self.air_wing_model = game_model.blue_air_wing_model

        self.setMinimumSize(1000, 440)
        self.setWindowTitle(f"Air Wing")
        # TODO: self.setWindowIcon()

        layout = QVBoxLayout()
        self.setLayout(layout)

        layout.addWidget(AirWingTabs(game_model))