using Facepunch;
using HarmonyLib;
using Newtonsoft.Json;
using Oxide.Core.Plugins;
using Rust.Ai.Gen2;
using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
using System.Text;
using UnityEngine;
using static Rust.Ai.Gen2.FSMComponent;

namespace Oxide.Plugins
{
    [Info("Remove Animals AI", "Whipers88 @CobaltStudios", "1.2.0")]
    [Description("Removing AI only for animals (not for bots)")]
    public class RemoveAnimalsAI : RustPlugin
    {
        static RemoveAnimalsAI _plugin;
        static HashSet<string> _frozenAnimals;
        static bool _allowUnlisted = true;

        #region Config

        private PluginConfig _config;

        private class PluginConfig
        {
            [JsonProperty("Animals that may not move (prefab short names)", ObjectCreationHandling = ObjectCreationHandling.Replace)]
            public List<string> FrozenAnimals = new List<string>
            {
                "bear", "polarbear", "wolf", "wolf2", "boar", "stag", "chicken",
                "tiger", "panther", "crocodile", "snake"
            };

            [JsonProperty("Let animals not in the list move (e.g. livestock and critters from the October 2026 update)")]
            public bool AllowUnlisted = true;
        }

        protected override void LoadDefaultConfig() => _config = new PluginConfig();

        protected override void LoadConfig()
        {
            base.LoadConfig();
            try
            {
                _config = Config.ReadObject<PluginConfig>();
                if (_config == null || _config.FrozenAnimals == null)
                    throw new Exception();
            }
            catch
            {
                PrintError("Config file is invalid, using default values");
                LoadDefaultConfig();
            }
            SaveConfig();
        }

        protected override void SaveConfig() => Config.WriteObject(_config);

        // Unknown entity (null) counts as frozen, which is how the plugin behaved before the config existed
        static bool IsFrozen(BaseEntity entity)
        {
            if (entity == null)
                return true;
            if (entity is ScientistNPC2)
                return false;
            if (_frozenAnimals == null || _frozenAnimals.Contains(entity.ShortPrefabName))
                return true;
            return !_allowUnlisted;
        }

        #endregion

        private void Init()
        {
            _plugin = this;
            _frozenAnimals = new HashSet<string>(_config.FrozenAnimals.Select(n => n.Trim()), StringComparer.OrdinalIgnoreCase);
            _allowUnlisted = _config.AllowUnlisted;

            // workQueue is no longer static (October 2026 patch), so each FSMComponent is pulled from its own queue
            List<FSMComponent> FSMComponents = Pool.Get<List<FSMComponent>>();
            List<BaseEntity> toKill = Pool.Get<List<BaseEntity>>();
            foreach (var networkable in BaseNetworkable.serverEntities)
            {
                var entity = networkable as BaseEntity;
                if (entity == null || !(entity is BaseNPC2 || entity is BaseAnimalNPC || entity is BaseNpc))
                    continue;

                if (!IsFrozen(entity))
                {
                    // Animals that are allowed to move may have been held back by an earlier version of the plugin
                    if (entity is BaseNPC2 && !(entity is ScientistNPC2) && entity.TryGetComponent<FSMComponent>(out var activeFsm) && !activeFsm.workQueue.Contains(activeFsm))
                        activeFsm.SetFsmActive(true);
                    continue;
                }

                if (entity is BaseNPC2)
                {
                    if (entity.TryGetComponent<FSMComponent>(out var fsmComponent))
                        FSMComponents.Add(fsmComponent);
                    toKill.Add(entity);
                }
                if (entity is BaseAnimalNPC baseAnimalNPC)
                {
                    if (baseAnimalNPC.HasBrain && baseAnimalNPC.TryGetComponent<BaseAIBrain>(out BaseAIBrain baseAIBrain))
                    {
                        UnityEngine.Object.Destroy(baseAIBrain);
                    }
                }
                if (entity is BaseNpc npc)
                {
                    npc.CancelInvoke(npc.TickAi);
                }
            }
            foreach (var item in FSMComponents)
            {
                if (item.workQueue.Contains(item))
                    item.workQueue.Remove(item);
            }
            foreach (var entity in toKill)
            {
                if (entity != null && !entity.IsDestroyed)
                    entity.Kill();
            }
            Pool.FreeUnmanaged(ref FSMComponents);
            Pool.FreeUnmanaged(ref toKill);
        }

        private void Unload()
        {
            _plugin = null;
        }

        // Lists every animal on the map with its prefab short name, so the config list can be checked against reality
        [ConsoleCommand("removeanimalsai.list")]
        private void ListAnimalsCommand(ConsoleSystem.Arg arg)
        {
            if (arg.Connection != null && !arg.IsAdmin)
                return;

            var counts = new SortedDictionary<string, int>(StringComparer.OrdinalIgnoreCase);
            var kinds = new Dictionary<string, string>(StringComparer.OrdinalIgnoreCase);
            var frozen = new Dictionary<string, bool>(StringComparer.OrdinalIgnoreCase);
            foreach (var networkable in BaseNetworkable.serverEntities)
            {
                var entity = networkable as BaseEntity;
                if (entity == null || entity is BasePlayer || entity is ScientistNPC2)
                    continue;

                string kind;
                if (entity is BaseNPC2 || entity.GetComponent<FSMComponent>() != null)
                    kind = "Gen2 AI";
                else if (entity.GetComponent<BaseAIBrain>() != null || entity is BaseAnimalNPC)
                    kind = "legacy AI";
                else if (entity is BaseNpc)
                    kind = "legacy NPC";
                else
                    continue;

                var name = entity.ShortPrefabName;
                counts.TryGetValue(name, out var count);
                counts[name] = count + 1;
                kinds[name] = kind;
                frozen[name] = IsFrozen(entity);
            }

            var sb = new StringBuilder();
            sb.AppendLine($"Animals on the map ({counts.Count} types):");
            foreach (var pair in counts)
                sb.AppendLine($"  {pair.Key,-24} x{pair.Value,-5} {kinds[pair.Key],-11} {(frozen[pair.Key] ? "FROZEN" : "moves")}");
            arg.ReplyWith(sb.ToString());
        }

        [HarmonyPatch(typeof(BaseNpc), nameof(BaseNpc.ServerInit)), AutoPatch]
        public static class ServerInitPatch
        {
            [HarmonyPostfix]
            private static void Postfix(BaseNpc __instance)
            {
                if (IsFrozen(__instance))
                    __instance.CancelInvoke(__instance.TickAi);
            }
        }

        [HarmonyPatch(typeof(FSMComponent), nameof(FSMComponent.SetFsmActive)), AutoPatch]
        public static class SetFsmActivePatch
        {
            [HarmonyPrefix]
            private static bool Prefix(FSMComponent __instance, bool newActive)
            {
                BaseEntity entity = __instance._baseEntity;
                if (entity == null)
                    entity = __instance.GetComponent<BaseEntity>();
                if (!IsFrozen(entity))
                    return true;
                return !newActive;
            }
        }

        // Used to skip the whole Animal queue; now only the frozen animals are taken out of it,
        // so legacy-AI animals that are allowed to move keep thinking
        [HarmonyPatch(typeof(AIThinkManager), nameof(AIThinkManager.ProcessQueue)), AutoPatch]
        public static class ProcessQueuePatch
        {
            static bool _lookupDone;
            static FieldInfo _queueField;
            static MethodInfo _queueRemove;
            static readonly List<object> _toRemove = new List<object>();

            [HarmonyPrefix]
            private static bool Prefix(AIThinkManager.QueueType queueType)
            {
                if (queueType != AIThinkManager.QueueType.Animal)
                    return true;

                var queue = GetAnimalQueue();
                if (queue == null)
                    return false; // queue not found: fall back to freezing the whole Animal queue

                foreach (var item in queue)
                {
                    var component = item as Component;
                    BaseEntity entity = component != null ? component.GetComponent<BaseEntity>() : null;
                    if (IsFrozen(entity))
                        _toRemove.Add(item);
                }
                foreach (var item in _toRemove)
                    _queueRemove.Invoke(_queueField.GetValue(null), new[] { item });
                _toRemove.Clear();
                return true;
            }

            static IEnumerable GetAnimalQueue()
            {
                if (!_lookupDone)
                {
                    _lookupDone = true;
                    foreach (var field in typeof(AIThinkManager).GetFields(BindingFlags.Static | BindingFlags.Public | BindingFlags.NonPublic))
                    {
                        if (field.Name.IndexOf("animal", StringComparison.OrdinalIgnoreCase) < 0 || !typeof(IEnumerable).IsAssignableFrom(field.FieldType))
                            continue;
                        var remove = field.FieldType.GetMethods().FirstOrDefault(m => m.Name == "Remove" && m.GetParameters().Length == 1);
                        if (remove == null)
                            continue;
                        _queueField = field;
                        _queueRemove = remove;
                        break;
                    }
                    if (_queueField == null)
                        _plugin?.PrintWarning("Could not find the animal think queue; all legacy-AI animals stay frozen regardless of config");
                }
                return _queueField?.GetValue(null) as IEnumerable;
            }
        }

        [HarmonyPatch(typeof(AnimalBrain), nameof(AnimalBrain.InitializeAI)), AutoPatch]
        public static class InitializeAIPatch
        {
            [HarmonyPrefix]
            private static bool Prefix(BaseAIBrain __instance)
            {
                return !IsFrozen(__instance.GetComponent<BaseEntity>());
            }
        }

        [HarmonyPatch(typeof(BaseCombatEntity), nameof(BaseCombatEntity.OnDied)), AutoPatch]
        public static class OnDiedRequeueFSMPatch
        {
            [HarmonyPrefix]
            private static void Prefix(BaseCombatEntity __instance)
            {
                if (__instance is BaseNPC2 npc && IsFrozen(npc) && !__instance.IsDestroyed && npc.TryGetComponent<FSMComponent>(out var fsm) && !fsm.workQueue.Contains(fsm))
                {
                    fsm.isRunning = true;
                    fsm.workQueue.Add(fsm);
                }
            }
        }

    }
}