// Copyright (C) 2026 SharpEmu Emulator Project
// SPDX-License-Identifier: GPL-2.0-or-later

using System.Linq;
using System.Collections.Generic;
using System.IO;

namespace SharpEmu.Libs.Kernel;

public static class KernelModuleRegistry
{
    private static readonly object _gate = new();
    private static readonly Dictionary<int, ModuleEntry> _modulesByHandle = new();
    private static readonly Dictionary<string, int> _handleByPath = new(StringComparer.OrdinalIgnoreCase);
    private static readonly Dictionary<string, int> _handleByName = new(StringComparer.OrdinalIgnoreCase);
    private static readonly Dictionary<int, int> _sysmoduleHandleById = new();
    private static int _nextHandle = 1;

    private static readonly Dictionary<int, string[]> KnownSysmoduleNames = new()
    {
        [0x0006] = new[] { "libSceFiber.sprx", "libSceFiber.prx" },
        [0x0021] = new[] { "libSceRudp.sprx", "libSceRudp.prx" },
        [0x0095] = new[] { "libSceIme.sprx", "libSceIme.prx" },
        [0x0096] = new[] { "libSceImeDialog.sprx", "libSceImeDialog.prx" },
        [0x00A9] = new[] { "libSceMouse.sprx", "libSceMouse.prx" },
    };

    public readonly record struct ModuleEntry(
        int Handle,
        string Name,
        string Path,
        ulong BaseAddress,
        ulong EndAddress,
        ulong EntryPoint,
        bool IsMain,
        bool IsSystemModule);

    public static void Reset()
    {
        lock (_gate)
        {
            _modulesByHandle.Clear();
            _handleByPath.Clear();
            _handleByName.Clear();
            _sysmoduleHandleById.Clear();
            _nextHandle = 1;
        }
    }

    public static int RegisterModule(
        string? modulePath,
        ulong baseAddress,
        ulong size,
        ulong entryPoint,
        bool isMain,
        bool isSystemModule = false)
    {
        var normalizedPath = NormalizePath(modulePath);
        lock (_gate)
        {
            if (!string.IsNullOrWhiteSpace(normalizedPath) &&
                _handleByPath.TryGetValue(normalizedPath, out var existingHandle) &&
                _modulesByHandle.TryGetValue(existingHandle, out var existing))
            {
                var updated = existing with
                {
                    BaseAddress = baseAddress,
                    EndAddress = ComputeEnd(baseAddress, size),
                    EntryPoint = entryPoint,
                    IsMain = existing.IsMain || isMain,
                    IsSystemModule = existing.IsSystemModule || isSystemModule,
                };
                _modulesByHandle[existingHandle] = updated;
                _handleByName[updated.Name] = existingHandle;
                return existingHandle;
            }

            var handle = _nextHandle++;
            var name = ResolveName(normalizedPath, handle);
            var entry = new ModuleEntry(
                Handle: handle,
                Name: name,
                Path: normalizedPath,
                BaseAddress: baseAddress,
                EndAddress: ComputeEnd(baseAddress, size),
                EntryPoint: entryPoint,
                IsMain: isMain,
                IsSystemModule: isSystemModule);
            _modulesByHandle[handle] = entry;
            if (!string.IsNullOrWhiteSpace(normalizedPath))
            {
                _handleByPath[normalizedPath] = handle;
            }

            _handleByName[name] = handle;
            return handle;
        }
    }

    public static int RegisterSyntheticModule(string moduleName, bool isSystemModule)
    {
        lock (_gate)
        {
            if (TryResolveHandleByNameLocked(moduleName, out var existingHandle))
            {
                var existing = _modulesByHandle[existingHandle];
                _modulesByHandle[existingHandle] = existing with { IsSystemModule = existing.IsSystemModule || isSystemModule };
                return existingHandle;
            }

            var handle = _nextHandle++;
            var fallbackName = string.IsNullOrWhiteSpace(moduleName)
                ? $"module_{handle:X4}.sprx"
                : moduleName;
            var entry = new ModuleEntry(
                Handle: handle,
                Name: fallbackName,
                Path: string.Empty,
                BaseAddress: 0,
                EndAddress: 0,
                EntryPoint: 0,
                IsMain: false,
                IsSystemModule: isSystemModule);
            _modulesByHandle[handle] = entry;
            _handleByName[fallbackName] = handle;
            return handle;
        }
    }

    public static int MarkSysmoduleLoaded(int sysmoduleId)
    {
        lock (_gate)
        {
            if (_sysmoduleHandleById.TryGetValue(sysmoduleId, out var existingHandle) &&
                _modulesByHandle.ContainsKey(existingHandle))
            {
                return existingHandle;
            }

            var handle = TryResolveKnownSysmoduleHandleLocked(sysmoduleId, out var knownHandle)
                ? knownHandle
                : RegisterSyntheticModule($"sysmodule_0x{sysmoduleId:X4}.sprx", isSystemModule: true);
            _sysmoduleHandleById[sysmoduleId] = handle;
            return handle;
        }
    }

    public static void MarkSysmoduleUnloaded(int sysmoduleId)
    {
        lock (_gate)
        {
            _sysmoduleHandleById.Remove(sysmoduleId);
        }
    }

    public static bool IsSysmoduleLoaded(int sysmoduleId)
    {
        lock (_gate)
        {
            return _sysmoduleHandleById.ContainsKey(sysmoduleId);
        }
    }

    public static bool TryGetModuleByAddress(ulong address, out ModuleEntry module)
    {
        lock (_gate)
        {
            ModuleEntry? best = null;
            foreach (var entry in _modulesByHandle.Values)
            {
                if (entry.BaseAddress == 0 || entry.EndAddress <= entry.BaseAddress)
                {
                    continue;
                }

                if (address < entry.BaseAddress || address >= entry.EndAddress)
                {
                    continue;
                }

                if (best is null || (entry.EndAddress - entry.BaseAddress) < (best.Value.EndAddress - best.Value.BaseAddress))
                {
                    best = entry;
                }
            }

            if (best is null)
            {
                module = default;
                return false;
            }

            module = best.Value;
            return true;
        }
    }

    public static bool TryGetModuleByHandle(int handle, out ModuleEntry module)
    {
        lock (_gate)
        {
            return _modulesByHandle.TryGetValue(handle, out module);
        }
    }

    public static bool TryFindByPathOrName(string? modulePathOrName, out ModuleEntry module)
    {
        module = default;
        if (string.IsNullOrWhiteSpace(modulePathOrName))
        {
            return false;
        }

        var normalizedPath = NormalizePath(modulePathOrName);
        var fileName = ResolveName(normalizedPath, 0);
        lock (_gate)
        {
            if (!string.IsNullOrWhiteSpace(normalizedPath) &&
                _handleByPath.TryGetValue(normalizedPath, out var handleByPath) &&
                _modulesByHandle.TryGetValue(handleByPath, out module))
            {
                return true;
            }

            if (TryResolveHandleByNameLocked(fileName, out var handleByName) &&
                _modulesByHandle.TryGetValue(handleByName, out module))
            {
                return true;
            }
        }

        return false;
    }

    public static int[] GetModuleHandles(bool includeSystemModules)
    {
        lock (_gate)
        {
            return _modulesByHandle.Values
                .Where(entry => includeSystemModules || !entry.IsSystemModule)
                .OrderBy(entry => entry.Handle)
                .Select(entry => entry.Handle)
                .ToArray();
        }
    }

    public static bool TryGetFirstModule(out ModuleEntry module)
    {
        lock (_gate)
        {
            if (_modulesByHandle.Count == 0)
            {
                module = default;
                return false;
            }

            var first = _modulesByHandle.Values.OrderBy(entry => entry.Handle).First();
            module = first;
            return true;
        }
    }

    private static bool TryResolveKnownSysmoduleHandleLocked(int sysmoduleId, out int handle)
    {
        handle = 0;
        if (!KnownSysmoduleNames.TryGetValue(sysmoduleId, out var candidates))
        {
            return false;
        }

        foreach (var candidate in candidates)
        {
            if (TryResolveHandleByNameLocked(candidate, out handle))
            {
                var existing = _modulesByHandle[handle];
                _modulesByHandle[handle] = existing with { IsSystemModule = true };
                return true;
            }
        }

        return false;
    }

    private static bool TryResolveHandleByNameLocked(string moduleName, out int handle)
    {
        handle = 0;
        if (string.IsNullOrWhiteSpace(moduleName))
        {
            return false;
        }

        if (_handleByName.TryGetValue(moduleName, out handle))
        {
            return true;
        }

        var altName = moduleName.EndsWith(".sprx", StringComparison.OrdinalIgnoreCase)
            ? moduleName[..^5] + ".prx"
            : (moduleName.EndsWith(".prx", StringComparison.OrdinalIgnoreCase)
                ? moduleName[..^4] + ".sprx"
                : string.Empty);
        return !string.IsNullOrWhiteSpace(altName) && _handleByName.TryGetValue(altName, out handle);
    }

    private static ulong ComputeEnd(ulong baseAddress, ulong size)
    {
        if (size == 0)
        {
            return baseAddress;
        }

        return unchecked(baseAddress + size);
    }

    private static string NormalizePath(string? modulePath)
    {
        if (string.IsNullOrWhiteSpace(modulePath))
        {
            return string.Empty;
        }

        try
        {
            return Path.GetFullPath(modulePath).Replace('\\', '/');
        }
        catch
        {
            return modulePath.Replace('\\', '/');
        }
    }

    private static string ResolveName(string normalizedPath, int handleHint)
    {
        if (!string.IsNullOrWhiteSpace(normalizedPath))
        {
            var fileName = Path.GetFileName(normalizedPath);
            if (!string.IsNullOrWhiteSpace(fileName))
            {
                return fileName;
            }
        }

        return handleHint <= 0
            ? "module.sprx"
            : $"module_{handleHint:X4}.sprx";
    }
}