#include <Windows.h>
#include <shlobj.h>
#include <thread>
#include <chrono>
#include <fstream>
#include <codecvt>
#include <locale>
#include "../client_backend/well_known.h"

#define IDM_SCAN_FILE 101
#define IDM_SCAN_FOLDER 102


std::wstring string_to_widestring(const std::string& str) {
    std::wstring_convert<std::codecvt_utf8_utf16<wchar_t>> converter;
    return converter.from_bytes(str);
}
void send_command(const std::string& command) {
    std::ofstream outputFile(MAIN_COM_PATH);
    if (outputFile.is_open()) {
        outputFile << command;
        outputFile.close();
    }
}

// Function to update the content of the text field with the provided text
void update_textfield(HWND hWndTextField, const std::string& text) {
    // Get the current text length
    int textLength = GetWindowTextLength(hWndTextField);

    // Set the selection to the end of the text field
    SendMessage(hWndTextField, EM_SETSEL, textLength, textLength);

    // Append the new text
    SendMessage(hWndTextField, EM_REPLACESEL, FALSE, (LPARAM)string_to_widestring(text).c_str());
}

void scan_file(HWND hWndTextField, const std::string& filePath) {
    // Remove the answer file
    std::remove(ANSWER_COM_PATH);
    // Display the scanned file path in the window
    update_textfield(hWndTextField, "Scanning file: " + filePath + "\r\n");
    bool answered = false;
    // Write command into com file
    //printf("%d\n",send_to_pipe("scanfile \"" + filePath + "\""));
    std::ofstream outputFile(MAIN_COM_PATH);
    if (outputFile.is_open()) {
        outputFile << "scanfile \"" << filePath << "\"";
        outputFile.close();
    }
    else {
        update_textfield(hWndTextField, "Error: Unable to talk to daemon!\n");
        return;
    }
    while (!answered) {
        // Wait for answer in file
        std::ifstream inputFile(ANSWER_COM_PATH);
        // The structure of the answer file is as follows:
        // found/not_found
        // filepath
        // hash
        // action_taken/no_action_taken
        if (inputFile.is_open()) {
            std::string status, scannedFilePath, hash, action;
            if (inputFile >> status) {
                inputFile.ignore(1); // Ignore space
                inputFile.ignore(1); // Ignore starting double quote
                if (status == "found" || status == "not_found") {
                    std::getline(inputFile, scannedFilePath, '\"'); // Read until closing double quote
                    inputFile.ignore(1); // Ignore space between filepath and hash
                    inputFile.ignore(1); // Ignore starting double quote
                    std::getline(inputFile, hash, ' '); // Read until space
                    std::getline(inputFile, action); // Read until end of line

                    answered = true;

                    if (status == "found") {
                        update_textfield(hWndTextField, "Virus found in file: " + scannedFilePath + "\r\n");
                        update_textfield(hWndTextField, "File: " + scannedFilePath + " is infected\r\n");
                        update_textfield(hWndTextField, "Hash: " + hash + "\r\n");
                        update_textfield(hWndTextField, "Action taken: " + action + "\r\n");

                    }
                    else {
                        update_textfield(hWndTextField, "No virus found in file: " + scannedFilePath + "\r\n");
                    }
                    update_textfield(hWndTextField, "------------------------------------------\r\n");
                }
            }
            else {
                answered = true;
                update_textfield(hWndTextField, "Error: Unable to talk to daemon!\n");
                update_textfield(hWndTextField, "------------------------------------------\n");
            }
            inputFile.close();
            std::remove(ANSWER_COM_PATH);
        }
        // Wait for 1 second before checking again
        std::this_thread::sleep_for(std::chrono::seconds(1));
    }
    // Remove the answer file
    std::remove(ANSWER_COM_PATH);
}

// Function to simulate folder scanning
void scan_folder(HWND hProgressBar,HWND hWndTextField, const std::string& folderPath) {
    int num_of_found = 0;
    // Remove the answer file
    std::remove(ANSWER_COM_PATH);
    // Display the scanned folder path in the window
    update_textfield(hWndTextField, "Scanning folder: " + folderPath + "\r\n");
    bool answered = false;
    // Write command into com file
    std::ofstream outputFile(MAIN_COM_PATH);
    if (outputFile.is_open()) {
        outputFile << "scanfolder \"" << folderPath << "\"";
        outputFile.close();
    }
    else {
        update_textfield(hWndTextField, "Error: Unable to talk to daemon!\n");
        return;
    }
    while (!answered) {
        // Wait for answer in file
        std::ifstream inputFile(ANSWER_COM_PATH);
        // The structure of the answer file is as follows:
        // found/not_found
        // filepath
        // hash
        // action_taken/no_action_taken
        if (inputFile.is_open()) {
            std::string status, scannedFilePath, hash, action;
            while (!inputFile.eof()) {
                if (inputFile >> status) {
                    if (status == "found" || status == "not_found") {
                        inputFile.ignore(1); // Ignore space
                        inputFile.ignore(1); // Ignore starting double quote
                        std::getline(inputFile, scannedFilePath, '\"'); // Read until closing double quote
                        inputFile.ignore(1); // Ignore space between filepath and hash
                        inputFile.ignore(1); // Ignore starting double quote
                        std::getline(inputFile, hash, ' '); // Read until space
                        std::getline(inputFile, action); // Read until end of line

                        //answered = true;

                        if (status == "found") {
                            update_textfield(hWndTextField, "Virus found in file: " + scannedFilePath + "\r\n");
                            update_textfield(hWndTextField, "File: " + scannedFilePath + " is infected\r\n");
                            update_textfield(hWndTextField, "Hash: " + hash + "\r\n");
                            update_textfield(hWndTextField, "Action taken: " + action + "\r\n");
                            num_of_found++;
                        }
                        else {
                            update_textfield(hWndTextField, "No virus found in file: " + scannedFilePath + "\r\n");
                        }
                        update_textfield(hWndTextField, "------------------------------------------\r\n");
                    }
                    else if (status == "progress") {
                        std::string progress;
                        inputFile.ignore(1); // Ignore space
						inputFile >> progress;
						SendMessage(hProgressBar, PBM_SETPOS, atoi(progress.c_str()), 0);
                    }
                    else if (status == "end")
                        answered = true;
                }
            }
            inputFile.close();
            Sleep(1000);//only see for new entrys ~ once a second
            std::ofstream(ANSWER_COM_PATH);//clear the file
        }
        // Wait for 1 second before checking again
        std::this_thread::sleep_for(std::chrono::seconds(1));
    }
    update_textfield(hWndTextField, "Folder scan completed\r\n");
    update_textfield(hWndTextField, "Number of infected files: " + std::to_string(num_of_found) + "\r\n");
    // Remove the answer file
    std::remove(ANSWER_COM_PATH);
}



std::string getFolderPath(HWND hWnd) {
    std::wstring selectedFolderPath;

    // Initialize COM
    CoInitializeEx(NULL, COINIT_APARTMENTTHREADED | COINIT_DISABLE_OLE1DDE);

    // Display the folder picker dialog
    BROWSEINFO browseInfo = { 0 };
    TCHAR selectedPath[MAX_PATH];
    browseInfo.hwndOwner = hWnd; // Set the owner window
    browseInfo.pidlRoot = NULL;  // Start from the desktop
    browseInfo.pszDisplayName = selectedPath;
    browseInfo.lpszTitle = TEXT("Select a folder");
    browseInfo.ulFlags = BIF_RETURNONLYFSDIRS | BIF_NEWDIALOGSTYLE;

    LPITEMIDLIST pidlSelected = SHBrowseForFolder(&browseInfo);
    if (pidlSelected != NULL) {
        SHGetPathFromIDList(pidlSelected, selectedPath);

        // Convert TCHAR array to std::string
        std::wstring_convert<std::codecvt_utf8_utf16<wchar_t>> converter;
        selectedFolderPath = selectedPath;

        // Free the PIDL
        IMalloc* pMalloc;
        if (SUCCEEDED(SHGetMalloc(&pMalloc))) {
            pMalloc->Free(pidlSelected);
            pMalloc->Release();
        }
    }

    // Uninitialize COM
    CoUninitialize();

    return std::string(selectedFolderPath.begin(), selectedFolderPath.end());
}

LRESULT CALLBACK WndProc(HWND hWnd, UINT message, WPARAM wParam, LPARAM lParam) {
    static HWND hWndTextField;
    static HWND hProgressBar;
    static HBRUSH hBackgroundBrush = CreateSolidBrush(RGB(255, 255, 255)); // White color
    RECT rect;
    GetClientRect(hWnd, &rect);
    int width = rect.right - rect.left;
    int height = rect.bottom - rect.top;
    switch (message) {
    case WM_CREATE:
    {

        // Create the "Scan File" button
        CreateWindowEx(NULL, L"BUTTON", L"Scan File",
            WS_TABSTOP | WS_VISIBLE | WS_CHILD | BS_DEFPUSHBUTTON,
            20, 10, 100, 30, hWnd, (HMENU)IDM_SCAN_FILE, GetModuleHandle(NULL), NULL);

        // Create the "Scan Folder" button
        CreateWindowEx(NULL, L"BUTTON", L"Scan Folder",
            WS_TABSTOP | WS_VISIBLE | WS_CHILD | BS_DEFPUSHBUTTON,
            20, 50, 100, 30, hWnd, (HMENU)IDM_SCAN_FOLDER, GetModuleHandle(NULL), NULL);

        // Create a multi-line edit control for displaying text
        hWndTextField = CreateWindowEx(WS_EX_CLIENTEDGE, L"EDIT", NULL,
            WS_CHILD | WS_VISIBLE | WS_VSCROLL | ES_MULTILINE | ES_AUTOVSCROLL | ES_READONLY,
            140, 10, width-140-20, height-10-50, hWnd, NULL, NULL, NULL);
        update_textfield(hWndTextField, "Welcome to Cyberhex endpoint protection!\r\n");

        hProgressBar = CreateWindowEx(0, PROGRESS_CLASS, NULL,
            WS_CHILD | WS_VISIBLE | PBS_SMOOTH,
            140, height-40, 200, 20, hWnd, NULL, NULL, NULL);
        SendMessage(hProgressBar, PBM_SETRANGE, 0, MAKELPARAM(0, 100));
        SendMessage(hProgressBar, PBM_SETSTEP, 1, 0);
    }
    break;
    case WM_SIZE:
    {
		// Resize the text field to fit the window
		MoveWindow(hWndTextField, 140, 10, width - 140 - 20, height - 10 - 50, TRUE);
        MoveWindow(hProgressBar, 140, height - 40, 200, 20, TRUE);
        break;
	}
    case WM_COMMAND:
    {
        int wmId = LOWORD(wParam);
        // Parse the menu selections:
        switch (wmId) {
        case IDM_SCAN_FILE:
        {
            // Open file dialog to select a file
            // Call scan_file function in a separate thread
            OPENFILENAME ofn;
            WCHAR szFile[MAX_PATH] = L""; // Use WCHAR for Unicode compatibility

            ZeroMemory(&ofn, sizeof(ofn));
            ofn.lStructSize = sizeof(ofn);
            ofn.hwndOwner = hWnd;
            ofn.lpstrFile = szFile;
            ofn.lpstrFile[0] = L'\0'; // Use wide character constant L'\0'
            ofn.nMaxFile = sizeof(szFile);
            ofn.lpstrFilter = L"All Files\0*.*\0"; // Use wide character string literal L""
            ofn.nFilterIndex = 1;
            ofn.lpstrFileTitle = NULL;
            ofn.lpstrInitialDir = NULL;
            ofn.Flags = OFN_PATHMUSTEXIST | OFN_FILEMUSTEXIST;

            if (GetOpenFileName(&ofn) == TRUE) {
                std::wstring selectedFile = ofn.lpstrFile; // Use std::wstring for wide characters
                std::string narrowSelectedFile(selectedFile.begin(), selectedFile.end());
                std::thread(scan_file, hWndTextField, narrowSelectedFile).detach();
            }
        }
        break;
        case IDM_SCAN_FOLDER:
        {
            // Open folder picker dialog
            // Call scan_folder function in a separate thread
            std::string selectedFolder = getFolderPath(hWnd);
            std::thread(scan_folder,hProgressBar, hWndTextField, selectedFolder).detach();
        }
        break;
        default:
            return DefWindowProc(hWnd, message, wParam, lParam);
        }
    }
    break;
 
    case WM_PAINT:
    {
        PAINTSTRUCT ps;
        HDC hdc = BeginPaint(hWnd, &ps);
        RECT rect;
        GetClientRect(hWnd, &rect);
        FillRect(hdc, &rect, hBackgroundBrush); // Fill the entire client area with white color
        EndPaint(hWnd, &ps);
    }
    break;
    case WM_DESTROY:
        PostQuitMessage(0);
        break;
    default:
        return DefWindowProc(hWnd, message, wParam, lParam);
    }
    return 0;
}

int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLine, int nCmdShow) {
    const wchar_t CLASS_NAME[] = L"Cyberhex endpoint protection frontend";

    WNDCLASS wc = { };

    wc.lpfnWndProc = WndProc;
    wc.hInstance = hInstance;
    wc.lpszClassName = CLASS_NAME;

    RegisterClass(&wc);

    HWND hWnd = CreateWindowEx(
        0,
        CLASS_NAME,
        L"Cyberhex endpoint protection",
        WS_OVERLAPPEDWINDOW,
        CW_USEDEFAULT, 
        CW_USEDEFAULT, 
        640, 
        480,
        NULL,
        NULL,
        hInstance,
        NULL
    );


    if (hWnd == NULL) {
        return 0;
    }

    ShowWindow(hWnd, nCmdShow);
    UpdateWindow(hWnd);

    MSG msg = { };
    while (GetMessage(&msg, NULL, 0, 0)) {
        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }

    return 0;
}