package org.kne.cloud.network.tun;
import java.lang.foreign.*;
import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;

import org.kne.cloud.network.tun.WindowsDNSManager.DnsInterfaceSettings;
import org.kne.cloud.network.tun.WindowsInterfaceConverter.GUID;
public class WintunDriver {
    
    // Memory layouts matching C types
    public static final ValueLayout.OfInt DWORD = ValueLayout.JAVA_INT;
    public static final ValueLayout.OfLong DWORD64 = ValueLayout.JAVA_LONG;
    public static final AddressLayout POINTER = ValueLayout.ADDRESS;
    public static final ValueLayout.OfByte BYTE = ValueLayout.JAVA_BYTE;
    public static final ValueLayout.OfShort WORD = ValueLayout.JAVA_SHORT;
    
    // Constants from the header
    public static final int WINTUN_MIN_RING_CAPACITY = 0x20000; // 128kiB
    public static final int WINTUN_MAX_RING_CAPACITY = 0x4000000; // 64MiB
    public static final int WINTUN_MAX_IP_PACKET_SIZE = 0xFFFF;
    
    // Opaque handle types (represented as MemorySegment)
    public static class WintunAdapterHandle extends NativeHandle {
        WintunAdapterHandle(MemorySegment address) {
            super(address);
        }
    }
    
    public static class WintunSessionHandle extends NativeHandle {
        WintunSessionHandle(MemorySegment address) {
            super(address);
        }
    }
    
    // Native library loading
    private static final SymbolLookup LOOKUP;
    private static final Linker LINKER = Linker.nativeLinker();
    
    static {
        // Load wintun.dll from system or current directory
        System.loadLibrary("wintun");
        LOOKUP = SymbolLookup.loaderLookup();
    }
    
    // Native function handles
    private static final MethodHandle WINTUN_CREATE_ADAPTER;
    private static final MethodHandle WINTUN_OPEN_ADAPTER;
    private static final MethodHandle WINTUN_CLOSE_ADAPTER;
    private static final MethodHandle WINTUN_DELETE_DRIVER;
    private static final MethodHandle WINTUN_GET_ADAPTER_LUID;
    private static final MethodHandle WINTUN_GET_RUNNING_DRIVER_VERSION;
    private static final MethodHandle WINTUN_SET_LOGGER;
    private static final MethodHandle WINTUN_START_SESSION;
    private static final MethodHandle WINTUN_END_SESSION;
    private static final MethodHandle WINTUN_GET_READ_WAIT_EVENT;
    private static final MethodHandle WINTUN_RECEIVE_PACKET;
    private static final MethodHandle WINTUN_RELEASE_RECEIVE_PACKET;
    private static final MethodHandle WINTUN_ALLOCATE_SEND_PACKET;
    private static final MethodHandle WINTUN_SEND_PACKET;
    
    static {
        try {
            // Define function descriptors matching the C signatures
            FunctionDescriptor createAdapterDesc = FunctionDescriptor.of(
                POINTER,    // return: WINTUN_ADAPTER_HANDLE
                POINTER,    // Name: LPCWSTR
                POINTER,    // TunnelType: LPCWSTR
                POINTER     // RequestedGUID: const GUID*
            );
            
            FunctionDescriptor openAdapterDesc = FunctionDescriptor.of(
                POINTER,    // return: WINTUN_ADAPTER_HANDLE
                POINTER     // Name: LPCWSTR
            );
            
            FunctionDescriptor closeAdapterDesc = FunctionDescriptor.ofVoid(
                POINTER     // Adapter: WINTUN_ADAPTER_HANDLE
            );
            
            // 修正这里：去掉额外的VOID参数
            FunctionDescriptor deleteDriverDesc = FunctionDescriptor.of(
                ValueLayout.JAVA_INT // BOOL
            );
            
            FunctionDescriptor getAdapterLuidDesc = FunctionDescriptor.ofVoid(
                POINTER,    // Adapter: WINTUN_ADAPTER_HANDLE
                POINTER     // Luid: NET_LUID*
            );
            
            // 修正这里：去掉额外的VOID参数
            FunctionDescriptor getDriverVersionDesc = FunctionDescriptor.of(
                DWORD      // DWORD
            );
            
            FunctionDescriptor setLoggerDesc = FunctionDescriptor.ofVoid(
                POINTER     // NewLogger: WINTUN_LOGGER_CALLBACK
            );
            
            FunctionDescriptor startSessionDesc = FunctionDescriptor.of(
                POINTER,    // return: WINTUN_SESSION_HANDLE
                POINTER,    // Adapter: WINTUN_ADAPTER_HANDLE
                DWORD       // Capacity: DWORD
            );
            
            FunctionDescriptor endSessionDesc = FunctionDescriptor.ofVoid(
                POINTER     // Session: WINTUN_SESSION_HANDLE
            );
            
            FunctionDescriptor getReadWaitEventDesc = FunctionDescriptor.of(
                POINTER,    // HANDLE
                POINTER     // Session: WINTUN_SESSION_HANDLE
            );
            
            FunctionDescriptor receivePacketDesc = FunctionDescriptor.of(
                POINTER,    // return: BYTE*
                POINTER,    // Session: WINTUN_SESSION_HANDLE
                POINTER     // PacketSize: DWORD*
            );
            
            FunctionDescriptor releaseReceivePacketDesc = FunctionDescriptor.ofVoid(
                POINTER,    // Session: WINTUN_SESSION_HANDLE
                POINTER     // Packet: const BYTE*
            );
            
            FunctionDescriptor allocateSendPacketDesc = FunctionDescriptor.of(
                POINTER,    // return: BYTE*
                POINTER,    // Session: WINTUN_SESSION_HANDLE
                DWORD       // PacketSize: DWORD
            );
            
            FunctionDescriptor sendPacketDesc = FunctionDescriptor.ofVoid(
                POINTER,    // Session: WINTUN_SESSION_HANDLE
                POINTER     // Packet: const BYTE*
            );
            
            // Look up the symbols
            WINTUN_CREATE_ADAPTER = LINKER.downcallHandle(
                LOOKUP.find("WintunCreateAdapter").orElseThrow(),
                createAdapterDesc
            );
            
            WINTUN_OPEN_ADAPTER = LINKER.downcallHandle(
                LOOKUP.find("WintunOpenAdapter").orElseThrow(),
                openAdapterDesc
            );
            
            WINTUN_CLOSE_ADAPTER = LINKER.downcallHandle(
                LOOKUP.find("WintunCloseAdapter").orElseThrow(),
                closeAdapterDesc
            );
            
            WINTUN_DELETE_DRIVER = LINKER.downcallHandle(
                LOOKUP.find("WintunDeleteDriver").orElseThrow(),
                deleteDriverDesc
            );
            
            WINTUN_GET_ADAPTER_LUID = LINKER.downcallHandle(
                LOOKUP.find("WintunGetAdapterLUID").orElseThrow(),
                getAdapterLuidDesc
            );
            
            WINTUN_GET_RUNNING_DRIVER_VERSION = LINKER.downcallHandle(
                LOOKUP.find("WintunGetRunningDriverVersion").orElseThrow(),
                getDriverVersionDesc
            );
            
            WINTUN_SET_LOGGER = LINKER.downcallHandle(
                LOOKUP.find("WintunSetLogger").orElseThrow(),
                setLoggerDesc
            );
            
            WINTUN_START_SESSION = LINKER.downcallHandle(
                LOOKUP.find("WintunStartSession").orElseThrow(),
                startSessionDesc
            );
            
            WINTUN_END_SESSION = LINKER.downcallHandle(
                LOOKUP.find("WintunEndSession").orElseThrow(),
                endSessionDesc
            );
            
            WINTUN_GET_READ_WAIT_EVENT = LINKER.downcallHandle(
                LOOKUP.find("WintunGetReadWaitEvent").orElseThrow(),
                getReadWaitEventDesc
            );
            
            WINTUN_RECEIVE_PACKET = LINKER.downcallHandle(
                LOOKUP.find("WintunReceivePacket").orElseThrow(),
                receivePacketDesc
            );
            
            WINTUN_RELEASE_RECEIVE_PACKET = LINKER.downcallHandle(
                LOOKUP.find("WintunReleaseReceivePacket").orElseThrow(),
                releaseReceivePacketDesc
            );
            
            WINTUN_ALLOCATE_SEND_PACKET = LINKER.downcallHandle(
                LOOKUP.find("WintunAllocateSendPacket").orElseThrow(),
                allocateSendPacketDesc
            );
            
            WINTUN_SEND_PACKET = LINKER.downcallHandle(
                LOOKUP.find("WintunSendPacket").orElseThrow(),
                sendPacketDesc
            );
            
        } catch (Throwable e) {
            throw new RuntimeException("Failed to initialize native function handles", e);
        }
    }
    
 // Helper method to convert Java String to Windows wide string (UTF-16 LE)
    private static MemorySegment toWideString(String str, Arena arena) {
        if (str == null) return MemorySegment.NULL;
        
        // Convert to UTF-16LE bytes with null terminator
        byte[] utf16Bytes = str.getBytes(java.nio.charset.StandardCharsets.UTF_16LE);
        MemorySegment segment = arena.allocate(utf16Bytes.length + 2);
        
        // 使用 copyFrom 替代原来的 copy 方法
        // 注意参数顺序：源数组 -> 目标MemorySegment
        segment.copyFrom(MemorySegment.ofArray(utf16Bytes));
        
        // 添加 null 终止符 (两个字节)
        segment.set(ValueLayout.JAVA_BYTE, utf16Bytes.length, (byte)0);
        segment.set(ValueLayout.JAVA_BYTE, utf16Bytes.length + 1, (byte)0);
        
        return segment;
    }
    
    // Java API methods
    
    public static WintunAdapterHandle createAdapter(String name, String tunnelType, MemorySegment requestedGuid) {
        try (Arena arena = Arena.ofConfined()) {
            MemorySegment nameSegment = toWideString(name, arena);
            MemorySegment tunnelTypeSegment = toWideString(tunnelType, arena);
            
            MemorySegment result = (MemorySegment) WINTUN_CREATE_ADAPTER.invokeExact(
                nameSegment,
                tunnelTypeSegment,
               (MemorySegment)( requestedGuid != null ? requestedGuid : MemorySegment.NULL)
            );
            
            if (result.equals(MemorySegment.NULL)) {
                throw new WindowsAPIException("Failed to create adapter");
            }
            
            return new WintunAdapterHandle(result);
        } catch (Throwable e) {
            throw new RuntimeException("Error calling WintunCreateAdapter", e);
        }
    }
    
    public static WintunAdapterHandle openAdapter(String name) {
        try (Arena arena = Arena.ofConfined()) {
            MemorySegment nameSegment = toWideString(name, arena);
            
            MemorySegment result = (MemorySegment) WINTUN_OPEN_ADAPTER.invokeExact(nameSegment);
            
            if (result.equals(MemorySegment.NULL)) {
                throw new WindowsAPIException("Failed to open adapter");
            }
            
            return new WintunAdapterHandle(result);
        } catch (Throwable e) {
            throw new RuntimeException("Error calling WintunOpenAdapter", e);
        }
    }
    
    public static void closeAdapter(WintunAdapterHandle adapter) {
        try {
            WINTUN_CLOSE_ADAPTER.invokeExact(adapter.address());
        } catch (Throwable e) {
            throw new RuntimeException("Error calling WintunCloseAdapter", e);
        }
    }
    
    public static boolean deleteDriver() {
        try {
            int result = (int) WINTUN_DELETE_DRIVER.invokeExact();
            return result != 0;
        } catch (Throwable e) {
            throw new RuntimeException("Error calling WintunDeleteDriver", e);
        }
    }
    
    public static long getAdapterLuid(WintunAdapterHandle adapter) {
        // 直接在全局Arena中分配
        
        try ( Arena ara=Arena.ofConfined()){
        MemorySegment luidm = ara.allocate(ValueLayout.JAVA_LONG);
            WINTUN_GET_ADAPTER_LUID.invokeExact(adapter.address(), luidm);
            return luidm.get(ValueLayout.JAVA_LONG, 0);
        } catch (Throwable e) {
            throw new RuntimeException("Error calling WintunGetAdapterLuid", e);
        }
    }
    
    public static int getRunningDriverVersion() {
        try {
            return (int) WINTUN_GET_RUNNING_DRIVER_VERSION.invokeExact();
        } catch (Throwable e) {
            throw new RuntimeException("Error calling WintunGetRunningDriverVersion", e);
        }
    }
    
    public static void setLogger(MemorySegment loggerCallback) {
        try {
            WINTUN_SET_LOGGER.invokeExact(loggerCallback);
        } catch (Throwable e) {
            throw new RuntimeException("Error calling WintunSetLogger", e);
        }
    }
    
    public static WintunSessionHandle startSession(WintunAdapterHandle adapter, int capacity) {
        try {
            if (capacity < WINTUN_MIN_RING_CAPACITY || capacity > WINTUN_MAX_RING_CAPACITY) {
                throw new IllegalArgumentException(
                    String.format("Capacity must be between %d and %d", 
                        WINTUN_MIN_RING_CAPACITY, WINTUN_MAX_RING_CAPACITY)
                );
            }
            
            // Check if capacity is power of two
            if ((capacity & (capacity - 1)) != 0) {
                throw new IllegalArgumentException("Capacity must be a power of two");
            }
            
            MemorySegment result = (MemorySegment) WINTUN_START_SESSION.invokeExact(
                adapter.address(),
                capacity
            );
            
            if (result.equals(MemorySegment.NULL)) {
                throw new WindowsAPIException("Failed to start session");
            }
            
            return new WintunSessionHandle(result);
        } catch (Throwable e) {
            throw new RuntimeException("Error calling WintunStartSession", e);
        }
    }
    
    public static void endSession(WintunSessionHandle session) {
        try {
            WINTUN_END_SESSION.invokeExact(session.address());
        } catch (Throwable e) {
            throw new RuntimeException("Error calling WintunEndSession", e);
        }
    }
    
    public static MemorySegment getReadWaitEvent(WintunSessionHandle session) {
        try {
            return (MemorySegment) WINTUN_GET_READ_WAIT_EVENT.invokeExact(session.address());
        } catch (Throwable e) {
            throw new RuntimeException("Error calling WintunGetReadWaitEvent", e);
        }
    }
    
    public static Packet receivePacket(WintunSessionHandle session) {
        try (Arena arena = Arena.ofConfined()) {
            MemorySegment packetSizePtr = arena.allocate(DWORD);
            
            MemorySegment packetPtr = (MemorySegment) WINTUN_RECEIVE_PACKET.invokeExact(
                session.address(),
                packetSizePtr
            );
            
            if (packetPtr.equals(MemorySegment.NULL)) {
                // Check for specific errors via GetLastError if needed
                return null;
            }
            
            int packetSize = packetSizePtr.get(DWORD, 0);
            MemorySegment packetData = packetPtr.reinterpret(packetSize);
            
            return new Packet(session, packetData, packetSize);
        } catch (Throwable e) {
            throw new RuntimeException("Error calling WintunReceivePacket", e);
        }
    }
    
    public static SendPacket allocateSendPacket(WintunSessionHandle session, int packetSize) {
        try {
            if (packetSize > WINTUN_MAX_IP_PACKET_SIZE) {
                throw new IllegalArgumentException(
                    String.format("Packet size exceeds maximum of %d", WINTUN_MAX_IP_PACKET_SIZE)
                );
            }
            
            MemorySegment packetPtr = (MemorySegment) WINTUN_ALLOCATE_SEND_PACKET.invokeExact(
                session.address(),
                packetSize
            );
            
            if (packetPtr.equals(MemorySegment.NULL)) {
                return null;
            }
            
            MemorySegment packetData = packetPtr.reinterpret(packetSize);
            return new SendPacket(session, packetData, packetSize);
        } catch (Throwable e) {
            throw new RuntimeException("Error calling WintunAllocateSendPacket", e);
        }
    }
    
    // Helper classes
    
    public static class NativeHandle {
        private final MemorySegment address;
        
        protected NativeHandle(MemorySegment address) {
            this.address = address;
        }
        
        public MemorySegment address() {
            return address;
        }
    }
    
    public static class Packet {
        private final WintunSessionHandle session;
        private final MemorySegment data;
        private final int size;
        
        Packet(WintunSessionHandle session, MemorySegment data, int size) {
            this.session = session;
            this.data = data;
            this.size = size;
        }
        
        public MemorySegment data() {
            return data;
        }
        
        public int size() {
            return size;
        }
        
        public void release() {
            try {
                WINTUN_RELEASE_RECEIVE_PACKET.invokeExact(session.address(), data);
            } catch (Throwable e) {
                throw new RuntimeException("Error releasing packet", e);
            }
        }
    }
    
    public static class SendPacket {
        private final WintunSessionHandle session;
        private final MemorySegment data;
        private final int size;
        
        SendPacket(WintunSessionHandle session, MemorySegment data, int size) {
            this.session = session;
            this.data = data;
            this.size = size;
        }
        
        public MemorySegment data() {
            return data;
        }
        
        public int size() {
            return size;
        }
        
        public void send() {
            try {
                WINTUN_SEND_PACKET.invokeExact(session.address(), data);
            } catch (Throwable e) {
                throw new RuntimeException("Error sending packet", e);
            }
        }
    }
    
 
    
    // Logger callback interface for Java
    @FunctionalInterface
    public interface LoggerCallback {
        void log(int level, long timestamp, String message);
    }
    
    // Create a native logger callback from Java interface
    public static MemorySegment createLoggerCallback(LoggerCallback callback, Arena arena) {
        return MemorySegment.ofAddress(0); // Implementation would use upcall stub
        // Note: Full implementation requires creating upcall stub with LINKER.upcallStub
    }
    
    // Example usage
    public static void main(String[] args) {
        try {
            // Check driver version
            int version = getRunningDriverVersion();
            System.out.println("Wintun driver version: " + version);
            
            // Create adapter
            WintunAdapterHandle adapter = createAdapter("MyWintun", "Wintun", null);
            
            
            
            // Start session
            WintunSessionHandle session = startSession(adapter, WINTUN_MIN_RING_CAPACITY);
            
            
            
         // Get adapter LUID
            GUID guid=null;
            long luid=getAdapterLuid(adapter);
            System.out.println("Adapter LUID: " + luid);
            guid=WindowsInterfaceConverter.convertLuidToGuid(luid);
            System.out.println("Adapter GUID: "+guid);
            MemorySegment wintunGuid=guid.toMemorySegment(Arena.ofAuto());
            
     
            
            System.out.println("\n2. Setting DNS for Wintun adapter:");
            
            DnsInterfaceSettings dis= new DnsInterfaceSettings();
            dis.setNameServers(List.of("8.8.8.8"));
            dis.setFlags(WindowsDNSManager.DNS_SETTING_NAMESERVER);
            WindowsDNSManager.setInterfaceDnsSettings(guid, dis);
            System.out.println(dis);
            
            // Get read wait event (could be used with Java NIO Selector)
            MemorySegment readEvent = getReadWaitEvent(session);
            
            try {
                // Example packet loop
                while (true) {
                    Packet packet = receivePacket(session);
                    if (packet != null) {
                        try {
                            // Process packet data
                            byte[] packetData = packet.data().toArray(ValueLayout.JAVA_BYTE);
                            System.out.println("Received packet size: " + packet.size());
                            
                            // Create response packet
                            SendPacket sendPacket = allocateSendPacket(session, packet.size());
                            if (sendPacket != null) {
                                // Copy data to send buffer - Java 23 方式
                                sendPacket.data().copyFrom(MemorySegment.ofArray(packetData));
                                sendPacket.send();
                            }
                        } finally {
                            packet.release();
                        }
                    }
                }
            } finally {
                // Cleanup
                endSession(session);
                closeAdapter(adapter);
            }
            
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}