using System;
using System.Collections;
using System.Collections.Generic;
using Pursuit.Code.ExtensionMethods;
using UnityEngine;

namespace Pursuit.Code
{
    public static class Memoizer
    {
        private static readonly Dictionary<object, IDictionary> Maps = new();
        private static readonly Dictionary<object, string> Labels = new();

        public static TMemoized Memoize<TArgs, TMemoized>(TArgs tuple, Func<TArgs, TMemoized> ctor, string label = null)
        {
            // We don't care what the return value is. We just want the memoized object.
            TryMemoize(tuple, ctor, out TMemoized result, label);
            return result;
        }

        /// <summary>
        ///     Attempts to memoize the result of a function. If the result has already been computed,
        ///     the function returns false and assigns the existing object into result. Otherwise, it
        ///     returns true and assigns the newly created object into result.
        /// </summary>
        /// <param name="tuple">The arguments to the constructor function</param>
        /// <param name="ctor">A function that constructs the object to be memoized</param>
        /// <param name="result">Either an existing object or a newly-created object</param>
        /// <param name="label">An optional debug label.</param>
        /// <typeparam name="TArgs">The type of the arguments to the function</typeparam>
        /// <typeparam name="TMemoized">The type of the object produced by the function</typeparam>
        /// <returns>Whether a new instance was created</returns>
        public static bool TryMemoize<TArgs, TMemoized>(TArgs tuple, Func<TArgs, TMemoized> ctor,
            out TMemoized result, string label = null)
        {
            if (!Maps.TryGetValue(ctor, out IDictionary untypedMap))
            {
                untypedMap = new Dictionary<TArgs, TMemoized>();
                Maps[ctor] = untypedMap;

                Type ctorType = ctor.GetType();
                var ctorArgs = ctorType.GetGenericArguments();

                label ??= ctorArgs[0].GetNameWithGenerics() + " => " + ctorArgs[1].GetNameWithGenerics();

                Labels[ctor] = label;
            }

            var map = (Dictionary<TArgs, TMemoized>)untypedMap;

            if (map.TryGetValue(tuple, out result)) return false;

            result = ctor(tuple);
            map[tuple] = result;

            return true;
        }

        public static void Clear()
        {
            foreach (IDictionary map in Maps.Values) map.Clear();
        }

        public static void ShowDebug()
        {
            int sum = 0;

            foreach ((object ctor, object untypedMap) in Maps)
            {
                var dict = (IDictionary)untypedMap;
                if (dict.Count == 0) continue;

                sum += dict.Count;
                GUILayout.Label($"{Labels[ctor]}: {dict.Count}");
            }

            GUILayout.Label($"Object count: {sum}");
        }
    }
}