#include <cstdlib>
#include <cstring>
#include <iostream>
#include <getopt.h>
#include <signal.h>
#include <fcntl.h>

#include <unistd.h>
#include <wayfire/debug.hpp>
#include "main.hpp"

#include <wayland-server.h>

#include "core/opengl-priv.hpp"
#include "wayfire/config-backend.hpp"
#include "core/plugin-loader.hpp"
#include "core/core-impl.hpp"
#include <wayfire/nonstd/wlroots.hpp>

static std::string get_version_string()
{
    return std::string(WAYFIRE_VERSION) + "-" + wf::version::git_commit +
           " (" + __DATE__ + ") branch " + wf::version::git_branch +
           " wlroots-" + WLR_VERSION_STR;
}

static void print_version_and_exit()
{
    std::cout << get_version_string() << std::endl;
    exit(0);
}

static void print_help()
{
    std::cout << "Wayfire: " << get_version_string() << std::endl;
    std::cout << "Usage: wayfire [OPTION]...\n" << std::endl;
    std::cout << " -c,  --config            specify config file to use " <<
        "(overrides WAYFIRE_CONFIG_FILE from the environment)" << std::endl;
    std::cout << " -B,  --config-backend    specify config backend to use" <<
        std::endl;
    std::cout << " -h,  --help              print this help" << std::endl;
    std::cout << " -d,  --debug             enable debug logging" << std::endl;
    std::cout <<
        " -D,  --damage-debug      enable additional debug for damaged regions" <<
        std::endl;
    std::cout << " -R,  --damage-rerender   rerender damaged regions" << std::endl;
    std::cout << " -l,  --legacy-wl-drm     use legacy drm for wayland clients" << std::endl;
    std::cout << " -v,  --version           print version and exit" << std::endl;
    exit(0);
}

static bool drop_permissions(void)
{
    if ((getuid() != geteuid()) || (getgid() != getegid()))
    {
        // Set the gid and uid in the correct order.
        if ((setgid(getgid()) != 0) || (setuid(getuid()) != 0))
        {
            LOGE("Unable to drop root, refusing to start");

            return false;
        }
    }

    if ((setgid(0) != -1) || (setuid(0) != -1))
    {
        LOGE("Unable to drop root (we shouldn't be able to "
             "restore it after setuid), refusing to start");

        return false;
    }

    return true;
}

static void wlr_log_handler(wlr_log_importance level,
    const char *fmt, va_list args)
{
    const int bufsize = 4 * 1024;
    char buffer[bufsize];
    vsnprintf(buffer, bufsize, fmt, args);

    wf::log::log_level_t wlevel;
    switch (level)
    {
      case WLR_ERROR:
        wlevel = wf::log::LOG_LEVEL_ERROR;
        break;

      case WLR_INFO:
        wlevel = wf::log::LOG_LEVEL_INFO;
        break;

      case WLR_DEBUG:
        wlevel = wf::log::LOG_LEVEL_DEBUG;
        break;

      default:
        return;
    }

    bool enabled =
        wf::log::enabled_categories.test((int)wf::log::logging_category::WLR);
    if ((wlevel == wf::log::LOG_LEVEL_DEBUG) && !enabled)
    {
        return;
    }

    wf::log::log_plain(wlevel, buffer);
}

static std::optional<int> exit_because_signal;
static void signal_handler(int signal)
{
    std::string error;
    switch (signal)
    {
      case SIGSEGV:
        error = "Segmentation fault";
        break;

      case SIGFPE:
        error = "Floating-point exception";
        break;

      case SIGABRT:
        error = "Fatal error(SIGABRT)";
        break;

      case SIGINT:
        exit_because_signal = SIGINT;
        wf::get_core().shutdown();
        return;

      case SIGTERM:
        exit_because_signal = SIGTERM;
        wf::get_core().shutdown();
        return;

      default:
        error = "Unknown";
    }

    LOGE("Fatal error: ", error);
    wf::print_trace(false);
    std::_Exit(-1);
}

static std::optional<std::string> choose_socket(wl_display *display)
{
    for (int i = 1; i <= 32; i++)
    {
        auto name = "wayland-" + std::to_string(i);
        if (wl_display_add_socket(display, name.c_str()) >= 0)
        {
            return name;
        }
    }

    return {};
}

static wf::config_backend_t *load_backend(const std::string& backend)
{
    std::string config_plugin(backend);

    if (backend.size())
    {
        std::vector<std::string> plugin_prefixes = wf::get_plugin_paths();
        config_plugin =
            wf::get_plugin_path_for_name(plugin_prefixes, backend).value_or("");
    }

    auto [_, init_ptr] = wf::get_new_instance_handle(config_plugin, false);

    if (!init_ptr)
    {
        return nullptr;
    }

    using backend_init_t = wf::config_backend_t * (*)();
    auto init = wf::union_cast<void*, backend_init_t>(init_ptr);
    return init();
}

static std::string_view get_category_name(wf::log::logging_category category)
{
    switch (category)
    {
      case wf::log::logging_category::TXN:
        return "txn";

      case wf::log::logging_category::TXNI:
        return "txni";

      case wf::log::logging_category::VIEWS:
        return "views";

      case wf::log::logging_category::WLR:
        return "wlroots";

      case wf::log::logging_category::SCANOUT:
        return "scanout";

      case wf::log::logging_category::POINTER:
        return "pointer";

      case wf::log::logging_category::WSET:
        return "wset";

      case wf::log::logging_category::KBD:
        return "kbd";

      case wf::log::logging_category::XWL:
        return "xwayland";

      case wf::log::logging_category::LSHELL:
        return "layer-shell";

      case wf::log::logging_category::IM:
        return "im";

      case wf::log::logging_category::RENDER:
        return "render";

      case wf::log::logging_category::INPUT_DEVICES:
        return "input-devices";

      case wf::log::logging_category::OUTPUT:
        return "output";

      default:
        wf::dassert(false);
        return "unknown";
    }
}

void parse_extended_debugging(const std::vector<std::string>& categories)
{
    for (const auto& cat : categories)
    {
        size_t idx = 0;
        const size_t total = (size_t)wf::log::logging_category::TOTAL;
        for (; idx < total; idx++)
        {
            if (get_category_name((wf::log::logging_category)idx) == cat)
            {
                LOGD("Enabling debugging category \"", cat, "\"");
                wf::log::enabled_categories.set(idx, 1);
                break;
            }
        }

        if (idx == total)
        {
            LOGW("Unknown debugging category \"", cat, "\"");
        }
    }
}

// Override assert() handler to be more useful and print a trace.
// extern "C" {
// void __assert_fail(const char* expr, const char *filename, unsigned int line, const char *assert_func)
// {
// LOGE("Assertion failed: ", expr, " at ", filename, ":", line, " in ", assert_func);
// wf::print_trace(false);
// std::_Exit(-1);
// }
// }

// Override exception handling to be more useful and print a trace.
// extern "C" {
// void __cxa_throw(void *ex, void *info, void (*dest)(void*))
// {
// wf::print_trace(false);
// std::_Exit(-1);
// }
// }
//
int main(int argc, char *argv[])
{
    wf::log::log_level_t log_level = wf::log::LOG_LEVEL_INFO;
    struct option opts[] = {
        {
            "config", required_argument, NULL, 'c'
        },
        {
            "config-backend", required_argument, NULL, 'B'
        },
        {"debug", optional_argument, NULL, 'd'},
        {"damage-debug", no_argument, NULL, 'D'},
        {"damage-rerender", no_argument, NULL, 'R'},
        {"legacy-wl-drm", no_argument, NULL, 'l'},
        {"with-great-power-comes-great-responsibility", no_argument, NULL, 'r'},
        {"help", no_argument, NULL, 'h'},
        {"version", no_argument, NULL, 'v'},
        {"exit-on-gles-error", no_argument, NULL, '$'},
        {0, 0, NULL, 0}
    };

    std::string config_file;
    std::string config_backend = WF_DEFAULT_CONFIG_BACKEND;
    std::vector<std::string> extended_debug_categories;
    bool allow_root = false;

    if (char *default_config_backend = getenv("WAYFIRE_DEFAULT_CONFIG_BACKEND"))
    {
        config_backend = default_config_backend;
    }

    if (!getenv("XDG_CURRENT_DESKTOP"))
    {
        setenv("XDG_CURRENT_DESKTOP", "wayfire", 1);
    }

    int c, i;
    while ((c = getopt_long(argc, argv, "c:B:d::DhRlrv", opts, &i)) != -1)
    {
        switch (c)
        {
          case 'c':
            config_file = optarg;
            break;

          case 'B':
            config_backend = optarg;
            break;

          case 'D':
            runtime_config.damage_debug = true;
            break;

          case 'R':
            runtime_config.no_damage_track = true;
            break;

          case 'l':
            runtime_config.legacy_wl_drm = true;
            break;

          case 'r':
            allow_root = true;
            break;

          case 'h':
            print_help();
            break;

          case '$':
            OpenGL::exit_on_gles_error = true;
            break;

          case 'd':
            log_level = wf::log::LOG_LEVEL_DEBUG;

            // Make sure to parse things like `-d txn`, which getopt does not.
            // According to documentation, for optional arguments, optarg will
            // be NULL so we need to manually check.
            if (!optarg && (NULL != argv[optind]) &&
                ('-' != argv[optind][0]))
            {
                optarg = argv[optind];
                ++optind;
            }

            if (optarg)
            {
                extended_debug_categories.push_back(optarg);
            }

            break;

          case 'v':
            print_version_and_exit();
            break;

          default:
            std::cerr << "Unrecognized command line argument " << optarg << "\n" <<
                std::endl;
        }
    }

    /* Don't crash on SIGPIPE, e.g., when doing IPC to a client whose fd has been closed. */
    signal(SIGPIPE, SIG_IGN);

    wf::log::initialize_logging(std::cout, log_level, wf::detect_color_mode());

    parse_extended_debugging(extended_debug_categories);
    wlr_log_init(WLR_DEBUG, wlr_log_handler);

#ifdef PRINT_TRACE
    /* In case of crash, print the stacktrace for debugging.
     * However, if ASAN is enabled, we'll get better stacktrace from there. */
    signal(SIGSEGV, signal_handler);
    signal(SIGFPE, signal_handler);
    signal(SIGABRT, signal_handler);
#endif

    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);

    std::set_terminate([] ()
    {
        std::cout << "Unhandled exception" << std::endl;
        wf::print_trace(false);
        std::abort();
    });

    LOGI("Starting wayfire: ", get_version_string());
    /* First create display and initialize safe-list's event loop, so that
     * wf objects (which depend on safe-list) can work */
    auto display = wl_display_create();
    auto& core   = wf::compositor_core_impl_t::allocate_core();

    core.argc = argc;
    core.argv = argv;

    /** TODO: move this to core_impl constructor */
    core.display = display;
    core.ev_loop = wl_display_get_event_loop(core.display);
    core.backend = wlr_backend_autocreate(core.ev_loop, &core.session);

    int drm_fd = -1;
    char *drm_device = getenv("WLR_RENDER_DRM_DEVICE");
    if (drm_device)
    {
        drm_fd = open(drm_device, O_RDWR | O_CLOEXEC);
    } else
    {
        drm_fd = wlr_backend_get_drm_fd(core.backend);
    }

    if (drm_fd < 0)
    {
#if WLR_HAS_UDMABUF_ALLOCATOR == 1
        LOGW("Failed to open DRM render device, consider specifying WLR_RENDER_DRM_DEVICE."
             "Trying SW rendering instead.");
#else
        LOGE("Failed to open DRM render device, consider specifying WLR_RENDER_DRM_DEVICE."
             "If you want to use software rendering, ensure that wlroots has been compiled with udmabuf "
             "allocator support (available in wlroots >= 0.19.0) and recompile Wayfire.");
        wl_display_destroy_clients(core.display);
        wl_display_destroy(core.display);
        return EXIT_FAILURE;
#endif
    }

    // core.renderer = wlr_vk_renderer_create_with_drm_fd(drm_fd);
    core.renderer = wlr_renderer_autocreate(core.backend);
    // core.renderer = wlr_pixman_renderer_create();
    // core.renderer = wlr_gles2_renderer_create_with_drm_fd(drm_fd);
    if (!core.renderer)
    {
        LOGE("Failed to create renderer");
        wl_display_destroy_clients(core.display);
        wl_display_destroy(core.display);
        return EXIT_FAILURE;
    }

    core.allocator = wlr_allocator_autocreate(core.backend, core.renderer);
    assert(core.allocator);

    if (core.is_gles2())
    {
        core.egl = wlr_gles2_renderer_get_egl(core.renderer);
        assert(core.egl);
    }

    if (!allow_root && !drop_permissions())
    {
        wl_display_destroy_clients(core.display);
        wl_display_destroy(core.display);

        return EXIT_FAILURE;
    }

    auto backend = load_backend(config_backend);
    if (!backend)
    {
        LOGE("Failed to load configuration backend!");
        wl_display_destroy_clients(core.display);
        wl_display_destroy(core.display);
        return EXIT_FAILURE;
    }

    LOGD("Using configuration backend: ", config_backend);
    core.config_backend = std::unique_ptr<wf::config_backend_t>(backend);
    core.config_backend->init(display, *core.config, config_file);
    core.init();

    auto socket = choose_socket(core.display);
    if (!socket)
    {
        LOGE("Failed to create wayland socket, exiting.");

        return -1;
    }

    core.wayland_display = socket.value();
    LOGI("Using socket name ", core.wayland_display);
    if (!wlr_backend_start(core.backend))
    {
        LOGE("Failed to initialize backend, exiting");
        wlr_backend_destroy(core.backend);
        wl_display_destroy(core.display);

        return -1;
    }

    setenv("WAYLAND_DISPLAY", core.wayland_display.c_str(), 1);
    core.post_init();

    wl_display_run(core.display);
    if (exit_because_signal == SIGINT)
    {
        LOGI("Got SIGINT, shutting down");
    } else if (exit_because_signal == SIGTERM)
    {
        LOGI("Got SIGTERM, shutting down");
    }

    wf::compositor_core_impl_t::deallocate_core();
    LOGI("Shutdown successful!");
    return EXIT_SUCCESS;
}
