wwdpublic/Content.Client/_White/Loadouts/LoadoutPicker.xaml.cs

using System.Diagnostics.CodeAnalysis;
using System.Linq;
using Content.Client._White.DatumContainer;
using Content.Shared.CCVar;
using Content.Shared.Clothing.Loadouts.Prototypes;
using Content.Shared.Clothing.Loadouts.Systems;
using Content.Shared.Customization.Systems;
using Content.Shared.Mind;
using Content.Shared.Preferences;
using Content.Shared.Roles;
using Robust.Client.AutoGenerated;
using Robust.Client.UserInterface;
using Robust.Client.UserInterface.Controls;
using Robust.Client.UserInterface.XAML;
using Robust.Shared.Configuration;
using Robust.Shared.Map;
using Robust.Shared.Prototypes;
using Robust.Shared.Utility;


namespace Content.Client._White.Loadouts;


[GenerateTypedNameReferences]
public sealed partial class LoadoutPicker : Control
{
    [Dependency] private readonly IPrototypeManager _prototypeManager = default!;
    [Dependency] private readonly IConfigurationManager _configManager = default!;
    [Dependency] private readonly IEntityManager _entityManager = default!;

    private readonly Dictionary<string, Loadout> _selectedLoadouts = [];
    private readonly List<LoadoutEntry> _loadoutEntries = new();

    private readonly LocalDatumContainer<string> _customColorTints;
    private readonly LocalDatumContainer<bool> _customHeirloom;
    private readonly LocalDatumContainer<string> _customName;
    private readonly LocalDatumContainer<string> _customDescription;
    private readonly LocalDatumContainer<string> _customContent;

    private List<Loadout> _checkpointLoadouts = new();
    private CharacterRequirementsArgs? _checkpointRequirements;

    public Action<List<Loadout>>? OnLoadoutsChanged;

    private int MaxLoadoutPoints => _configManager.GetCVar(CCVars.GameLoadoutsPoints);

    private int _loadoutPoints = 0;

    public CharacterRequirementsArgs CharacterRequirementsArgs = default!;
    private ProtoId<LoadoutCategoryPrototype>? _selectedLoadoutCategory;
    private Dictionary<ProtoId<LoadoutCategoryPrototype>, List<LoadoutPrototype>> _loadoutCache = [];
    private Dictionary<string, string> _loadoutDisplayNames = new();

    public int LoadoutPoint
    {
        get => _loadoutPoints;

        set
        {
            LoadoutPoints.Text = Loc.GetString("marking-points-remaining", ("points", value));
            _loadoutPoints = value;
        }
    }

    public LoadoutPicker()
    {
        RobustXamlLoader.Load(this);
        IoCManager.InjectDependencies(this);

        _customName = new("custom_name");
        _customDescription = new("custom_description");
        _customContent = new("custom_content");
        _customHeirloom = new("custom_heirloom");
        _customColorTints = new("custom_color_tints");

        SaveButton.OnPressed += SaveButtonPressed;
        SpecialColorTintToggle.OnPressed += SpecialColorTintTogglePressed;
        ResetButton.OnPressed += ResetButtonPressed;
        LoadoutSearch.OnTextChanged += args => Populate(args.Text);
        _prototypeManager.PrototypesReloaded += OnPrototypesReloaded;

        InitializeCache();
        InitializeCategories();
    }

    private void InitializeCache()
    {
        CacheRootCategories();
        _loadoutCache.Clear();
        _loadoutDisplayNames.Clear();

        foreach (var loadoutPrototype in _prototypeManager.EnumeratePrototypes<LoadoutPrototype>())
        {
            if (!_loadoutCache.TryGetValue(loadoutPrototype.Category, out var loadoutList))
            {
                loadoutList = new();
                _loadoutCache.Add(loadoutPrototype.Category, loadoutList);
            }

            loadoutList.Add(loadoutPrototype);

            var displayName = Loc.GetString($"loadout-name-{loadoutPrototype.ID}");
            if (displayName == $"loadout-name-{loadoutPrototype.ID}" && loadoutPrototype.Items.Any())
            {
                var tempEntity = _entityManager.SpawnEntity(loadoutPrototype.Items.First(), MapCoordinates.Nullspace);
                displayName = _entityManager.GetComponent<MetaDataComponent>(tempEntity).EntityName;
                _entityManager.DeleteEntity(tempEntity);
            }
            _loadoutDisplayNames[loadoutPrototype.ID] = displayName;
        }
    }

    private void OnPrototypesReloaded(PrototypesReloadedEventArgs obj)
    {
        if (obj.WasModified<LoadoutPrototype>())
            InitializeCache();
    }

    private void Populate(string searchText)
    {
        if (string.IsNullOrEmpty(searchText))
        {
            EntryBackButton.Visible = CurrentEntry.Parent != null;

            if (_selectedLoadoutCategory != null)
                LoadCategoryButtons(_selectedLoadoutCategory.Value);
            else
            {
                ClearLoadoutCategoryButtons();
                CurrentEntry = CurrentEntry;
            }

            return;
        }

        ClearLoadoutCategoryButtons();
        ClearupEdit();

        EntryBackButton.Visible = true;

        foreach (var (categoryId, loadouts) in _loadoutCache)
        {
            foreach (var loadoutPrototype in loadouts)
            {
                var displayName = _loadoutDisplayNames.GetValueOrDefault(loadoutPrototype.ID, loadoutPrototype.ID);

                if (!loadoutPrototype.ID.Contains(searchText, StringComparison.OrdinalIgnoreCase) &&
                    !displayName.Contains(searchText, StringComparison.OrdinalIgnoreCase))
                    continue;

                var loadoutEntry = new LoadoutEntry();

                if (_selectedLoadouts.TryGetValue(loadoutPrototype.ID, out var loadout))
                {
                    loadoutEntry.SetLoadout(loadout);
                    loadoutEntry.Selected = true;
                }
                else
                {
                    var newLoadout = new Loadout(
                        loadoutPrototype.ID,
                        loadoutPrototype.CustomName ? _customName.GetValueOrDefault(loadoutPrototype.ID) : null,
                        loadoutPrototype.CustomDescription ? _customDescription.GetValueOrDefault(loadoutPrototype.ID) : null,
                        loadoutPrototype.CustomContent ? _customContent.GetValueOrDefault(loadoutPrototype.ID) : null,
                        loadoutPrototype.CustomColorTint ? _customColorTints.GetValueOrDefault(loadoutPrototype.ID) : null,
                        loadoutPrototype.CanBeHeirloom ? _customHeirloom.GetValueOrDefault(loadoutPrototype.ID) : null
                    );
                    loadoutEntry.SetLoadout(newLoadout);
                }

                loadoutEntry.OnEditLoadoutRequired += OnEntryEditLoadoutRequired;
                loadoutEntry.OnLoadoutDirty += OnEntryLoadoutDirty;
                loadoutEntry.EnsureIsWearable(CharacterRequirementsArgs, LoadoutPoint);
                loadoutEntry.ShowUnusable = _showUnusable;

                _loadoutEntries.Add(loadoutEntry);
            }
        }

        SortAndPasteEntries();
    }

    private void ResetButtonPressed(BaseButton.ButtonEventArgs obj)
    {
        if (_checkpointRequirements is null)
            return;

        SetData(_checkpointLoadouts, _checkpointRequirements);
    }

    public void SetCheckpoint()
    {
        _checkpointLoadouts = _selectedLoadouts.Values.ToList();
        _checkpointRequirements = CharacterRequirementsArgs;
    }

    public bool IsCategoryValid(ProtoId<LoadoutCategoryPrototype> category) => _loadoutCache.ContainsKey(category);

    public void SetData(IEnumerable<Loadout> selectedPreferenceList, CharacterRequirementsArgs characterRequirements)
    {
        ClearLoadouts();
        ClearupEdit();

        var currentSearchText = LoadoutSearch.Text;

        LoadoutSearch.Clear();

        CharacterRequirementsArgs = characterRequirements;

        foreach (var preference in selectedPreferenceList)
        {
            if (!_prototypeManager.TryIndex<LoadoutPrototype>(preference.LoadoutName, out var proto))
            {
                // Remove this if not needed
                Logger.Error($"Cannot add loadout to selected loadouts: prototype {preference.LoadoutName} not found");
                continue;
            }

            var loadoutEntry = CreateEntry(preference.LoadoutName);

            loadoutEntry.Selected = true;
            _selectedLoadouts.Add(preference.LoadoutName, loadoutEntry.Loadout);

            loadoutEntry.EnsureIsWearable(CharacterRequirementsArgs, LoadoutPoint);

            if (loadoutEntry.CanWear)
                LoadoutPoint -= loadoutEntry.Cost;

            if (proto.CustomContent && loadoutEntry.Loadout.CustomContent != preference.CustomContent)
            {
                // Remove this if not needed
                Logger.Warning("CustomContent mismatch, syncing...");
                loadoutEntry.Loadout.CustomContent = preference.CustomContent;
                _customContent.SetValue(preference.LoadoutName, preference.CustomContent);
            }

            if (proto.CustomName && loadoutEntry.Loadout.CustomName != preference.CustomName)
            {
                // Remove this if not needed
                Logger.Warning("CustomName mismatch, syncing...");
                loadoutEntry.Loadout.CustomName = preference.CustomName;
                _customName.SetValue(preference.LoadoutName, preference.CustomName);
            }

            if (proto.CustomDescription && loadoutEntry.Loadout.CustomDescription != preference.CustomDescription)
            {
                // Remove this if not needed
                Logger.Warning("CustomDescription mismatch, syncing...");
                loadoutEntry.Loadout.CustomDescription = preference.CustomDescription;
                _customDescription.SetValue(preference.LoadoutName, preference.CustomDescription);
            }

            if (proto.CanBeHeirloom && loadoutEntry.Loadout.CustomHeirloom != preference.CustomHeirloom)
            {
                // Remove this if not needed
                Logger.Warning("CustomHeirloom mismatch, syncing...");
                loadoutEntry.Loadout.CustomHeirloom = preference.CustomHeirloom;

                if (preference.CustomHeirloom is null)
                    _customHeirloom.RemoveValue(preference.LoadoutName);
                else
                    _customHeirloom.SetValue(preference.LoadoutName, preference.CustomHeirloom.Value);
            }

            if (proto.CustomColorTint && loadoutEntry.Loadout.CustomColorTint != preference.CustomColorTint)
            {
                // Remove this if not needed
                Logger.Warning("CustomColorTint mismatch, syncing...");
                loadoutEntry.Loadout.CustomColorTint = preference.CustomColorTint;
                _customColorTints.SetValue(preference.LoadoutName, preference.CustomColorTint);
            }
        }

        CurrentEntry = CurrentEntry;
        if (!string.IsNullOrEmpty(currentSearchText))
            LoadoutSearch.Text = currentSearchText;
    }

    public bool LoadCategoryButtons(ProtoId<LoadoutCategoryPrototype> loadoutCategoryPrototype)
    {
        ClearLoadoutCategoryButtons();
        ClearupEdit();

        var loadoutPrototypes = GroupLoadoutsByGroup(loadoutCategoryPrototype).ToList();
        if (loadoutPrototypes.Count == 0)
            return false;

        foreach (var loadoutPrototype in loadoutPrototypes)
        {
            if (!string.IsNullOrEmpty(LoadoutSearch.Text))
            {
                var displayName = _loadoutDisplayNames.GetValueOrDefault(loadoutPrototype.ID, loadoutPrototype.ID);

                if (!loadoutPrototype.ID.Contains(LoadoutSearch.Text, StringComparison.OrdinalIgnoreCase) &&
                    !displayName.Contains(LoadoutSearch.Text, StringComparison.OrdinalIgnoreCase))
                    continue;
            }

            var loadoutEntry = new LoadoutEntry();

            if (_selectedLoadouts.TryGetValue(loadoutPrototype.ID, out var loadout))
            {
                loadoutEntry.SetLoadout(loadout);
                loadoutEntry.Selected = true;
            }
            else
            {
                var newLoadout = new Loadout(
                    loadoutPrototype.ID,
                    loadoutPrototype.CustomName ? _customName.GetValueOrDefault(loadoutPrototype.ID) : null,
                    loadoutPrototype.CustomDescription ? _customDescription.GetValueOrDefault(loadoutPrototype.ID) : null,
                    loadoutPrototype.CustomContent ? _customContent.GetValueOrDefault(loadoutPrototype.ID) : null,
                    loadoutPrototype.CustomColorTint ? _customColorTints.GetValueOrDefault(loadoutPrototype.ID) : null,
                    loadoutPrototype.CanBeHeirloom ? _customHeirloom.GetValueOrDefault(loadoutPrototype.ID) : null
                );
                loadoutEntry.SetLoadout(newLoadout);
            }

            loadoutEntry.OnEditLoadoutRequired += OnEntryEditLoadoutRequired;
            loadoutEntry.OnLoadoutDirty += OnEntryLoadoutDirty;
            loadoutEntry.EnsureIsWearable(CharacterRequirementsArgs, LoadoutPoint);
            loadoutEntry.ShowUnusable = _showUnusable;

            _loadoutEntries.Add(loadoutEntry);
        }

        SortAndPasteEntries();

        _selectedLoadoutCategory = loadoutCategoryPrototype;
        return true;
    }

    private void Dirty()
    {
        OnLoadoutsChanged?.Invoke(_selectedLoadouts.Select(x => x.Value).ToList());
    }

    private LoadoutEntry CreateEntry(string loadoutName)
    {
        if (!_prototypeManager.TryIndex<LoadoutPrototype>(loadoutName, out var prototype))
        {
            throw new Exception("Could not find a prototype " + loadoutName);
        }

        var loadout = new Loadout(
            loadoutName,
            prototype.CustomName ? _customName.GetValueOrDefault(loadoutName) : null,
            prototype.CustomDescription ? _customDescription.GetValueOrDefault(loadoutName) : null,
            prototype.CustomContent ? _customContent.GetValueOrDefault(loadoutName) : null,
            prototype.CustomColorTint ? _customColorTints.GetValueOrDefault(loadoutName) : null,
            prototype.CanBeHeirloom ? _customHeirloom.GetValueOrDefault(loadoutName) : null
        );

        var entry = new LoadoutEntry();
        entry.SetLoadout(loadout);
        return entry;
    }

    private void SpecialColorTintTogglePressed(BaseButton.ButtonEventArgs obj)
    {
        if (_currPrototype == null || !_currPrototype.CustomColorTint)
            return;

        ColorEdit.Visible = SpecialColorTintToggle.Pressed;
    }

    private void OnEntryLoadoutDirty(LoadoutEntry entry)
    {
        if (TryFreeLoadout(entry.Loadout))
        {
            entry.Selected = false;
            Dirty();
            return;
        }

        entry.Selected = TrySelectLoadout(entry);

        Dirty();
    }

    private bool TrySelectLoadout(LoadoutEntry entry)
    {
        if (entry.EnsureIsWearable(CharacterRequirementsArgs, LoadoutPoint))
        {
            _selectedLoadouts.Add(entry.Loadout.LoadoutName, entry.Loadout);
            return true;
        }
        return false;
    }

    private bool TryFreeLoadout(Loadout loadout)
    {
        if (!_selectedLoadouts.Remove(loadout.LoadoutName, out _))
            return false;

        if (_prototypeManager.TryIndex<LoadoutPrototype>(loadout.LoadoutName, out var prototype))
            LoadoutPoint += prototype.Cost;

        return true;
    }

    private bool _showUnusable;
    public bool ShowUnusable
    {
        get => _showUnusable;
        set
        {
            _showUnusable = value;
            foreach (var entry in _loadoutEntries)
                entry.ShowUnusable = value;
            UpdateCategoriesVisibility();
        }
    }

    public bool IsCategoryVisiblePublic(ProtoId<LoadoutCategoryPrototype> categoryId)
    {
        return IsCategoryVisible(categoryId);
    }

    private void OnEntryEditLoadoutRequired(LoadoutEntry entry)
    {
        EditLoadout(entry);
    }

    private LoadoutPrototype? _currPrototype;
    private LoadoutEntry? _currEdit;

    public void EditLoadout(LoadoutEntry loadoutEntry)
    {
        ClearupEdit();
        var loadout = loadoutEntry.Loadout;
        if (!_prototypeManager.TryIndex<LoadoutPrototype>(loadout.LoadoutName, out var loadoutPrototype))
        {
            Logger.Error($"Unable to load loadout: unknown prototype {loadout.LoadoutName}");
            return;
        }

        if (loadoutPrototype.CustomName)
        {
            NameEdit.Text = loadout.CustomName ?? "";
            if (loadoutPrototype.Items.Any())
            {
                var tempEntity = _entityManager.SpawnEntity(loadoutPrototype.Items.First(), MapCoordinates.Nullspace);
                var itemName = _entityManager.GetComponent<MetaDataComponent>(tempEntity).EntityName;
                NameEdit.PlaceHolder = itemName;
                _entityManager.DeleteEntity(tempEntity);
            }
        }

        if (loadoutPrototype.CustomDescription)
        {
            DescriptionEdit.TextRope = new Rope.Leaf(loadout.CustomDescription ?? "");
            if (loadoutPrototype.Items.Any())
            {
                var tempEntity = _entityManager.SpawnEntity(loadoutPrototype.Items.First(), MapCoordinates.Nullspace);
                var itemDesc = _entityManager.GetComponent<MetaDataComponent>(tempEntity).EntityDescription;
                DescriptionEdit.Placeholder = new Rope.Leaf(itemDesc);
                _entityManager.DeleteEntity(tempEntity);
            }
        }

        if (loadoutPrototype.CustomContent)
        {
            BookTextEdit.TextRope = new Rope.Leaf(loadout.CustomContent ?? "");
            BookTextEdit.Placeholder = new Rope.Leaf(Loc.GetString("humanoid-profile-editor-loadouts-customize-book-text-placeholder"));
        }
        LoadoutConfigContainer.Visible = true;

        SpecialName.Visible = loadoutPrototype.CustomName;
        SpecialDescription.Visible = loadoutPrototype.CustomDescription;
        SpecialBookText.Visible = loadoutPrototype.CustomContent;
        SpecialColorTintToggle.Visible = loadoutPrototype.CustomColorTint;

        if (loadoutPrototype.CustomName)
            NameEdit.Text = loadout.CustomName ?? "";

        if (loadoutPrototype.CustomDescription)
            DescriptionEdit.TextRope = new Rope.Leaf(loadout.CustomDescription ?? "");

        if (loadoutPrototype.CustomContent)
            BookTextEdit.TextRope = new Rope.Leaf(loadout.CustomContent ?? "");

        if (loadoutPrototype.CustomColorTint)
        {
            if (loadout.CustomColorTint is not null)
            {
                SpecialColorTintToggle.Pressed = true;
                ColorEdit.Color = Color.FromHex(loadout.CustomColorTint);
                ColorEdit.Visible = true;
            }
            else
            {
                SpecialColorTintToggle.Pressed = false;
                ColorEdit.Visible = false;
            }
        }

        _currEdit = loadoutEntry;
        _currPrototype = loadoutPrototype;
    }

    public void ClearCustomValues()
    {
        _customName.Clear();
        _customDescription.Clear();
        _customContent.Clear();
        _customColorTints.Clear();
        _customHeirloom.Clear();
    }

    private void ClearLoadoutCategoryButtons()
    {
        foreach (var entry in _loadoutEntries)
        {
            entry.OnEditLoadoutRequired -= OnEntryEditLoadoutRequired;
            entry.OnLoadoutDirty -= OnEntryLoadoutDirty;
        }

        _loadoutEntries.Clear();
        Loadouts.Children.Clear();
    }

    private void SortAndPasteEntries()
    {
        Loadouts.Children.Clear();
        _loadoutEntries.Sort();
        foreach (var entry in _loadoutEntries)
        {
            Loadouts.Children.Add(entry);
        }
    }

    private void SaveButtonPressed(BaseButton.ButtonEventArgs obj)
    {
        if (_currEdit == null)
            return;

        var oldValue = _currEdit.Loadout;
        oldValue.CustomName = TextOrNull(NameEdit.Text);
        oldValue.CustomContent = TextOrNull(BookTextEdit.TextRope);
        oldValue.CustomDescription = TextOrNull(DescriptionEdit.TextRope);
        oldValue.CustomColorTint = _currPrototype?.CustomColorTint == true && SpecialColorTintToggle.Pressed
            ? ColorEdit.Color.ToHex()
            : null;

        if (oldValue.CustomName != null) _customName.SetValue(oldValue.LoadoutName, oldValue.CustomName);
        else _customName.RemoveValue(oldValue.LoadoutName);
        if (oldValue.CustomContent != null) _customContent.SetValue(oldValue.LoadoutName, oldValue.CustomContent);
        else _customContent.RemoveValue(oldValue.LoadoutName);
        if (oldValue.CustomDescription != null) _customDescription.SetValue(oldValue.LoadoutName, oldValue.CustomDescription);
        else _customDescription.RemoveValue(oldValue.LoadoutName);
        if (oldValue.CustomColorTint != null) _customColorTints.SetValue(oldValue.LoadoutName, oldValue.CustomColorTint);
        else _customColorTints.RemoveValue(oldValue.LoadoutName);

        ClearupEdit();
        Dirty();
    }

    private void ClearupEdit()
    {
        LoadoutConfigContainer.Visible = false;
        _currEdit = null;
        _currPrototype = null;

        NameEdit.Text = "";
        DescriptionEdit.TextRope = Rope.Leaf.Empty;
        BookTextEdit.TextRope = Rope.Leaf.Empty;
        SpecialColorTintToggle.Pressed = false;
        ColorEdit.Visible = false;
    }

    private string? TextOrNull(Rope.Node node)
    {
        var text = Rope.Collapse(node);
        return string.IsNullOrEmpty(text) ? null : text;
    }

    private string? TextOrNull(string text)
    {
        if (string.IsNullOrEmpty(text))
            return null;

        return text;
    }

    private void ClearLoadouts()
    {
        ClearLoadoutCategoryButtons();
        _selectedLoadouts.Clear();
        LoadoutPoint = MaxLoadoutPoints;
    }

    private IEnumerable<LoadoutPrototype> GroupLoadoutsByGroup(ProtoId<LoadoutCategoryPrototype> loadoutCategoryPrototype)
    {
        if (_loadoutCache.TryGetValue(loadoutCategoryPrototype, out var loadouts))
            return loadouts;
        return [];
    }
}

public sealed class CharacterRequirementsArgs(
    JobPrototype job,
    HumanoidCharacterProfile profile,
    IReadOnlyDictionary<string, TimeSpan> playTimes,
    bool whitelisted,
    int depth = 0,
    MindComponent? mind = null,
    IDependencyCollection? dependencies = null)
{
    public JobPrototype Job { get; set; } = job;
    public HumanoidCharacterProfile Profile { get; set; } = profile;
    public IReadOnlyDictionary<string, TimeSpan> PlayTimes { get; set; } = playTimes;
    public bool Whitelisted { get; set; } = whitelisted;
    public int Depth { get; set; } = depth;
    public MindComponent? Mind { get; set; } = mind;
    public IDependencyCollection? Dependencies { get; set; } = dependencies;

    public bool IsValid(CharacterRequirement requirement, IPrototype prototype, [NotNullWhen(false)] out string? reason)
    {
        Dependencies ??= IoCManager.Instance!;
        var valid = requirement.IsValid(
            Job,
            Profile,
            PlayTimes,
            Whitelisted,
            prototype,
            Dependencies.Resolve<IEntityManager>(),
            Dependencies.Resolve<IPrototypeManager>(),
            Dependencies.Resolve<IConfigurationManager>(),
            out reason,
            Depth,
            Mind);
        return requirement.Inverted ? !valid : valid;
    }
}