#import "Basic";
#import "Windows";
#import "Windows_Utf8";
#import "Thread";
#import "String";

#load "win32.jai";
#load "dialog.jai";
#load "updater.jai";

// Custom constants
TIMER_PROCESS_CHECK :: 1;

IDI_APPLICATION   :: cast(*u16) 32512;
IDI_SHIELD        :: cast(*u16) 32518;

App_State :: struct {
    hwnd: HWND;
    singbox_running: bool;
    singbox_process_handle: HANDLE;
    singbox_thread_handle: HANDLE;
    singbox_job_handle: HANDLE;

    updater_thread: Thread;
    updater_data: Updater_Thread_Data;
    use_system_proxy: bool = true;
}

// Custom log print sending logs to OutputDebugStringW
log_print :: (format_string: string, args: .. Any) {
    formatted := tprint(format_string, .. args);
    OutputDebugStringW(utf8_to_wide(formatted));
}

file_exists :: (path: string) -> bool {
    INVALID_FILE_ATTRIBUTES :: 0xFFFFFFFF;
    wide_path := utf8_to_wide(path);
    attribs := GetFileAttributesW(wide_path);
    return attribs != INVALID_FILE_ATTRIBUTES && !(attribs & FILE_ATTRIBUTE_DIRECTORY);
}

// Hidden Process Spawning with Job Object configuration and stdout/stderr redirection
create_process_hidden :: (cmd: string, log_file: string = "") -> HANDLE, PROCESS_INFORMATION, bool {
    startup_info: STARTUPINFOW;
    startup_info.cb = size_of(STARTUPINFOW);
    startup_info.dwFlags = STARTF_USESHOWWINDOW;
    startup_info.wShowWindow = SW_HIDE;

    file_handle: HANDLE = INVALID_HANDLE_VALUE;

    if log_file {
        sa: SECURITY_ATTRIBUTES;
        sa.nLength = size_of(SECURITY_ATTRIBUTES);
        sa.bInheritHandle = 1; // TRUE
        sa.lpSecurityDescriptor = null;

        wide_log := utf8_to_wide(log_file);

        file_handle = CreateFileW(
            wide_log,
            GENERIC_WRITE,
            FILE_SHARE_READ,
            *sa,
            CREATE_ALWAYS,
            FILE_ATTRIBUTE_NORMAL,
            null
        );

        if file_handle != INVALID_HANDLE_VALUE {
            startup_info.dwFlags |= STARTF_USESTDHANDLES;
            startup_info.hStdOutput = file_handle;
            startup_info.hStdError = file_handle;
        }
    }

    process_info: PROCESS_INFORMATION;

    CREATE_NO_WINDOW :: 0x08000000;

    job_handle := CreateJobObjectA(null, null);
    if job_handle {
        job_info: JOBOBJECT_EXTENDED_LIMIT_INFORMATION;
        job_info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
        SetInformationJobObject(job_handle, .ExtendedLimitInformation, *job_info, size_of(type_of(job_info)));
    }

    cmd_wide := utf8_to_wide(cmd);

    inherit_handles := ifx file_handle != INVALID_HANDLE_VALUE then cast(BOOL) 1 else cast(BOOL) 0;

    success := CreateProcessW(
        null,
        cmd_wide,
        null,
        null,
        inherit_handles,
        CREATE_NO_WINDOW,
        null,
        null,
        *startup_info,
        *process_info
    );

    if success && job_handle {
        AssignProcessToJobObject(job_handle, process_info.hProcess);
    }

    if file_handle != INVALID_HANDLE_VALUE {
        CloseHandle(file_handle);
    }

    return job_handle, process_info, cast(bool) success;
}

set_tray_tip :: (nid: *NOTIFYICONDATAW, text: string) {
    wide_text := utf8_to_wide(text);
    i := 0;
    while wide_text[i] != 0 && i < 127 {
        nid.szTip[i] = wide_text[i];
        i += 1;
    }
    nid.szTip[i] = 0;
}

update_tray :: (app: *App_State) {
    nid: NOTIFYICONDATAW;
    nid.cbSize = size_of(NOTIFYICONDATAW);
    nid.hWnd = app.hwnd;
    nid.uID = 1;
    nid.uFlags = NIF_ICON | NIF_TIP;

    if app.singbox_running {
        nid.hIcon = LoadIconW(null, IDI_SHIELD);
        set_tray_tip(*nid, "Sing-box: Running");
    } else {
        nid.hIcon = LoadIconW(null, IDI_APPLICATION);
        set_tray_tip(*nid, "Sing-box: Stopped");
    }

    Shell_NotifyIconW(NIM_MODIFY, *nid);
}

start_singbox :: (app: *App_State) -> bool {
    if app.singbox_running return true;

    if !file_exists("config.json") {
        if file_exists("url.txt") {
            changed, success, err_msg := perform_update();
            if !success {
                msg := tprint("Could not download configuration.\nError: %", err_msg);
                MessageBoxW(app.hwnd, utf8_to_wide(msg), utf8_to_wide("Sing-box Tray"), MB_ICONERROR);
                return false;
            }
        } else {
            MessageBoxW(app.hwnd, utf8_to_wide("No configuration found. Please configure the Config URL first."), utf8_to_wide("Sing-box Tray"), MB_ICONWARNING);
            return false;
        }
    }

    if file_exists("config.json") {
        config_data, read_ok := read_entire_file("config.json");
        if read_ok {
            modified := modify_config_inbounds(config_data);
            if modified != config_data {
                write_entire_file("config.json", modified);
                log_print("Existing config.json sanitized to use mixed inbound.\n");
            }
            free(config_data);
            free(modified);
        }
    }

    exe_path := "sing-box.exe";
    // Check if sing-box.exe is not in current folder, check in the sing-box subfolder
    if !file_exists("sing-box.exe") {
        if file_exists("sing-box/sing-box.exe") {
            exe_path = "sing-box\\sing-box.exe";
        }
    }

    cmd := tprint("% run -c config.json", exe_path);
    job, pi, ok := create_process_hidden(cmd, "sing-box.log");
    if !ok {
        msg := tprint("Failed to start %. Please ensure sing-box.exe is in the same folder or in the 'sing-box' subfolder.", exe_path);
        MessageBoxW(app.hwnd, utf8_to_wide(msg), utf8_to_wide("Sing-box Tray"), MB_ICONERROR);
        return false;
    }

    app.singbox_process_handle = pi.hProcess;
    app.singbox_thread_handle = pi.hThread;
    app.singbox_job_handle = job;
    app.singbox_running = true;

    if app.use_system_proxy {
        set_windows_system_proxy(true, "127.0.0.1:20122");
        log_print("System proxy enabled.\n");
    }

    update_tray(app);
    log_print("Sing-box started successfully.\n");
    return true;
}

stop_singbox :: (app: *App_State) {
    if !app.singbox_running return;

    TerminateProcess(app.singbox_process_handle, 0);
    CloseHandle(app.singbox_process_handle);
    CloseHandle(app.singbox_thread_handle);
    CloseHandle(app.singbox_job_handle);

    app.singbox_process_handle = null;
    app.singbox_thread_handle = null;
    app.singbox_job_handle = null;
    app.singbox_running = false;

    set_windows_system_proxy(false, "");
    log_print("System proxy disabled.\n");

    update_tray(app);
    log_print("Sing-box stopped.\n");
}

check_process_status :: (app: *App_State) {
    if !app.singbox_running return;

    exit_code: u32;
    success := GetExitCodeProcess(app.singbox_process_handle, *exit_code);
    STILL_ACTIVE :: 259;
    if success && exit_code != STILL_ACTIVE {
        log_print("Sing-box process exited unexpectedly with code %.\n", exit_code);

        CloseHandle(app.singbox_process_handle);
        CloseHandle(app.singbox_thread_handle);
        CloseHandle(app.singbox_job_handle);

        app.singbox_process_handle = null;
        app.singbox_thread_handle = null;
        app.singbox_job_handle = null;
        app.singbox_running = false;

        set_windows_system_proxy(false, "");
        log_print("System proxy disabled due to unexpected exit.\n");

        update_tray(app);
    }
}

trigger_immediate_update :: (app: *App_State) {
    changed, success, err_msg := perform_update();
    if success {
        if changed {
            log_print("Config updated, restarting Sing-box...\n");
            if app.singbox_running {
                stop_singbox(app);
                start_singbox(app);
            } else {
                start_singbox(app);
            }
            MessageBoxW(app.hwnd, utf8_to_wide("Configuration updated and Sing-box restarted successfully."), utf8_to_wide("Sing-box Tray"), MB_ICONINFORMATION);
        } else {
            if !app.singbox_running {
                start_singbox(app);
            }
            MessageBoxW(app.hwnd, utf8_to_wide("Configuration is already up-to-date."), utf8_to_wide("Sing-box Tray"), MB_ICONINFORMATION);
        }
    } else {
        msg := tprint("Failed to download configuration.\nError: %", err_msg);
        MessageBoxW(app.hwnd, utf8_to_wide(msg), utf8_to_wide("Sing-box Tray"), MB_ICONERROR);
    }
}

show_context_menu :: (app: *App_State) {
    hMenu := CreatePopupMenu();
    if !hMenu return;
    defer DestroyMenu(hMenu);

    status_str := ifx app.singbox_running then "Sing-box: Running" else "Sing-box: Stopped";
    AppendMenuW(hMenu, MF_STRING | MF_GRAYED | MF_DISABLED, CMD_STATUS, utf8_to_wide(status_str));

    AppendMenuW(hMenu, MF_SEPARATOR, 0, null);

    toggle_str := ifx app.singbox_running then "Stop Sing-box" else "Start Sing-box";
    AppendMenuW(hMenu, MF_STRING, CMD_START_STOP, utf8_to_wide(toggle_str));

    AppendMenuW(hMenu, MF_STRING, CMD_SET_URL, utf8_to_wide("Configure URL..."));
    AppendMenuW(hMenu, MF_STRING, CMD_UPDATE_NOW, utf8_to_wide("Update Config Now"));

    AppendMenuW(hMenu, MF_SEPARATOR, 0, null);

    proxy_flags := MF_STRING;
    sys_proxy_flags := MF_STRING;
    if app.use_system_proxy {
        sys_proxy_flags |= MF_CHECKED;
    } else {
        proxy_flags |= MF_CHECKED;
    }
    AppendMenuW(hMenu, proxy_flags, CMD_MODE_PROXY, utf8_to_wide("Proxy Mode (Local only)"));
    AppendMenuW(hMenu, sys_proxy_flags, CMD_MODE_SYS_PROXY, utf8_to_wide("System Proxy Mode"));

    AppendMenuW(hMenu, MF_SEPARATOR, 0, null);

    AppendMenuW(hMenu, MF_STRING, CMD_EXIT, utf8_to_wide("Exit"));

    cursor_pos: POINT;
    GetCursorPos(*cursor_pos);

    SetForegroundWindow(app.hwnd);

    track_flags := TPM_LEFTALIGN | TPM_RIGHTBUTTON | TPM_RETURNCMD;
    cmd := TrackPopupMenu(hMenu, track_flags, cursor_pos.x, cursor_pos.y, 0, app.hwnd, null);

    if cmd == {
        case CMD_START_STOP; {
            if app.singbox_running {
                stop_singbox(app);
            } else {
                start_singbox(app);
            }
        }
        case CMD_SET_URL; {
            changed := show_config_url_dialog(app.hwnd);
            if changed {
                log_print("URL configured, triggering download...\n");
                trigger_immediate_update(app);
            }
        }
        case CMD_UPDATE_NOW; {
            trigger_immediate_update(app);
        }
        case CMD_MODE_PROXY; {
            if app.use_system_proxy {
                app.use_system_proxy = false;
                write_entire_file("mode.txt", "proxy");
                log_print("Switched to Proxy Mode.\n");
                if app.singbox_running {
                    set_windows_system_proxy(false, "");
                    log_print("System proxy disabled.\n");
                }
            }
        }
        case CMD_MODE_SYS_PROXY; {
            if !app.use_system_proxy {
                app.use_system_proxy = true;
                write_entire_file("mode.txt", "system_proxy");
                log_print("Switched to System Proxy Mode.\n");
                if app.singbox_running {
                    set_windows_system_proxy(true, "127.0.0.1:20122");
                    log_print("System proxy enabled.\n");
                }
            }
        }
        case CMD_EXIT; {
            DestroyWindow(app.hwnd);
        }
    }

    PostMessageW(app.hwnd, WM_NULL, 0, 0);
}

app_wnd_proc :: (hwnd: HWND, msg: u32, wparam: WPARAM, lparam: LPARAM) -> LRESULT #c_call {
    ctx: #Context;
    push_context ctx {
        app := cast(*App_State) GetWindowLongPtrW(hwnd, GWLP_USERDATA);

        if msg == {
            case WM_CREATE; {
                create_struct := cast(*CREATESTRUCTW) lparam;
                app = cast(*App_State) create_struct.lpCreateParams;
                SetWindowLongPtrW(hwnd, GWLP_USERDATA, cast(LONG_PTR) app);
                return 0;
            }
            case WM_TRAY_CALLBACK; {
                if lparam == WM_RBUTTONUP || lparam == WM_LBUTTONUP {
                    show_context_menu(app);
                }
                return 0;
            }
            case WM_RESTART_SINGBOX; {
                log_print("WM_RESTART_SINGBOX received, restarting sing-box...\n");
                if app.singbox_running {
                    stop_singbox(app);
                    start_singbox(app);
                }
                return 0;
            }
            case WM_TIMER; {
                if wparam == TIMER_PROCESS_CHECK {
                    check_process_status(app);
                }
                return 0;
            }
            case WM_DESTROY; {
                PostQuitMessage(0);
                return 0;
            }
        }
        return DefWindowProcW(hwnd, msg, wparam, lparam);
    }
}

main :: () {
    hInstance := GetModuleHandleW(null);
    class_name := utf8_to_wide("SingboxTrayControllerClass");

    wclass: WNDCLASSEXW;
    wclass.cbSize = size_of(WNDCLASSEXW);
    wclass.lpfnWndProc = app_wnd_proc;
    wclass.hInstance = hInstance;
    wclass.hCursor = LoadCursorW(null, IDC_ARROW);
    wclass.lpszClassName = class_name;

    RegisterClassExW(*wclass);
    defer UnregisterClassW(class_name, hInstance);

    app_state: App_State;

    hwnd := CreateWindowExW(
        0,
        class_name,
        utf8_to_wide("Singbox Tray Controller"),
        0,
        0, 0, 0, 0,
        null,
        null,
        hInstance,
        *app_state
    );

    if !hwnd {
        log_print("Failed to create hidden window!\n");
        return;
    }

    app_state.hwnd = hwnd;

    // Load persisted mode
    mode_data, mode_ok := read_entire_file("mode.txt");
    if mode_ok {
        mode_str := trim(mode_data);
        if mode_str == "proxy" {
            app_state.use_system_proxy = false;
        }
        free(mode_data);
    }

    nid: NOTIFYICONDATAW;
    nid.cbSize = size_of(NOTIFYICONDATAW);
    nid.hWnd = hwnd;
    nid.uID = 1;
    nid.uFlags = NIF_MESSAGE | NIF_ICON | NIF_TIP;
    nid.uCallbackMessage = WM_TRAY_CALLBACK;
    nid.hIcon = LoadIconW(null, IDI_APPLICATION);
    set_tray_tip(*nid, "Sing-box: Stopped");

    if !Shell_NotifyIconW(NIM_ADD, *nid) {
        log_print("Failed to register tray icon!\n");
        return;
    }
    defer Shell_NotifyIconW(NIM_DELETE, *nid);

    if file_exists("config.json") {
        start_singbox(*app_state);
    }

    SetTimer(hwnd, TIMER_PROCESS_CHECK, 1000, null);
    defer KillTimer(hwnd, TIMER_PROCESS_CHECK);

    app_state.updater_data.hwnd = hwnd;
    app_state.updater_data.update_interval_seconds = 3600;
    app_state.updater_data.stop_event = CreateEventW(null, TRUE, FALSE, null);

    thread_init(*app_state.updater_thread, updater_thread_proc);
    app_state.updater_thread.data = *app_state.updater_data;
    thread_start(*app_state.updater_thread);

    msg: MSG;
    while GetMessageW(*msg, null, 0, 0) {
        TranslateMessage(*msg);
        DispatchMessageW(*msg);
    }

    log_print("Exiting application...\n");

    if app_state.updater_data.stop_event {
        SetEvent(app_state.updater_data.stop_event);
        thread_is_done(*app_state.updater_thread, -1);
        thread_deinit(*app_state.updater_thread);
        CloseHandle(app_state.updater_data.stop_event);
    }

    stop_singbox(*app_state);
}