LightlessClient/LightlessSync/FileCache/TransientResourceManager.cs

using LightlessSync.API.Data.Enum;
using LightlessSync.LightlessConfiguration;
using LightlessSync.LightlessConfiguration.Configurations;
using LightlessSync.PlayerData.Data;
using LightlessSync.PlayerData.Handlers;
using LightlessSync.PlayerData.Factories;
using LightlessSync.Services;
using LightlessSync.Services.ActorTracking;
using LightlessSync.Services.Mediator;
using LightlessSync.Utils;
using Microsoft.Extensions.Logging;
using System.Collections.Concurrent;
using DalamudObjectKind = Dalamud.Game.ClientState.Objects.Enums.ObjectKind;

namespace LightlessSync.FileCache;

public sealed class TransientResourceManager : DisposableMediatorSubscriberBase
{
    private readonly object _cacheAdditionLock = new();
    private readonly HashSet<string> _cachedHandledPaths = new(StringComparer.Ordinal);
    private readonly TransientConfigService _configurationService;
    private readonly DalamudUtilService _dalamudUtil;
    private readonly ActorObjectService _actorObjectService;
    private readonly GameObjectHandlerFactory _gameObjectHandlerFactory;
    private readonly object _ownedHandlerLock = new();
    private readonly string[] _handledFileTypes = ["tmb", "pap", "avfx", "atex", "sklb", "eid", "phyb", "scd", "skp", "shpk", "kdb"];
    private readonly string[] _handledRecordingFileTypes = ["tex", "mdl", "mtrl"];
    private readonly HashSet<GameObjectHandler> _playerRelatedPointers = [];
    private readonly object _playerRelatedLock = new();
    private readonly ConcurrentDictionary<nint, GameObjectHandler> _playerRelatedByAddress = new();
    private readonly Dictionary<nint, GameObjectHandler> _ownedHandlers = new();
    private ConcurrentDictionary<nint, ObjectKind> _cachedFrameAddresses = new();
    private ConcurrentDictionary<ObjectKind, HashSet<string>>? _semiTransientResources = null;
    private uint _lastClassJobId = uint.MaxValue;
    public bool IsTransientRecording { get; private set; } = false;

    public TransientResourceManager(ILogger<TransientResourceManager> logger, TransientConfigService configurationService,
            DalamudUtilService dalamudUtil, LightlessMediator mediator, ActorObjectService actorObjectService, GameObjectHandlerFactory gameObjectHandlerFactory) : base(logger, mediator)
    {
        _configurationService = configurationService;
        _dalamudUtil = dalamudUtil;
        _actorObjectService = actorObjectService;
        _gameObjectHandlerFactory = gameObjectHandlerFactory;
        Mediator.Subscribe<PenumbraResourceLoadMessage>(this, Manager_PenumbraResourceLoadEvent);
        Mediator.Subscribe<ActorTrackedMessage>(this, msg => HandleActorTracked(msg.Descriptor));
        Mediator.Subscribe<ActorUntrackedMessage>(this, msg => HandleActorUntracked(msg.Descriptor));
        Mediator.Subscribe<PenumbraModSettingChangedMessage>(this, (_) => Manager_PenumbraModSettingChanged());
        Mediator.Subscribe<PriorityFrameworkUpdateMessage>(this, (_) => DalamudUtil_FrameworkUpdate());
        Mediator.Subscribe<GameObjectHandlerCreatedMessage>(this, (msg) =>
        {
            if (!msg.OwnedObject) return;
            lock (_playerRelatedLock)
            {
                _playerRelatedPointers.Add(msg.GameObjectHandler);
            }
        });
        Mediator.Subscribe<GameObjectHandlerDestroyedMessage>(this, (msg) =>
        {
            if (!msg.OwnedObject) return;
            lock (_playerRelatedLock)
            {
                _playerRelatedPointers.Remove(msg.GameObjectHandler);
            }
        });

        foreach (var descriptor in _actorObjectService.ObjectDescriptors)
        {
            HandleActorTracked(descriptor);
        }
    }

    private TransientConfig.TransientPlayerConfig PlayerConfig
    {
        get
        {
            if (!_configurationService.Current.TransientConfigs.TryGetValue(PlayerPersistentDataKey, out var transientConfig))
            {
                _configurationService.Current.TransientConfigs[PlayerPersistentDataKey] = transientConfig = new();
            }

            return transientConfig;
        }
    }

    private string PlayerPersistentDataKey => _dalamudUtil.GetPlayerName() + "_" + _dalamudUtil.GetHomeWorldId();
    private ConcurrentDictionary<ObjectKind, HashSet<string>> SemiTransientResources
    {
        get
        {
            if (_semiTransientResources == null)
            {
                _semiTransientResources = new();
                PlayerConfig.JobSpecificCache.TryGetValue(_dalamudUtil.ClassJobId, out var jobSpecificData);
                _semiTransientResources[ObjectKind.Player] = PlayerConfig.GlobalPersistentCache.Concat(jobSpecificData ?? [])
                    .ToHashSet(StringComparer.OrdinalIgnoreCase);
                PlayerConfig.JobSpecificPetCache.TryGetValue(_dalamudUtil.ClassJobId, out var petSpecificData);
                _semiTransientResources[ObjectKind.Pet] = new HashSet<string>(
                    petSpecificData ?? [],
                    StringComparer.OrdinalIgnoreCase);
            }

            return _semiTransientResources;
        }
    }
    private ConcurrentDictionary<ObjectKind, HashSet<string>> TransientResources { get; } = new();

    public void CleanUpSemiTransientResources(ObjectKind objectKind, List<FileReplacement>? fileReplacement = null)
    {
        if (!SemiTransientResources.TryGetValue(objectKind, out HashSet<string>? value))
            return;

        if (fileReplacement == null)
        {
            value.Clear();
            return;
        }

        int removedPaths = 0;
        foreach (var replacement in fileReplacement.Where(p => !p.HasFileReplacement).SelectMany(p => p.GamePaths).ToList())
        {
            removedPaths += PlayerConfig.RemovePath(replacement, objectKind);
            value.Remove(replacement);
        }

        if (removedPaths > 0)
        {
            Logger.LogTrace("Removed {amount} of SemiTransient paths during CleanUp, Saving from {name}", removedPaths, nameof(CleanUpSemiTransientResources));
            // force reload semi transient resources
            _configurationService.Save();
        }
    }

    public HashSet<string> GetSemiTransientResources(ObjectKind objectKind)
    {
        SemiTransientResources.TryGetValue(objectKind, out var result);

        return result ?? new HashSet<string>(StringComparer.OrdinalIgnoreCase);
    }

    public void PersistTransientResources(ObjectKind objectKind)
    {
        if (!SemiTransientResources.TryGetValue(objectKind, out HashSet<string>? semiTransientResources))
        {
            SemiTransientResources[objectKind] = semiTransientResources = new(StringComparer.OrdinalIgnoreCase);
        }

        if (!TransientResources.TryGetValue(objectKind, out var resources))
        {
            return;
        }

        List<string> transientResources;
        lock (resources)
        {
            transientResources = resources.ToList();
        }

        Logger.LogDebug("Persisting {count} transient resources", transientResources.Count);
        List<string> newlyAddedGamePaths;
        lock (semiTransientResources)
        {
            newlyAddedGamePaths = transientResources.Except(semiTransientResources, StringComparer.OrdinalIgnoreCase).ToList();
            foreach (var gamePath in transientResources)
            {
                semiTransientResources.Add(gamePath);
            }
        }

        bool saveConfig = false;
        if (objectKind == ObjectKind.Player && newlyAddedGamePaths.Count != 0)
        {
            saveConfig = true;
            foreach (var item in newlyAddedGamePaths.Where(f => !string.IsNullOrEmpty(f)))
            {
                PlayerConfig.AddOrElevate(_dalamudUtil.ClassJobId, item);
            }
        }
        else if (objectKind == ObjectKind.Pet && newlyAddedGamePaths.Count != 0)
        {
            saveConfig = true;

            if (!PlayerConfig.JobSpecificPetCache.TryGetValue(_dalamudUtil.ClassJobId, out var petPerma))
            {
                PlayerConfig.JobSpecificPetCache[_dalamudUtil.ClassJobId] = petPerma = [];
            }

            foreach (var item in newlyAddedGamePaths.Where(f => !string.IsNullOrEmpty(f)))
            {
                petPerma.Add(item);
            }
        }

        if (saveConfig)
        {
            Logger.LogTrace("Saving transient.json from {method}", nameof(PersistTransientResources));
            _configurationService.Save();
        }

        lock (resources)
        {
            resources.Clear();
        }
    }

    public void RemoveTransientResource(ObjectKind objectKind, string path)
    {
        var normalizedPath = NormalizeGamePath(path);
        if (SemiTransientResources.TryGetValue(objectKind, out var resources))
        {
            resources.Remove(normalizedPath);
            if (objectKind == ObjectKind.Player)
            {
                PlayerConfig.RemovePath(normalizedPath, objectKind);
                Logger.LogTrace("Saving transient.json from {method}", nameof(RemoveTransientResource));
                _configurationService.Save();
            }
        }
    }

    internal bool AddTransientResource(ObjectKind objectKind, string item)
    {
        var normalizedItem = NormalizeGamePath(item);
        if (SemiTransientResources.TryGetValue(objectKind, out var semiTransient) && semiTransient != null && semiTransient.Contains(normalizedItem))
            return false;

        if (!TransientResources.TryGetValue(objectKind, out HashSet<string>? transientResource))
        {
            transientResource = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
            TransientResources[objectKind] = transientResource;
        }

        return transientResource.Add(normalizedItem);
    }

    internal void ClearTransientPaths(ObjectKind objectKind, List<string> list)
    {
        // ignore all recording only datatypes
        int recordingOnlyRemoved = list.RemoveAll(entry => _handledRecordingFileTypes.Any(ext => entry.EndsWith(ext, StringComparison.OrdinalIgnoreCase)));
        if (recordingOnlyRemoved > 0)
        {
            Logger.LogTrace("Ignored {0} game paths when clearing transients", recordingOnlyRemoved);
        }

        if (TransientResources.TryGetValue(objectKind, out var set))
        {
            foreach (var file in set.Where(p => list.Contains(p, StringComparer.OrdinalIgnoreCase)))
            {
                Logger.LogTrace("Removing From Transient: {file}", file);
            }

            int removed = set.RemoveWhere(p => list.Contains(p, StringComparer.OrdinalIgnoreCase));
            Logger.LogDebug("Removed {removed} previously existing transient paths", removed);
        }

        bool reloadSemiTransient = false;
        if (objectKind == ObjectKind.Player && SemiTransientResources.TryGetValue(objectKind, out var semiset))
        {
            foreach (var file in semiset.Where(p => list.Contains(p, StringComparer.OrdinalIgnoreCase)))
            {
                Logger.LogTrace("Removing From SemiTransient: {file}", file);
                PlayerConfig.RemovePath(file, objectKind);
            }

            int removed = semiset.RemoveWhere(p => list.Contains(p, StringComparer.OrdinalIgnoreCase));
            Logger.LogDebug("Removed {removed} previously existing semi transient paths", removed);
            if (removed > 0)
            {
                reloadSemiTransient = true;
                Logger.LogTrace("Saving transient.json from {method}", nameof(ClearTransientPaths));
                _configurationService.Save();
            }
        }

        if (reloadSemiTransient)
            _semiTransientResources = null;
    }

    protected override void Dispose(bool disposing)
    {
        base.Dispose(disposing);

        TransientResources.Clear();
        SemiTransientResources.Clear();

        lock (_ownedHandlerLock)
        {
            foreach (var handler in _ownedHandlers.Values)
            {
                handler.Dispose();
            }
            _ownedHandlers.Clear();
        }
    }

    private void DalamudUtil_FrameworkUpdate()
    {
        _ = Task.Run(() => RefreshPlayerRelatedAddressMap());

        lock (_cacheAdditionLock)
        {
            _cachedHandledPaths.Clear();
        }

        if (_lastClassJobId != _dalamudUtil.ClassJobId)
        {
            UpdateClassJobCache();
        }

        CleanupAbsentObjects();
    }

    private void RefreshPlayerRelatedAddressMap()
    {
        var tempMap = new ConcurrentDictionary<nint, GameObjectHandler>();
        var updatedFrameAddresses = new ConcurrentDictionary<nint, ObjectKind>();

        lock (_playerRelatedLock)
        {
            foreach (var handler in _playerRelatedPointers)
            {
                var address = (nint)handler.Address;
                if (address != nint.Zero)
                {
                    tempMap[address] = handler;
                    updatedFrameAddresses[address] = handler.ObjectKind;
                }
            }
        }

        _playerRelatedByAddress.Clear();
        foreach (var kvp in tempMap)
        {
            _playerRelatedByAddress[kvp.Key] = kvp.Value;
        }

        _cachedFrameAddresses.Clear();
        foreach (var kvp in updatedFrameAddresses)
        {
            _cachedFrameAddresses[kvp.Key] = kvp.Value;
        }
    }

    private void UpdateClassJobCache()
    {
        _lastClassJobId = _dalamudUtil.ClassJobId;
        if (SemiTransientResources.TryGetValue(ObjectKind.Pet, out HashSet<string>? value))
        {
            value?.Clear();
        }

        PlayerConfig.JobSpecificCache.TryGetValue(_dalamudUtil.ClassJobId, out var jobSpecificData);
        SemiTransientResources[ObjectKind.Player] = PlayerConfig.GlobalPersistentCache
            .Concat(jobSpecificData ?? [])
            .ToHashSet(StringComparer.OrdinalIgnoreCase);

        PlayerConfig.JobSpecificPetCache.TryGetValue(_dalamudUtil.ClassJobId, out var petSpecificData);
        SemiTransientResources[ObjectKind.Pet] = new HashSet<string>(
            petSpecificData ?? [],
            StringComparer.OrdinalIgnoreCase);
    }

    private void CleanupAbsentObjects()
    {
        foreach (var kind in Enum.GetValues(typeof(ObjectKind)).Cast<ObjectKind>())
        {
            if (!_cachedFrameAddresses.Any(k => k.Value == kind) && TransientResources.Remove(kind, out _))
            {
                Logger.LogDebug("Object not present anymore: {kind}", kind.ToString());
            }
        }
    }

    private void Manager_PenumbraModSettingChanged()
    {
        _ = Task.Run(() =>
        {
            Logger.LogDebug("Penumbra Mod Settings changed, verifying SemiTransientResources");
            lock (_playerRelatedLock)
            {
                foreach (var item in _playerRelatedPointers)
                {
                    Mediator.Publish(new TransientResourceChangedMessage(item.Address));
                }
            }
        });
    }

    public void RebuildSemiTransientResources()
    {
        _semiTransientResources = null;
    }

    private void HandleActorTracked(ActorObjectService.ActorDescriptor descriptor)
    {
        if (descriptor.IsInGpose)
            return;

        if (descriptor.OwnedKind is not ObjectKind ownedKind)
            return;

        if (Logger.IsEnabled(LogLevel.Debug))
        {
            Logger.LogDebug("ActorObject tracked: {kind} addr={address:X} name={name}", ownedKind, descriptor.Address, descriptor.Name);
        }

        _cachedFrameAddresses[descriptor.Address] = ownedKind;

        lock (_ownedHandlerLock)
        {
            if (_ownedHandlers.ContainsKey(descriptor.Address))
                return;

            _ = CreateOwnedHandlerAsync(descriptor, ownedKind);
        }
    }

    private void HandleActorUntracked(ActorObjectService.ActorDescriptor descriptor)
    {
        if (Logger.IsEnabled(LogLevel.Debug))
        {
            var kindLabel = descriptor.OwnedKind?.ToString()
                ?? (descriptor.ObjectKind == DalamudObjectKind.Player ? ObjectKind.Player.ToString() : "<none>");
            Logger.LogDebug("ActorObject untracked: addr={address:X} name={name} kind={kind}", descriptor.Address, descriptor.Name, kindLabel);
        }

        _cachedFrameAddresses.TryRemove(descriptor.Address, out _);

        if (descriptor.OwnedKind is not ObjectKind)
            return;

        lock (_ownedHandlerLock)
        {
            if (_ownedHandlers.Remove(descriptor.Address, out var handler))
            {
                handler.Dispose();
            }
        }
    }

    private async Task CreateOwnedHandlerAsync(ActorObjectService.ActorDescriptor descriptor, ObjectKind kind)
    {
        try
        {
            var handler = await _gameObjectHandlerFactory.Create(
                kind,
                () =>
                {
                    if (!string.IsNullOrEmpty(descriptor.HashedContentId) &&
                        _actorObjectService.TryGetValidatedActorByHash(descriptor.HashedContentId, out var current) &&
                        current.OwnedKind == kind)
                    {
                        return current.Address;
                    }

                    return descriptor.Address;
                },
                true).ConfigureAwait(false);

            if (handler.Address == IntPtr.Zero)
            {
                handler.Dispose();
                return;
            }

            lock (_ownedHandlerLock)
            {
                if (!_cachedFrameAddresses.ContainsKey(descriptor.Address))
                {
                    Logger.LogDebug("ActorObject handler discarded (stale): addr={address:X}", descriptor.Address);
                    handler.Dispose();
                    return;
                }

                _ownedHandlers[descriptor.Address] = handler;
            }

            Logger.LogDebug("ActorObject handler created: {kind} addr={address:X}", kind, descriptor.Address);
        }
        catch (Exception ex)
        {
            Logger.LogError(ex, "Failed to create owned handler for {kind} at {address:X}", kind, descriptor.Address);
        }
    }

    private static string NormalizeGamePath(string path)
    {
        if (string.IsNullOrEmpty(path))
            return string.Empty;

        return path.Replace("\\", "/", StringComparison.Ordinal).ToLowerInvariant();
    }

    private static string NormalizeFilePath(string path)
    {
        if (string.IsNullOrEmpty(path))
            return string.Empty;

        if (path.StartsWith("|", StringComparison.Ordinal))
        {
            var lastPipe = path.LastIndexOf('|');
            if (lastPipe >= 0 && lastPipe + 1 < path.Length)
            {
                path = path[(lastPipe + 1)..];
            }
        }

        return NormalizeGamePath(path);
    }

    private static bool HasHandledFileType(string gamePath, string[] handledTypes)
    {
        for (var i = 0; i < handledTypes.Length; i++)
        {
            if (gamePath.EndsWith(handledTypes[i], StringComparison.Ordinal))
                return true;
        }

        return false;
    }

    private void Manager_PenumbraResourceLoadEvent(PenumbraResourceLoadMessage msg)
    {
        var gamePath = msg.GamePath.ToLowerInvariant();
        var gameObjectAddress = msg.GameObject;
        var filePath = msg.FilePath;

        // ignore files already processed this frame
        if (_cachedHandledPaths.Contains(gamePath)) return;

        lock (_cacheAdditionLock)
        {
            _cachedHandledPaths.Add(gamePath);
        }

        // replace individual mtrl stuff
        if (filePath.StartsWith("|", StringComparison.OrdinalIgnoreCase))
        {
            filePath = filePath.Split("|")[2];
        }
        // replace filepath
        filePath = filePath.ToLowerInvariant().Replace("\\", "/", StringComparison.OrdinalIgnoreCase);

        // ignore files that are the same
        var replacedGamePath = gamePath.ToLowerInvariant().Replace("\\", "/", StringComparison.OrdinalIgnoreCase);
        if (string.Equals(filePath, replacedGamePath, StringComparison.OrdinalIgnoreCase))
        {
            return;
        }

        // ignore files to not handle
        var handledTypes = IsTransientRecording ? _handledRecordingFileTypes.Concat(_handledFileTypes) : _handledFileTypes;
        if (!handledTypes.Any(type => gamePath.EndsWith(type, StringComparison.OrdinalIgnoreCase)))
        {
            lock (_cacheAdditionLock)
            {
                _cachedHandledPaths.Add(gamePath);
            }
            return;
        }

        // ignore files not belonging to anything player related
        if (!_cachedFrameAddresses.TryGetValue(gameObjectAddress, out var objectKind))
        {
            lock (_cacheAdditionLock)
            {
                _cachedHandledPaths.Add(gamePath);
            }
            return;
        }

        // ^ all of the code above is just to sanitize the data

        if (!TransientResources.TryGetValue(objectKind, out HashSet<string>? transientResources))
        {
            transientResources = new(StringComparer.OrdinalIgnoreCase);
            TransientResources[objectKind] = transientResources;
        }

        _playerRelatedByAddress.TryGetValue(gameObjectAddress, out var owner);
        bool alreadyTransient = false;

        bool transientContains = transientResources.Contains(replacedGamePath);
        bool semiTransientContains = SemiTransientResources.SelectMany(k => k.Value)
            .Any(f => string.Equals(f, gamePath, StringComparison.OrdinalIgnoreCase));
        if (transientContains || semiTransientContains)
        {
            if (!IsTransientRecording)
                Logger.LogTrace("Not adding {replacedPath} => {filePath}, Reason: Transient: {contains}, SemiTransient: {contains2}", replacedGamePath, filePath,
                    transientContains, semiTransientContains);
            alreadyTransient = true;
        }
        else
        {
            if (!IsTransientRecording)
            {
                bool isAdded = transientResources.Add(replacedGamePath);
                if (isAdded)
                {
                    Logger.LogDebug("Adding {replacedGamePath} for {gameObject} ({filePath})", replacedGamePath, owner?.ToString() ?? gameObjectAddress.ToString("X"), filePath);
                    SendTransients(gameObjectAddress, objectKind);
                }
            }
        }

        if (owner != null && IsTransientRecording)
        {
            _recordedTransients.Add(new TransientRecord(owner, replacedGamePath, filePath, alreadyTransient) { AddTransient = !alreadyTransient });
        }
    }

    private void SendTransients(nint gameObject, ObjectKind objectKind)
    {
        _sendTransientCts.Cancel();
        _sendTransientCts = new();
        var token = _sendTransientCts.Token;

        _ = Task.Run(async () =>
        {
            try
            {
                await Task.Delay(TimeSpan.FromSeconds(5), token).ConfigureAwait(false);

                if (TransientResources.TryGetValue(objectKind, out var values) && values.Any())
                {
                    Logger.LogTrace("Sending Transients for {kind}", objectKind);
                    Mediator.Publish(new TransientResourceChangedMessage(gameObject));
                }
            }
            catch (TaskCanceledException)
            {

            }
            catch (System.OperationCanceledException)
            {

            }
        });
    }

    public void StartRecording(CancellationToken token)
    {
        if (IsTransientRecording) return;
        _recordedTransients.Clear();
        IsTransientRecording = true;
        RecordTimeRemaining.Value = TimeSpan.FromSeconds(150);
        _ = Task.Run(async () =>
        {
            try
            {
                while (RecordTimeRemaining.Value > TimeSpan.Zero && !token.IsCancellationRequested)
                {
                    await Task.Delay(TimeSpan.FromSeconds(1), token).ConfigureAwait(false);
                    RecordTimeRemaining.Value = RecordTimeRemaining.Value.Subtract(TimeSpan.FromSeconds(1));
                }
            }
            finally
            {
                IsTransientRecording = false;
            }
        });
    }

    public async Task WaitForRecording(CancellationToken token)
    {
        while (IsTransientRecording)
        {
            await Task.Delay(TimeSpan.FromSeconds(1), token).ConfigureAwait(false);
        }
    }

    internal void SaveRecording()
    {
        HashSet<nint> addedTransients = [];
        foreach (var item in _recordedTransients)
        {
            if (!item.AddTransient || item.AlreadyTransient) continue;
            if (!TransientResources.TryGetValue(item.Owner.ObjectKind, out var transient))
            {
                TransientResources[item.Owner.ObjectKind] = transient = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
            }

            Logger.LogTrace("Adding recorded: {gamePath} => {filePath}", item.GamePath, item.FilePath);

            transient.Add(item.GamePath);
            addedTransients.Add(item.Owner.Address);
        }

        _recordedTransients.Clear();

        foreach (var item in addedTransients)
        {
            Mediator.Publish(new TransientResourceChangedMessage(item));
        }
    }

    private readonly HashSet<TransientRecord> _recordedTransients = [];
    public IReadOnlySet<TransientRecord> RecordedTransients => _recordedTransients;

    public ValueProgress<TimeSpan> RecordTimeRemaining { get; } = new();
    private CancellationTokenSource _sendTransientCts = new();

    public record TransientRecord(GameObjectHandler Owner, string GamePath, string FilePath, bool AlreadyTransient)
    {
        public bool AddTransient { get; set; }
    }
}