StarHorizon_Public/Content.Client/_Mono/FireControl/UI/FireControlWindow.xaml.cs

// Copyright Rane (elijahrane@gmail.com) 2025
// All rights reserved. Relicensed under AGPL with permission

using Content.Client.UserInterface.Controls;
using Content.Shared._Mono.FireControl;
using Content.Shared._Mono.ShipGuns;
using Robust.Client.AutoGenerated;
using Robust.Client.UserInterface.XAML;
using Robust.Client.UserInterface.Controls;

namespace Content.Client._Mono.FireControl.UI;

[GenerateTypedNameReferences]
public sealed partial class FireControlWindow : FancyWindow
{
    [Dependency] private readonly IEntityManager _entityManager = default!;

    public FireControlNavControl Radar => NavRadar;
    public Action? OnServerRefresh;
    public Action? OnWeaponSelectionChanged;

    public readonly Dictionary<NetEntity, Button> WeaponsList = new();

    // Dictionary to store weapon name to entity mapping for easier filtering
    private readonly Dictionary<string, NetEntity> _weaponNameToEntity = new();

    // Dictionary to store weapon entity to type mapping
    private readonly Dictionary<NetEntity, ShipGunType> _weaponTypes = new();

    public FireControlWindow()
    {
        RobustXamlLoader.Load(this);
        IoCManager.InjectDependencies(this);
        RefreshButton.OnPressed += _ => OnServerRefresh?.Invoke();
        SelectAllButton.OnPressed += SelectAllWeapons;
        UnselectAllButton.OnPressed += UnselectAllWeapons;
        SelectBallisticButton.OnPressed += SelectBallisticWeapons;
        SelectEnergyButton.OnPressed += SelectEnergyWeapons;
        SelectMissileButton.OnPressed += SelectMissileWeapons;
    }

    private void SelectAllWeapons(BaseButton.ButtonEventArgs args)
    {
        foreach (var button in WeaponsList.Values)
        {
            button.Pressed = true;
        }

        OnWeaponSelectionChanged?.Invoke();
    }

    private void UnselectAllWeapons(BaseButton.ButtonEventArgs args)
    {
        foreach (var button in WeaponsList.Values)
        {
            button.Pressed = false;
        }

        OnWeaponSelectionChanged?.Invoke();
    }

    private void SelectBallisticWeapons(BaseButton.ButtonEventArgs args)
    {
        // First unselect all weapons
        foreach (var button in WeaponsList.Values)
        {
            button.Pressed = false;
        }

        // Then select only ballistic weapons
        foreach (var kvp in WeaponsList)
        {
            var weaponEntity = kvp.Key;
            var button = kvp.Value;

            if (_weaponTypes.TryGetValue(weaponEntity, out var type) && type == ShipGunType.Ballistic)
            {
                button.Pressed = true;
            }
        }

        OnWeaponSelectionChanged?.Invoke();
    }

    private void SelectEnergyWeapons(BaseButton.ButtonEventArgs args)
    {
        // First unselect all weapons
        foreach (var button in WeaponsList.Values)
        {
            button.Pressed = false;
        }

        // Then select only energy weapons
        foreach (var kvp in WeaponsList)
        {
            var weaponEntity = kvp.Key;
            var button = kvp.Value;

            if (_weaponTypes.TryGetValue(weaponEntity, out var type) && type == ShipGunType.Energy)
            {
                button.Pressed = true;
            }
        }

        OnWeaponSelectionChanged?.Invoke();
    }

    private void SelectMissileWeapons(BaseButton.ButtonEventArgs args)
    {
        // First unselect all weapons
        foreach (var button in WeaponsList.Values)
        {
            button.Pressed = false;
        }

        // Then select only missile weapons
        foreach (var kvp in WeaponsList)
        {
            var weaponEntity = kvp.Key;
            var button = kvp.Value;

            if (_weaponTypes.TryGetValue(weaponEntity, out var type) && type == ShipGunType.Missile)
            {
                button.Pressed = true;
            }
        }

        OnWeaponSelectionChanged?.Invoke();
    }

    public void UpdateStatus(FireControlConsoleBoundInterfaceState state)
    {
        NavRadar.UpdateState(state.NavState);

        if (state.Connected)
        {
            ServerStatus.Text = Loc.GetString("gunnery-window-connected");
            ServerStatus.FontColorOverride = Color.Green;
        }
        else
        {
            RefreshButton.Disabled = false;
            ServerStatus.Text = Loc.GetString("gunnery-window-disconnected");
            ServerStatus.FontColorOverride = Color.Red;
        }

        UpdateWeaponsList(state);

        // Update the category buttons state based on whether weapons of that type are available
        bool hasBallisticWeapons = false;
        bool hasEnergyWeapons = false;
        bool hasMissileWeapons = false;

        foreach (var kvp in _weaponTypes)
        {
            var type = kvp.Value;

            if (type == ShipGunType.Ballistic)
                hasBallisticWeapons = true;
            else if (type == ShipGunType.Energy)
                hasEnergyWeapons = true;
            else if (type == ShipGunType.Missile)
                hasMissileWeapons = true;

            if (hasBallisticWeapons && hasEnergyWeapons && hasMissileWeapons)
                break;
        }

        SelectBallisticButton.Disabled = !hasBallisticWeapons;
        SelectEnergyButton.Disabled = !hasEnergyWeapons;
        SelectMissileButton.Disabled = !hasMissileWeapons;
    }

    private void UpdateWeaponsList(FireControlConsoleBoundInterfaceState state)
    {
        if (state.FireControllables == null)
        {
            WeaponsList.Clear();
            ControllablesBox.DisposeAllChildren();
            _weaponNameToEntity.Clear();
            _weaponTypes.Clear();
            return;
        }

        var toRemove = new HashSet<NetEntity>(WeaponsList.Keys);
        _weaponNameToEntity.Clear();
        _weaponTypes.Clear();

        foreach (var controllable in state.FireControllables)
        {
            // Store the weapon name and entity for filtering
            _weaponNameToEntity[controllable.Name] = controllable.NetEntity;

            // Store the weapon type if available
            var entity = _entityManager.GetEntity(controllable.NetEntity);
            if (_entityManager.TryGetComponent<ShipGunTypeComponent>(entity, out var typeComp))
            {
                _weaponTypes[controllable.NetEntity] = typeComp.Type;
            }

            if (WeaponsList.TryGetValue(controllable.NetEntity, out var existingButton))
            {
                toRemove.Remove(controllable.NetEntity);
            }
            else
            {
                var button = new Button
                {
                    ToggleMode = true,
                    Text = controllable.Name,
                    StyleClasses = { "ButtonSquare OpenRight" },
                    HorizontalExpand = true,
                    Margin = new Thickness(4, 1)
                };

                button.OnToggled += _ => OnWeaponSelectionChanged?.Invoke();

                ControllablesBox.AddChild(button);
                WeaponsList.Add(controllable.NetEntity, button);
            }
        }

        foreach (var entity in toRemove)
        {
            if (WeaponsList.TryGetValue(entity, out var button))
            {
                ControllablesBox.RemoveChild(button);
                WeaponsList.Remove(entity);
            }
        }

        SelectAllButton.Disabled = WeaponsList.Count == 0;
        UnselectAllButton.Disabled = WeaponsList.Count == 0;
    }
}