// Copyright 2026 Aaron R Robinson
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is furnished
// to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
// INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
// PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
// HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
// SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

using System;
using System.Collections.Generic;
using System.Collections.Specialized;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Runtime.InteropServices;
using System.Runtime.Versioning;
using System.Text;

namespace DNNE
{
    internal static class RustEmitter
    {
        private static readonly RustTypeProvider s_typeProvider = new RustTypeProvider();

        // Rust keywords that are valid C# identifiers but reserved in Rust.
        // C# reserved keywords (for example, 'if', 'for', 'return') are excluded
        // because they cannot appear as parameter names in managed assemblies.
        // See https://doc.rust-lang.org/reference/keywords.html
        private static readonly HashSet<string> s_rustKeywords = new(StringComparer.Ordinal)
        {
            "async", "await", "crate", "dyn", "fn", "impl", "let", "loop",
            "match", "mod", "move", "mut", "pub", "self", "Self", "super",
            "trait", "type", "use", "where", "yield",
            // Reserved for future use
            "become", "box", "final", "macro", "priv", "unsized",
        };

        /// <summary>
        /// Returns a safe Rust identifier for the given name. If the name is a
        /// Rust keyword it is prefixed with <c>r#</c>; otherwise it is returned
        /// unchanged.
        /// </summary>
        private static string SafeRustIdentifier(string name)
            => s_rustKeywords.Contains(name) ? $"r#{name}" : name;

        public static void Emit(TextWriter outputStream, string assemblyName, IEnumerable<ExportedMethod> exports, IEnumerable<string> additionalCodeStatements)
        {
            // Emit preamble
            outputStream.WriteLine(
$@"//
// Auto-generated by dnne-gen
//
// .NET Assembly: {assemblyName}
//

#![allow(non_snake_case)]
#![allow(non_upper_case_globals)]

use core::ffi::c_void;
use core::sync::atomic::{{AtomicPtr, Ordering}};

//
// Forward declarations
//

use crate::platform::get_callable_managed_function;
use crate::platform::get_fast_callable_managed_function;");

            // Emit additional code statements as comments
            if (additionalCodeStatements.Any())
            {
                outputStream.WriteLine(
$@"
//
// Additional code provided by user
//");
                foreach (var stmt in additionalCodeStatements)
                {
                    outputStream.WriteLine(stmt);
                }
            }

            // Emit string table
            outputStream.WriteLine(
@"
//
// String constants
//");
            int count = 1;
            var map = new StringDictionary();
            foreach (var method in exports)
            {
                if (map.ContainsKey(method.EnclosingTypeName))
                {
                    continue;
                }

                string id = $"T{count++}_NAME";
                var typeNameWithAssembly = $"{method.EnclosingTypeName}, {assemblyName}";
                outputStream.WriteLine(
$@"
const {id}: &[u8] = b""{typeNameWithAssembly}\0"";");
                map.Add(method.EnclosingTypeName, id);
            }

            // Emit exports
            outputStream.WriteLine(
@"
//
// Exports
//");
            var declsig = new StringBuilder();
            var callsig = new StringBuilder();
            var typesig = new StringBuilder();
            foreach (var export in exports)
            {
                string cfgGuard = GetPlatformCfg(export.Platforms);
                string cfgLine = string.IsNullOrEmpty(cfgGuard) ? "" : $"{cfgGuard}\n";

                // Create declaration and call signatures.
                declsig.Clear();
                callsig.Clear();
                typesig.Clear();
                string delim = "";
                for (int i = 0; i < export.ArgumentTypes.Length; ++i)
                {
                    var argName = SafeRustIdentifier(export.ArgumentNames[i] ?? $"arg{i}");
                    declsig.AppendFormat("{0}{1}: {2}", delim, argName, export.ArgumentTypes[i]);
                    callsig.AppendFormat("{0}{1}", delim, argName);
                    typesig.AppendFormat("{0}{1}", delim, export.ArgumentTypes[i]);
                    delim = ", ";
                }

                // Return type handling
                bool isVoid = export.ReturnType == "c_void";
                string returnAnnotation = isVoid ? "" : $" -> {export.ReturnType}";
                string fnReturnAnnotation = isVoid ? "" : $" -> {export.ReturnType}";

                string callConv = s_typeProvider.MapCallConv(export.CallingConvention);

                string classNameConstant = map[export.EnclosingTypeName];
                Debug.Assert(!string.IsNullOrEmpty(classNameConstant));

                // Generate the acquire managed function based on the export type.
                string acquireManagedFunction;
                if (export.Type == ExportType.Export)
                {
                    var delegateType = $"{export.EnclosingTypeName}+{export.MethodName}Delegate, {assemblyName}";
                    acquireManagedFunction =
$@"        let method_name = b""{export.MethodName}\0"".as_ptr();
        let delegate_type = b""{delegateType}\0"".as_ptr();
        let new_ptr = get_callable_managed_function({classNameConstant}.as_ptr(), method_name, delegate_type);";
                }
                else
                {
                    Debug.Assert(export.Type == ExportType.UnmanagedCallersOnly);
                    acquireManagedFunction =
$@"        let method_name = b""{export.MethodName}\0"".as_ptr();
        let new_ptr = get_fast_callable_managed_function({classNameConstant}.as_ptr(), method_name);";
                }

                string ptrName = $"{export.ExportName}_ptr";

                // Emit export
                outputStream.WriteLine(
$@"
// Computed from {export.EnclosingTypeName}{Type.Delimiter}{export.MethodName}{export.XmlDoc}
{cfgLine}static {ptrName}: AtomicPtr<c_void> = AtomicPtr::new(core::ptr::null_mut());

{cfgLine}pub unsafe fn {export.ExportName}({declsig}){returnAnnotation} {{
    let ptr = {ptrName}.load(Ordering::Acquire);
    let f: unsafe {callConv} fn({typesig}){fnReturnAnnotation} = if !ptr.is_null() {{
        core::mem::transmute(ptr)
    }} else {{
{acquireManagedFunction}
        {ptrName}.store(new_ptr, Ordering::Release);
        core::mem::transmute(new_ptr)
    }};
    f({callsig})
}}");
            }
        }

        private static string GetPlatformCfg(in PlatformSupport platformSupport)
        {
            var conditions = new List<string>();
            AddCfgConditions(platformSupport.Assembly, conditions);
            AddCfgConditions(platformSupport.Module, conditions);
            AddCfgConditions(platformSupport.Type, conditions);
            AddCfgConditions(platformSupport.Method, conditions);

            if (conditions.Count == 0)
            {
                return string.Empty;
            }

            if (conditions.Count == 1)
            {
                return $"#[cfg({conditions[0]})]";
            }

            return $"#[cfg(all({string.Join(", ", conditions)}))]";
        }

        private static void AddCfgConditions(in Scope scope, List<string> conditions)
        {
            var supportList = scope.Support?.ToList() ?? new List<OSPlatform>();
            var noSupportList = scope.NoSupport?.ToList() ?? new List<OSPlatform>();

            if (supportList.Count > 0)
            {
                var supported = supportList.Select(MapOSPlatformToRustCfg).ToList();
                if (supported.Count == 1)
                {
                    conditions.Add(supported[0]);
                }
                else
                {
                    conditions.Add($"any({string.Join(", ", supported)})");
                }
            }

            if (noSupportList.Count > 0)
            {
                var unsupported = noSupportList.Select(MapOSPlatformToRustCfg).ToList();
                if (unsupported.Count == 1)
                {
                    conditions.Add($"not({unsupported[0]})");
                }
                else
                {
                    conditions.Add($"not(any({string.Join(", ", unsupported)}))");
                }
            }
        }

        private static string MapOSPlatformToRustCfg(OSPlatform os)
        {
            var name = os.ToString();
            if (name.StartsWith("DNNE_"))
            {
                name = name.Substring(5);
            }

            return name.ToUpperInvariant() switch
            {
                "WINDOWS" => @"target_os = ""windows""",
                "OSX" => @"target_os = ""macos""",
                "LINUX" => @"target_os = ""linux""",
                "FREEBSD" => @"target_os = ""freebsd""",
                _ => name.ToLowerInvariant(),
            };
        }
    }
}
