XFEWarframe

using System.Runtime.InteropServices;
using System.Text.Json.Nodes;
using XFEWarframe.Client;

namespace XFEWarframe.Bridge;

internal static unsafe partial class ExperimentalAbilityRuntime
{
    private sealed class BankState(Plan plan)
    {
        public Plan Plan { get; } = plan;
        public int[] Selected { get; } = new int[4];
        public bool[] Pending { get; } = new bool[4];
        // Each candidate owns a strong reference. Selection is independent of the
        // lifetime of a buff/channel and never destroys a previously cast ability.
        public nint[][] References { get; } = plan.Banks!.Select(b => new nint[b.Length]).ToArray();
        public bool[] ForwardedPress { get; } = new bool[4];
        public AbilityBankGesture[] Gestures { get; } = Enumerable.Range(0, 4).Select(_ => new AbilityBankGesture()).ToArray();
    }
    private static readonly Dictionary<nint, BankState> BankStates = [];
    private static nint input, updateSuit, activate, deactivate, abilityUpdate, referenceCopy, emptyReference;

    internal static JsonObject Snapshot(nint suit)
    {
        var slots = new JsonArray(); var retained = new JsonArray(); var state = State(suit);
        for (var slot = 0; slot < 4; slot++)
        {
            var ability = Instance(suit, slot); if (ability == 0) continue;
            var scriptReference = *(nint*)(ability + 0x570);
            slots.Add((JsonNode)new JsonObject { ["slot"] = slot, ["instance"] = $"0x{ability:X}",
                ["active"] = *(byte*)(ability + 0x678) != 0, ["rank"] = *(byte*)(ability + 0x690),
                ["selected"] = state?.Selected[slot] ?? 0, ["count"] = state?.Plan.Banks?[slot].Length ?? 1,
                ["script"] = $"0x{(scriptReference == 0 ? 0 : *(nint*)scriptReference):X}",
                ["scriptType"] = $"0x{*(nint*)(ability + 0x578):X}" });
        }
        if (state is not null)
            for (var slot = 0; slot < 4; slot++)
                for (var choice = 0; choice < state.References[slot].Length; choice++)
                {
                    var reference = state.References[slot][choice]; var ability = reference == 0 ? 0 : *(nint*)reference;
                    if (ability != 0) retained.Add((JsonNode)new JsonObject { ["slot"] = slot, ["choice"] = choice,
                        ["instance"] = $"0x{ability:X}", ["active"] = *(byte*)(ability + 0x678) != 0,
                        ["type"] = TypePath(*(nint*)(ability + 8)), ["selected"] = choice == state.Selected[slot] });
                }
        var avatar = Owner(suit); var controller = avatar == 0 ? 0 : AnimationController(avatar);
        return new JsonObject { ["suit"] = $"0x{suit:X}", ["slots"] = slots, ["retained"] = retained,
            ["hud"] = HudSnapshot(),
            ["infestationKing"] = KingSnapshot(),
            ["branding"] = GameBranding.Snapshot(),
            ["soundEvents"] = GameTestRuntime.SoundSnapshot(),
            ["animationController"] = controller == 0 ? null : TypePath(*(nint*)(controller + 8)),
            ["animationInstance"] = $"0x{controller:X}", ["donorAnimationActive"] = AnimationLeases.ContainsKey(suit) };
    }

    private static void PrepareBanks(Func<string, nint> address, Func<string, nint, nint> hook, List<NativeHook> hooks)
    {
        activate = address("TryActivate"); deactivate = address("Deactivate");
        abilityUpdate = address("AbilityUpdate"); referenceCopy = address("StrongReferenceCopy");
        emptyReference = activate - ExperimentalAbilityCompatibility.Entries.Single(e => e.Name == "TryActivate").Rva + 0x272F3B0;
        input = hook("SuitInput", (nint)(delegate* unmanaged<nint, int, byte, byte>)&Input);
        updateSuit = hook("SuitUpdate", (nint)(delegate* unmanaged<nint, float, void>)&UpdateSuit);
        // This call appends the replaced ability's mSecondaryScripts and mAbilityEvents.
        // The engine already checks that the types differ before reaching this call.
        // Keep the new ability's own scripts; host-slot resource checks must not migrate.
        hooks.Add(NativeHook.RedirectCall(address("OverrideInheritScripts"),
            (nint)(delegate* unmanaged<nint, nint, void>)&KeepDonorScripts));
    }
    [UnmanagedCallersOnly] private static void KeepDonorScripts(nint ability, nint replacedType) { }

    private static bool BeginObject(nint reader, string name, bool reading)
    {
        Name(reader, name);
        ((delegate* unmanaged<nint, byte, void>)Virtual(reader, 0x158))(reader, 12);
        if (reading && ((delegate* unmanaged<nint, byte>)Virtual(reader, 0x48))(reader) == 0) return false;
        ((delegate* unmanaged<nint, void>)Virtual(reader, 0x58))(reader); return true;
    }
    private static void EndObject(nint reader) => ((delegate* unmanaged<nint, void>)Virtual(reader, 0x60))(reader);
    private static string Key(int value) => value.ToString(System.Globalization.CultureInfo.InvariantCulture);
    private static nint[][]? ReadBanks(nint reader)
    {
        if (!BeginObject(reader, "XFEAbilityBanks", true)) return null;
        var banks = Enumerable.Range(0, 4).Select(_ => Array.Empty<nint>()).ToArray();
        try
        {
            for (var slot = 0; slot < 4; slot++)
            {
                if (!BeginObject(reader, Key(slot), true)) continue;
                var choices = new List<nint>(3);
                try
                {
                    var gap = false;
                    for (var index = 0; index < 3; index++)
                    {
                        Name(reader, Key(index)); nint ability = 0;
                        ((delegate* unmanaged<nint*, nint, nint, void>)typeSerialize)(&ability, reader, typeBase);
                        if (ability == 0) { gap = true; continue; }
                        if (gap || choices.Contains(ability) || *(nint*)((delegate* unmanaged<nint*, nint>)validate)(&ability) == 0)
                            throw new InvalidDataException("Invalid experimental ability bank.");
                        choices.Add(ability);
                    }
                }
                finally { EndObject(reader); }
                banks[slot] = choices.ToArray();
            }
        }
        finally { EndObject(reader); }
        return banks;
    }
    private static void WriteBanks(nint reader, nint[][] banks)
    {
        BeginObject(reader, "XFEAbilityBanks", false);
        try
        {
            for (var slot = 0; slot < 4; slot++)
            {
                if (banks[slot].Length == 0) continue;
                BeginObject(reader, Key(slot), false);
                try
                {
                    for (var index = 0; index < banks[slot].Length; index++)
                    {
                        Name(reader, Key(index)); var ability = banks[slot][index];
                        ((delegate* unmanaged<nint*, nint, nint, void>)typeSerialize)(&ability, reader, typeBase);
                    }
                }
                finally { EndObject(reader); }
            }
        }
        finally { EndObject(reader); }
    }
    private static nint Instance(nint suit, int slot)
    {
        var instances = *(nint*)(suit + 0x3448); var bytes = *(uint*)(suit + 0x3450);
        if (instances == 0 || slot < 0 || slot >= bytes / 8) return 0;
        var reference = *(nint*)(instances + slot * 8);
        return reference == 0 ? 0 : *(nint*)reference;
    }
    private static BankState? State(nint suit)
    {
        lock (Gate)
        {
            if (BankStates.TryGetValue(suit, out var existing)) return existing;
            if (!Suits.TryGetValue(suit, out var plan) || plan.Banks is null) return null;
            var created = new BankState(plan); BankStates[suit] = created; return created;
        }
    }
    [UnmanagedCallersOnly] private static byte Input(nint suit, int code, byte pressed)
    {
        if (code is >= 0x91 and <= 0x94 && !GameTestRuntime.AcceptAbilityInput(suit)) return 1;
        try
        {
            var slot = code - 0x91; var state = slot is >= 0 and < 4 ? State(suit) : null;
            if (state?.Plan.Banks?[slot].Length > 1 && Instance(suit, slot) != 0)
            {
                if (pressed != 0) state.Gestures[slot].Press();
                else if (state.Gestures[slot].Release() && !state.Pending[slot])
                {
                    state.Pending[slot] = true;
                    CycleWhenReady(suit, state, slot);
                }
                else if (state.ForwardedPress[slot])
                {
                    state.ForwardedPress[slot] = false;
                    ((delegate* unmanaged<nint, int, byte, byte>)input)(suit, code, 0);
                }
                return 1;
            }
        }
        catch (Exception ex) { log("Ability bank input: " + ex.GetType().Name); }
        return ((delegate* unmanaged<nint, int, byte, byte>)input)(suit, code, pressed);
    }
    [UnmanagedCallersOnly] private static void UpdateSuit(nint suit, float seconds)
    {
        ((delegate* unmanaged<nint, float, void>)updateSuit)(suit, seconds);
        try
        {
            AdvanceDependencyStartup(suit, seconds);
            RestoreDonorAnimations(suit);
            BankState? state; lock (Gate) BankStates.TryGetValue(suit, out state);
            if (state is null) return;
            for (var slot = 0; slot < 4; slot++)
            {
                UpdateRetained(suit, state, slot, seconds);
                CycleWhenReady(suit, state, slot);
                if (state.Gestures[slot].Tick(seconds) && !state.Pending[slot])
                {
                    // The native activation path still checks energy, rank, target, silence,
                    // cooldown and the new ability's own evaluation script.
                    state.ForwardedPress[slot] = true;
                    ((delegate* unmanaged<nint, int, byte, byte>)input)(suit, slot + 0x91, 1);
                    if (trace) log($"Ability bank: held slot {slot + 1}, choice {state.Selected[slot] + 1}.");
                }
            }
        }
        catch (Exception ex) { log("Ability bank update: " + ex.GetType().Name); }
    }
    private static void CycleWhenReady(nint suit, BankState state, int slot)
    {
        if (!state.Pending[slot]) return;
        var previous = Instance(suit, slot);
        if (previous == 0 || DependencyStartupPending(suit, previous)) return;
        var oldRank = *(byte*)(previous + 0x690);
        var bank = state.Plan.Banks![slot]; var next = (state.Selected[slot] + 1) % bank.Length;
        RetainCurrent(suit, state, slot);
        var saved = state.References[slot][next];
        if (saved == 0)
        {
            ApplyOne(suit, bank[next], (byte)slot);
            state.Selected[slot] = next;
            RetainCurrent(suit, state, slot);
        }
        else
        {
            var destination = *(nint*)(suit + 0x3448) + slot * 8;
            ((delegate* unmanaged<nint, nint*, void>)referenceCopy)(destination, &saved);
            state.Selected[slot] = next;
        }
        var current = Instance(suit, slot);
        if (current == 0) return;
        *(byte*)(current + 0x690) = oldRank;
        state.Pending[slot] = false;
        if (trace) log($"Ability bank: cycled slot {slot + 1} to choice {next + 1}/{bank.Length}.");
    }

    private static void RetainCurrent(nint suit, BankState state, int slot)
    {
        var choice = state.Selected[slot];
        if (state.References[slot].Length == 0 || state.References[slot][choice] != 0) return;
        var reference = *(nint*)(*(nint*)(suit + 0x3448) + slot * 8);
        ++*(int*)(reference + 8); ++*(int*)(reference + 12);
        state.References[slot][choice] = reference;
    }

    private static void UpdateRetained(nint suit, BankState state, int slot, float seconds)
    {
        // The original suit update has already ticked the selected candidate.
        // Temporarily expose each hidden candidate in its real slot while ticking it;
        // scripts that query their owner see their own instance during the call.
        var destination = *(nint*)(suit + 0x3448) + slot * 8;
        var selected = *(nint*)destination;
        try
        {
            foreach (var reference in state.References[slot])
            {
                if (reference == 0 || reference == selected || *(nint*)reference == 0) continue;
                *(nint*)destination = reference;
                ((delegate* unmanaged<nint, float, void>)abilityUpdate)(*(nint*)reference, seconds);
            }
        }
        finally { *(nint*)destination = selected; }
    }

    private static bool HasAbility(nint suit, nint ability) => Enumerable.Range(0, 4).Any(slot => Instance(suit, slot) == ability) ||
        BankStates.TryGetValue(suit, out var bank) && bank.References.Any(b => b.Any(r => r != 0 && *(nint*)r == ability));

    private static void ReleaseBanks(nint suit)
    {
        if (!BankStates.Remove(suit, out var state)) return;
        for (var slot = 0; slot < 4; slot++)
        {
            state.Gestures[slot].Cancel();
            if (state.ForwardedPress[slot]) ((delegate* unmanaged<nint, int, byte, byte>)input)(suit, slot + 0x91, 0);
            var destination = *(nint*)(suit + 0x3448) + slot * 8;
            var selected = *(nint*)destination;
            try
            {
                foreach (var reference in state.References[slot])
                    if (reference != 0 && *(nint*)reference != 0 && *(byte*)(*(nint*)reference + 0x678) != 0)
                    {
                        *(nint*)destination = reference;
                        ((delegate* unmanaged<nint, byte, void>)deactivate)(suit, (byte)slot);
                    }
            }
            finally { *(nint*)destination = selected; }
            for (var choice = 0; choice < state.References[slot].Length; choice++)
            {
                var reference = state.References[slot][choice];
                if (reference == 0) continue;
                var empty = emptyReference;
                ((delegate* unmanaged<nint*, nint*, void>)referenceCopy)(&reference, &empty);
                // Drop the sentinel reference acquired by native assignment.
                --*(int*)(emptyReference + 8); --*(int*)(emptyReference + 12);
                state.References[slot][choice] = 0;
            }
        }
    }
}