#include <jni.h>
#include <stdio.h>
#include <wchar.h>
#include <winsock2.h>
#include <windows.h>
#include "wireguard.h"
#include <string.h>
#include <ws2tcpip.h>
#include "tech_nevets_tvpn_wg_WireGuardJNI.h"

#pragma comment(lib, "Rpcrt4.lib")

const wchar_t *APP_DIR =  L"C:\\Program Files\\TVPN\\";

typedef struct {
    WIREGUARD_INTERFACE Interface;
    WIREGUARD_PEER Peer;
    WIREGUARD_ALLOWED_IP AllowedIP;
} Config;

static WIREGUARD_CREATE_ADAPTER_FUNC *WireGuardCreateAdapter;
static WIREGUARD_OPEN_ADAPTER_FUNC *WireGuardOpenAdapter;
static WIREGUARD_CLOSE_ADAPTER_FUNC *WireGuardCloseAdapter;
static WIREGUARD_GET_ADAPTER_LUID_FUNC *WireGuardGetAdapterLUID;
static WIREGUARD_GET_RUNNING_DRIVER_VERSION_FUNC *WireGuardGetRunningDriverVersion;
static WIREGUARD_DELETE_DRIVER_FUNC *WireGuardDeleteDriver;
static WIREGUARD_SET_LOGGER_FUNC *WireGuardSetLogger;
static WIREGUARD_SET_ADAPTER_LOGGING_FUNC *WireGuardSetAdapterLogging;
static WIREGUARD_GET_ADAPTER_STATE_FUNC *WireGuardGetAdapterState;
static WIREGUARD_SET_ADAPTER_STATE_FUNC *WireGuardSetAdapterState;
static WIREGUARD_GET_CONFIGURATION_FUNC *WireGuardGetConfiguration;
static WIREGUARD_SET_CONFIGURATION_FUNC *WireGuardSetConfiguration;

static HMODULE InitializeWireGuardNT(void) {
    printf("[C] Starting InitializeWireGuardNT!\n");
    HMODULE WireGuardDll = LoadLibraryW(L"wireguard.dll");
    if (!WireGuardDll) {
        printf("[C] Error loading wireguard.dll\n");
        return NULL;
    }
#define X(Name) ((*(FARPROC *)&Name = GetProcAddress(WireGuardDll, #Name)) == NULL)
    if (X(WireGuardCreateAdapter) || X(WireGuardOpenAdapter) || X(WireGuardCloseAdapter) ||
        X(WireGuardGetAdapterLUID) || X(WireGuardGetRunningDriverVersion) || X(WireGuardDeleteDriver) ||
        X(WireGuardSetLogger) || X(WireGuardSetAdapterLogging) || X(WireGuardGetAdapterState) ||
        X(WireGuardSetAdapterState) || X(WireGuardGetConfiguration) || X(WireGuardSetConfiguration))
#undef X
    {
        DWORD LastError = GetLastError();
        fprintf(stderr, "[C] Failed to load wireguard.dll, error code: %lu\n", LastError);
        FreeLibrary(WireGuardDll);
        SetLastError(LastError);
        return NULL;
    }
    printf("[C] Finished Init WG.dll\n");
    return WireGuardDll;
}

static GUID generateGUID(void) {
    UUID uuid;
    RPC_STATUS status = UuidCreate(&uuid);

    // Handle GUID generation error
    if (status != RPC_S_OK) {
        fprintf(stderr, "[C] Failed to create GUID\n");
        exit(1);  // Exiting here to indicate a critical error; you may handle differently
    }

    // Copy the generated UUID into a GUID struct
    GUID guid = { uuid.Data1, uuid.Data2, uuid.Data3 };
    for (int i = 0; i < 8; i++) {
        guid.Data4[i] = uuid.Data4[i];
    }

    return guid;
}

JNIEXPORT jint JNICALL Java_tech_nevets_tvpn_wg_WireGuardJNI_installAdapter(JNIEnv *env, jobject obj) {
    HMODULE WireGuardDLL = InitializeWireGuardNT();

    GUID guid = generateGUID();
    WIREGUARD_ADAPTER_HANDLE Adapter = WireGuardCreateAdapter(L"TVPNet", L"WireGuard", &guid);

    struct {
        WIREGUARD_INTERFACE Interface;
        WIREGUARD_PEER DemoServer;
        WIREGUARD_ALLOWED_IP AllV4;
    } Config = {
            .Interface = {
                    .Flags = WIREGUARD_INTERFACE_HAS_PRIVATE_KEY,
                    .PeersCount = 1
            },
            .DemoServer = {
                    .Flags = WIREGUARD_PEER_HAS_PUBLIC_KEY | WIREGUARD_PEER_HAS_ENDPOINT,
                    .AllowedIPsCount = 1
            },
            .AllV4 = {
                    .AddressFamily = AF_INET
            }
    };

    WireGuardSetConfiguration(Adapter, &Config.Interface, sizeof(Config));
    WireGuardSetAdapterState(Adapter, WIREGUARD_ADAPTER_STATE_UP);

    FreeLibrary(WireGuardDLL);
    return 0;
}

JNIEXPORT jint JNICALL Java_tech_nevets_tvpn_wg_WireGuardJNI_removeAdapter(JNIEnv *env, jobject obj) {
    HMODULE WireGuardDLL = InitializeWireGuardNT();

    WIREGUARD_ADAPTER_HANDLE Adapter = WireGuardOpenAdapter(L"TVPNet");

    printf("[C] Closing Adapter...\n");
    WireGuardCloseAdapter(Adapter);

    FreeLibrary(WireGuardDLL);
    return 0;
}

JNIEXPORT jint JNICALL Java_tech_nevets_tvpn_wg_WireGuardJNI_installTunnel(JNIEnv *env, jobject obj, jstring tunnelName, jstring pathToConfig) {
    printf("[C] Started installTunnel\n");

    const jchar *rawStr = (*env)->GetStringChars(env, tunnelName, 0); //TODO mem cleanup
    LPWSTR tnlName = (LPWSTR) rawStr;

    wprintf(L"[C] Tunnel Name: %ls\n", tnlName);

    wchar_t *servicePrefix = L"WireGuardTunnel$";
    wchar_t *serviceName = malloc(sizeof(wchar_t) * (wcslen(servicePrefix) + wcslen(tnlName) + 1)); //TODO mem cleanup
    wcscpy(serviceName, servicePrefix);
    wcscat(serviceName, tnlName);

    wchar_t *displayNamePrefix = L"WireGuard Tunnel: ";
    wchar_t *displayName = malloc(sizeof(wchar_t) * (wcslen(displayNamePrefix) + wcslen(tnlName) + 1)); //TODO mem cleanup
    wcscpy(displayName, displayNamePrefix);
    wcscat(displayName, tnlName);

    printf("[C] Opening Service Manager\n");

    SC_HANDLE hSCManager = OpenSCManager(
            NULL,
            NULL,
            SC_MANAGER_ALL_ACCESS
    );
    if (hSCManager == NULL) {
        printf("[C] OpenSCManager failed (%lu)\n", GetLastError());
        return 3;
    }

    wchar_t *params = L"TVPN.exe /installtunnelservice ";
    wchar_t *conf = L".conf";
    wchar_t *exePath = malloc(sizeof(wchar_t) * (wcslen(APP_DIR) + wcslen(params) + wcslen(tnlName) + wcslen(conf) + 1)); //TODO mem cleanup
    wcscpy(exePath, APP_DIR);
    wcscat(exePath, params);
    wcscat(exePath, tnlName);
    wcscat(exePath, conf);

    wprintf(L"[C] Creating Service: %ls\n", exePath);

    SC_HANDLE hService = CreateServiceW(
            hSCManager,
            serviceName,
            displayName,
            SERVICE_ALL_ACCESS,
            SERVICE_WIN32_OWN_PROCESS,
            SERVICE_AUTO_START,
            SERVICE_ERROR_NORMAL,
            exePath,
            NULL,
            NULL,
            L"Nsi\0TcpIp\0\0",
            NULL,
            NULL
    );

    DWORD sid_info = SERVICE_SID_TYPE_UNRESTRICTED;

    printf("[C] Changing Service Config\n");

    if (!ChangeServiceConfig2W(hService, SERVICE_CONFIG_SERVICE_SID_INFO, &sid_info)) {
        printf("[C] ChangeServiceConfig2W failed (%lu)\n", GetLastError());
        (*env)->ReleaseStringChars(env, tunnelName, rawStr);
        free(serviceName);
        free(displayName);
        free(exePath);
        CloseServiceHandle(hSCManager);
        CloseServiceHandle(hService);
        return 4;
    }

    printf("[C] Cleaning up creation\n");

    CloseServiceHandle(hSCManager);
    CloseServiceHandle(hService);
    free(serviceName);
    free(displayName);
    free(exePath);
    (*env)->ReleaseStringChars(env, tunnelName, rawStr);

    printf("[C] Loading tunnel.dll\n");

    HMODULE tunnel_lib = LoadLibrary("tunnel.dll");
    if (!tunnel_lib) abort();

    BOOL (_cdecl *tunnel_proc) (_In_ LPCWSTR conf_file);
    *(FARPROC*) &tunnel_proc = GetProcAddress(tunnel_lib, "WireGuardTunnelService");
    if (!tunnel_proc) abort();

    const jchar *rawString = (*env)->GetStringChars(env, pathToConfig, NULL);
    LPCWSTR configPath = (LPCWSTR) rawString;

    //TODO figure out how to handle this
    //(*env)->ReleaseStringChars(env, tunnelName, rawStr);

    wprintf(L"[C] Calling function with arg: %ls\n", configPath);

    return tunnel_proc(configPath);
}

JNIEXPORT jint JNICALL Java_tech_nevets_tvpn_wg_WireGuardJNI_removeTunnel(JNIEnv *env, jobject obj, jstring tunnelName) {
    const jchar *nativeString = (*env)->GetStringChars(env, tunnelName, NULL);

    wchar_t *initialServiceName = L"WireGuardTunnel$";
    wchar_t *tnlName = malloc(sizeof(wchar_t) * (wcslen(initialServiceName) + wcslen(nativeString) + 1));
    wcscpy(tnlName, initialServiceName);
    wcscat(tnlName, nativeString);

    SC_HANDLE hSCManager = OpenSCManager(
            NULL,
            NULL,
            SC_MANAGER_ALL_ACCESS
            );
    if (hSCManager == NULL) {
        // Handle error
        return 1;
    }

    SC_HANDLE hService = OpenServiceW(hSCManager, tnlName, DELETE);
    if (hService != NULL) {
        if (!DeleteService(hService)) {
            // Handle error
            return 2;
        }
        CloseServiceHandle(hService);
    } else {
        // Handle error
        return 3;
    }

    CloseServiceHandle(hSCManager);
    (*env)->ReleaseStringChars(env, tunnelName, nativeString);
    free(tnlName);

    return 0;
}

JNIEXPORT jlong JNICALL Java_tech_nevets_tvpn_wg_WireGuardJNI_getWGVersion(JNIEnv *env, jobject obj) {
    WSADATA WsaData;
    if (WSAStartup(MAKEWORD(2, 2), &WsaData)) {
        fprintf(stderr, "[C] Failed to initialize Winsock\n");
        return 0;
    }
    HMODULE WireGuardDLL = InitializeWireGuardNT();
    printf("[C] Post Init WGNT\n");

    if (WireGuardDLL == NULL) {
        fprintf(stderr, "[C] Failed to initialize WireGuardNT library\n");
        return 1;
    }

    DWORD Version = WireGuardGetRunningDriverVersion();
    printf("[C] Before Printed version\n");


    //printf("[C] WG Ver: %lu\n", Version);
    wprintf(L"WireGuardNT v%u.%u loaded", (Version >> 16) & 0xff, (Version >> 0) & 0xff);

    FreeLibrary(WireGuardDLL);
    return Version;
}

JNIEXPORT jlong JNICALL Java_tech_nevets_tvpn_wg_WireGuardJNI_parseConfig(JNIEnv *env, jobject obj, jobject conf) {
    // Allocate memory for Config struct
    Config *config = (Config *)malloc(sizeof(Config));
    if (config == NULL) {
        return 0; // Memory allocation failed
    }

    // Get WGConf class and fields
    jclass wgConfClass = (*env)->GetObjectClass(env, conf);

    // Step 1: Fill Interface fields
    jfieldID ifaceFieldID = (*env)->GetFieldID(env, wgConfClass, "iface", "Ltech/nevets/tvpn/wg/WGConf$Interface;");
    jobject ifaceObj = (*env)->GetObjectField(env, conf, ifaceFieldID);

    jclass ifaceClass = (*env)->GetObjectClass(env, ifaceObj);
    config->Interface.Flags = (WIREGUARD_INTERFACE_FLAG)(*env)->GetIntField(env, ifaceObj, (*env)->GetFieldID(env, ifaceClass, "flags", "I"));
    config->Interface.ListenPort = (WORD)(*env)->GetIntField(env, ifaceObj, (*env)->GetFieldID(env, ifaceClass, "port", "I"));

    // Get and copy privateKey
    jbyteArray privateKeyArray = (jbyteArray)(*env)->GetObjectField(env, ifaceObj, (*env)->GetFieldID(env, ifaceClass, "privateKey", "[B"));
    (*env)->GetByteArrayRegion(env, privateKeyArray, 0, WIREGUARD_KEY_LENGTH, (jbyte *)config->Interface.PrivateKey);

    // Get and copy publicKey
    jbyteArray publicKeyArray = (jbyteArray)(*env)->GetObjectField(env, ifaceObj, (*env)->GetFieldID(env, ifaceClass, "publicKey", "[B"));
    (*env)->GetByteArrayRegion(env, publicKeyArray, 0, WIREGUARD_KEY_LENGTH, (jbyte *)config->Interface.PublicKey);

    config->Interface.PeersCount = (DWORD)(*env)->GetIntField(env, ifaceObj, (*env)->GetFieldID(env, ifaceClass, "peersCount", "I"));

    // Step 2: Fill Peer fields
    jfieldID peerFieldID = (*env)->GetFieldID(env, wgConfClass, "peer", "Ltech/nevets/tvpn/wg/WGConf$Peer;");
    jobject peerObj = (*env)->GetObjectField(env, conf, peerFieldID);

    jclass peerClass = (*env)->GetObjectClass(env, peerObj);
    config->Peer.Flags = (WIREGUARD_PEER_FLAG)(*env)->GetIntField(env, peerObj, (*env)->GetFieldID(env, peerClass, "flags", "I"));
    config->Peer.Reserved = (DWORD)(*env)->GetIntField(env, peerObj, (*env)->GetFieldID(env, peerClass, "reserved", "I"));

    // Get and copy publicKey
    jbyteArray peerPublicKeyArray = (jbyteArray)(*env)->GetObjectField(env, peerObj, (*env)->GetFieldID(env, peerClass, "publicKey", "[B"));
    (*env)->GetByteArrayRegion(env, peerPublicKeyArray, 0, WIREGUARD_KEY_LENGTH, (jbyte *)config->Peer.PublicKey);

    // Get and copy presharedKey
    jbyteArray presharedKeyArray = (jbyteArray)(*env)->GetObjectField(env, peerObj, (*env)->GetFieldID(env, peerClass, "presharedKey", "[B"));
    (*env)->GetByteArrayRegion(env, presharedKeyArray, 0, WIREGUARD_KEY_LENGTH, (jbyte *)config->Peer.PresharedKey);

    config->Peer.PersistentKeepalive = (WORD)(*env)->GetIntField(env, peerObj, (*env)->GetFieldID(env, peerClass, "persistentKeepalive", "I"));
    config->Peer.AllowedIPsCount = (DWORD)(*env)->GetIntField(env, peerObj, (*env)->GetFieldID(env, peerClass, "allowedIPsCount", "I"));

    // Endpoint is more complex: requires parsing from byte array and setting `SOCKADDR_INET`
    jbyteArray endpointArray = (jbyteArray)(*env)->GetObjectField(env, peerObj, (*env)->GetFieldID(env, peerClass, "endpoint", "[B"));
    jsize endpointLen = (*env)->GetArrayLength(env, endpointArray);
    jbyte *endpointBytes = (*env)->GetByteArrayElements(env, endpointArray, NULL);

    if (endpointLen == sizeof(SOCKADDR_INET)) {
        memcpy(&config->Peer.Endpoint, endpointBytes, sizeof(SOCKADDR_INET));
    }
    (*env)->ReleaseByteArrayElements(env, endpointArray, endpointBytes, 0);

    // Step 3: Fill AllowedIP fields
    jfieldID allowedIPFieldID = (*env)->GetFieldID(env, wgConfClass, "allowedIP", "Ltech/nevets/tvpn/wg/WGConf$AllowedIP;");
    jobject allowedIPObj = (*env)->GetObjectField(env, conf, allowedIPFieldID);

    jclass allowedIPClass = (*env)->GetObjectClass(env, allowedIPObj);
    jstring addressStr = (jstring)(*env)->GetObjectField(env, allowedIPObj, (*env)->GetFieldID(env, allowedIPClass, "address", "Ljava/lang/String;"));
    const char *addressChars = (*env)->GetStringUTFChars(env, addressStr, NULL);

    // Assuming address parsing logic here, converting to either IN_ADDR or IN6_ADDR based on `addressFamily`
    jstring familyStr = (jstring)(*env)->GetObjectField(env, allowedIPObj, (*env)->GetFieldID(env, allowedIPClass, "addressFamily", "Ljava/lang/String;"));
    const char *familyChars = (*env)->GetStringUTFChars(env, familyStr, NULL);

    if (strcmp(familyChars, "AF_INET") == 0) {
        config->AllowedIP.AddressFamily = AF_INET;
        inet_pton(AF_INET, addressChars, &config->AllowedIP.Address.V4);
    } else if (strcmp(familyChars, "AF_INET6") == 0) {
        config->AllowedIP.AddressFamily = AF_INET6;
        inet_pton(AF_INET6, addressChars, &config->AllowedIP.Address.V6);
    }
    (*env)->ReleaseStringUTFChars(env, familyStr, familyChars);
    (*env)->ReleaseStringUTFChars(env, addressStr, addressChars);

    config->AllowedIP.Cidr = (BYTE)(*env)->GetByteField(env, allowedIPObj, (*env)->GetFieldID(env, allowedIPClass, "cidr", "B"));

    // Return the pointer to Config struct as jlong
    return (jlong)(uintptr_t)config;

}