using System;
using System.Collections.Generic;

namespace VMdemo.Simulation
{
    public enum PageTableProbeStatus
    {
        MissingL1Table = 0,
        MissingL2Entry = 1,
        NotPresentEntry = 2,
        PresentEntry = 3
    }

    public readonly struct PageTableProbeResult
    {
        public PageTableProbeResult(PageTableProbeStatus status, bool hasEntry, PageTableEntry entry)
        {
            Status = status;
            HasEntry = hasEntry;
            Entry = entry;
        }

        public PageTableProbeStatus Status { get; }
        public bool HasEntry { get; }
        public PageTableEntry Entry { get; }
    }

    public class TwoLevelPageTable
    {
        private readonly Dictionary<ulong, Dictionary<ulong, PageTableEntry>> _table;
        private readonly ulong _l1Mask;
        private readonly ulong _l2Mask;
        private readonly ulong _vpnMask;

        public TwoLevelPageTable(int l1Bits, int l2Bits)
        {
            if (l1Bits < 0 || l1Bits > 64)
            {
                throw new ArgumentOutOfRangeException(nameof(l1Bits), "l1Bits 必须在 0 到 64 之间。");
            }

            if (l2Bits < 0 || l2Bits > 64)
            {
                throw new ArgumentOutOfRangeException(nameof(l2Bits), "l2Bits 必须在 0 到 64 之间。");
            }

            if (l1Bits + l2Bits <= 0 || l1Bits + l2Bits > 64)
            {
                throw new ArgumentOutOfRangeException(nameof(l2Bits), "l1Bits + l2Bits 必须在 1 到 64 之间。");
            }

            L1Bits = l1Bits;
            L2Bits = l2Bits;
            _table = new Dictionary<ulong, Dictionary<ulong, PageTableEntry>>();
            _l1Mask = GetBitMask(l1Bits);
            _l2Mask = GetBitMask(l2Bits);
            _vpnMask = GetBitMask(l1Bits + l2Bits);
        }

        public int L1Bits { get; }
        public int L2Bits { get; }

        public bool TryGetPresentPfn(ulong l1Index, ulong l2Index, out ulong pfn)
        {
            ValidateIndices(l1Index, l2Index);

            if (!_table.TryGetValue(l1Index, out var l2Table))
            {
                pfn = 0UL;
                return false;
            }

            if (!l2Table.TryGetValue(l2Index, out var entry) || !entry.Present)
            {
                pfn = 0UL;
                return false;
            }

            pfn = entry.Pfn;
            return true;
        }

        public bool TryGetEntry(ulong l1Index, ulong l2Index, out PageTableEntry entry)
        {
            ValidateIndices(l1Index, l2Index);

            if (_table.TryGetValue(l1Index, out var l2Table) &&
                l2Table.TryGetValue(l2Index, out var value))
            {
                entry = value;
                return true;
            }

            entry = PageTableEntry.NotPresent;
            return false;
        }

        public PageTableProbeResult Probe(ulong l1Index, ulong l2Index)
        {
            ValidateIndices(l1Index, l2Index);

            if (!_table.TryGetValue(l1Index, out var l2Table))
            {
                return new PageTableProbeResult(
                    PageTableProbeStatus.MissingL1Table,
                    hasEntry: false,
                    entry: PageTableEntry.NotPresent);
            }

            if (!l2Table.TryGetValue(l2Index, out var entry))
            {
                return new PageTableProbeResult(
                    PageTableProbeStatus.MissingL2Entry,
                    hasEntry: false,
                    entry: PageTableEntry.NotPresent);
            }

            if (!entry.Present)
            {
                return new PageTableProbeResult(
                    PageTableProbeStatus.NotPresentEntry,
                    hasEntry: true,
                    entry: entry);
            }

            return new PageTableProbeResult(
                PageTableProbeStatus.PresentEntry,
                hasEntry: true,
                entry: entry);
        }

        public void SetPresent(ulong l1Index, ulong l2Index, ulong pfn, bool dirty = false)
        {
            ValidateIndices(l1Index, l2Index);

            if (!_table.TryGetValue(l1Index, out var l2Table))
            {
                l2Table = new Dictionary<ulong, PageTableEntry>();
                _table[l1Index] = l2Table;
            }

            l2Table[l2Index] = new PageTableEntry(pfn, true, dirty);
        }

        public bool MarkNotPresent(ulong l1Index, ulong l2Index)
        {
            ValidateIndices(l1Index, l2Index);

            if (!_table.TryGetValue(l1Index, out var l2Table) ||
                !l2Table.TryGetValue(l2Index, out var entry))
            {
                return false;
            }

            if (!entry.Present)
            {
                return false;
            }

            l2Table[l2Index] = entry.MarkNotPresent();
            return true;
        }

        public void GetIndicesFromVpn(ulong vpn, out ulong l1Index, out ulong l2Index)
        {
            if (L1Bits + L2Bits < 64 && (vpn & ~_vpnMask) != 0UL)
            {
                throw new ArgumentOutOfRangeException(nameof(vpn), "vpn 超出页表可表示范围。");
            }

            l2Index = vpn & _l2Mask;
            l1Index = L2Bits >= 64 ? 0UL : vpn >> L2Bits;
            ValidateIndices(l1Index, l2Index);
        }

        private void ValidateIndices(ulong l1Index, ulong l2Index)
        {
            if (!IsWithinBits(l1Index, L1Bits, _l1Mask))
            {
                throw new ArgumentOutOfRangeException(nameof(l1Index), "l1Index 超出 L1 位宽范围。");
            }

            if (!IsWithinBits(l2Index, L2Bits, _l2Mask))
            {
                throw new ArgumentOutOfRangeException(nameof(l2Index), "l2Index 超出 L2 位宽范围。");
            }
        }

        private static bool IsWithinBits(ulong value, int bits, ulong mask)
        {
            if (bits == 64)
            {
                return true;
            }

            if (bits == 0)
            {
                return value == 0UL;
            }

            return (value & ~mask) == 0UL;
        }

        private static ulong GetBitMask(int bits)
        {
            if (bits <= 0)
            {
                return 0UL;
            }

            if (bits >= 64)
            {
                return ulong.MaxValue;
            }

            return (1UL << bits) - 1UL;
        }
    }
}