using lkfxdok.c;
using System;
using System.CodeDom;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading.Tasks;

namespace lkfxdok.LKF
{

    public struct MemoryBlock
    {
        public string key;
        public int size;
    }

    public struct MemoryKeyCounter
    {
        public int Count;
        public int AllocCount;
        public int Size;
    }

    public class MemoryTracker : IDisposable
    {
        LKFMemoryDelegate memoryDelegate;
        bool wasAssigned = false;
        public bool KeepKeys = true;

        public Dictionary<long, MemoryBlock> Blocks { get; internal set; } = new Dictionary<long, MemoryBlock>();
        public Dictionary<string, MemoryKeyCounter> Keys { get; internal set; } = new Dictionary<string, MemoryKeyCounter>();

        public Dictionary<string, MemoryKeyCounter> FreedBlocksWithoutAlloc { get; internal set; } = new Dictionary<string, MemoryKeyCounter>();

        public long BlockCount { get; internal set; } = 0;
        public long ByteCount { get; internal set; } = 0;

        public int FreedBlockCount => FreedBlocksWithoutAlloc?.Count ?? 0;

        public MemoryTracker(bool keepKeys = true, int delegateFlags = 1)
        {
            memoryDelegate = new LKFMemoryDelegate((op, key, ptr, num, elem) =>
            {
                callbackFunction(op, key, ptr, num, elem);
                //Console.WriteLine($"{op}({key}): {num * elem}");
            }, delegateFlags);
            this.KeepKeys = keepKeys;
 
        }

        public void Assign()
        {
            wasAssigned = true;
            memoryDelegate.Assign();
        }

        public void Deassign()
        {
            wasAssigned = false;
            memoryDelegate.Deassign();
        }

        public void Dispose()
        {
            if (wasAssigned)
            {
                memoryDelegate.Deassign();
            }
        }

        public void Reset()
        {
            Blocks = new Dictionary<long, MemoryBlock>();
            Keys = new Dictionary<string, MemoryKeyCounter>();
            FreedBlocksWithoutAlloc = new Dictionary<string, MemoryKeyCounter>();
            BlockCount = 0;
            ByteCount = 0;
        }

        private void callbackFunction(string op, string key, long ptr, int num, int elem)
        {
            if (op.StartsWith("free"))
            {
                int freeSize;
                string allocKey;

                if (Blocks.TryGetValue(ptr, out var block))
                {
                    Blocks.Remove(ptr);
                    ByteCount -= block.size;
                    freeSize = -block.size;
                    allocKey = block.key;
                    BlockCount--;
                }
                else
                {
                    allocKey = key;
                    freeSize = (num > 0 ? num * elem : elem);
                    if (FreedBlocksWithoutAlloc.TryGetValue(key, out var freeKeyCounter))
                    {
                        freeKeyCounter.Count--;
                        freeKeyCounter.Size += freeSize;
                    }
                    else
                    {
                        FreedBlocksWithoutAlloc[key] = new MemoryKeyCounter { Count = -1, Size = freeSize };
                    }
                    freeSize = -freeSize;
                }
                if (Keys.TryGetValue(allocKey, out var keyCounter))
                {
                    if (keyCounter.Count > 1 || KeepKeys)
                    {
                        keyCounter.Count--;
                        keyCounter.Size += freeSize;
                        Keys[allocKey] = keyCounter;
                    }
                    else
                    {
                        Keys.Remove(allocKey);
                    }
                }
            }
            else
            {
                ByteCount += num * elem;
                BlockCount++;

                Blocks[ptr] = new MemoryBlock() { key = key, size = num * elem };
                if (Keys.TryGetValue(key, out var keyCounter))
                {
                    keyCounter.Count++;
                    keyCounter.AllocCount++;
                    keyCounter.Size += num * elem;
                    Keys[key] = keyCounter;
                }
                else
                {
                    Keys[key] = new MemoryKeyCounter() { Count = 1, AllocCount = 1, Size = num * elem };
                }
            }

            //Console.WriteLine($"Blocks = {blocks.Count}, Keys = {keys.Count}");
        }

        public IEnumerable<string> TraceMemoryBlocksPerKey()
        {
            var list = new List<string>();
            foreach (var kv in Keys)
            {
                list.Add ($"{kv.Key}: {kv.Value.Count}x, {kv.Value.Size} byte (alloc-count: {kv.Value.AllocCount})");
            }
            return list;
        }

        public IEnumerable<string> TraceFreedBlocksPerKey()
        {
            if (FreedBlocksWithoutAlloc.Count == 0)
            {
                return null;
            }

            var list = new List<string>();
            foreach (var kv in FreedBlocksWithoutAlloc)
            {
                list.Add($"{kv.Key}: {kv.Value.Count}x, {kv.Value.Size} byte (alloc-count: {kv.Value.AllocCount})");
            }
            return list;
        }

        public void PrintMemoryBlocksPerKey()
        {
            foreach (var s in TraceMemoryBlocksPerKey()) Console.WriteLine(s);
        }

        public void PrintFreedBlocksPerKey()
        {
            var strings = TraceFreedBlocksPerKey();
            if (strings != null)
            {
                Console.WriteLine("**** FREED BLOCKS THAT WERE NOT ALLOCATED:");
                foreach (var s in TraceMemoryBlocksPerKey()) Console.WriteLine(s);
            }
        }

    }
}
