/*
 * Copyright (c) 2022, Florent Castelli <florent.castelli@gmail.com>
 * Copyright (c) 2022, Sam Atkins <atkinssj@serenityos.org>
 * Copyright (c) 2022, Tobias Christiansen <tobyase@serenityos.org>
 * Copyright (c) 2022, Linus Groh <linusg@serenityos.org>
 * Copyright (c) 2022-2025, Tim Flynn <trflynn89@ladybird.org>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <AK/HashMap.h>
#include <AK/JsonObject.h>
#if !defined(AK_OS_MACOS)
#    include <LibCore/LocalServer.h>
#    include <LibCore/Socket.h>
#    include <LibCore/StandardPaths.h>
#else
#    include <LibIPC/TransportBootstrapMach.h>
#    include <LibWebView/Utilities.h>
#endif
#include <LibCore/System.h>
#include <LibIPC/Transport.h>
#include <LibWeb/Crypto/Crypto.h>
#include <LibWeb/WebDriver/Proxy.h>
#include <LibWeb/WebDriver/TimeoutsConfiguration.h>
#include <LibWeb/WebDriver/UserPrompt.h>
#include <WebDriver/Session.h>

namespace WebDriver {

static HashMap<String, NonnullRefPtr<Session>> s_sessions;
static HashMap<String, NonnullRefPtr<Session>> s_http_sessions;

// https://w3c.github.io/webdriver/#dfn-create-a-session
ErrorOr<NonnullRefPtr<Session>> Session::create(NonnullRefPtr<Client> client, JsonObject& capabilities, Web::WebDriver::SessionFlags flags)
{
    // 1. Let session id be the result of generating a UUID.
    auto session_id = Web::Crypto::generate_random_uuid();

    // 2. Let session be a new session with session ID session id, and HTTP flag flags contains "http".
    auto session = adopt_ref(*new Session(client, capabilities, move(session_id), flags));
    TRY(session->start(client->launch_browser_callback()));

    // 3. Let proxy be the result of getting property "proxy" from capabilities and run the substeps of the first matching statement:
    // -> proxy is a proxy configuration object
    if (auto proxy = capabilities.get_object("proxy"sv); proxy.has_value()) {
        // Take implementation-defined steps to set the user agent proxy using the extracted proxy configuration. If the
        // defined proxy cannot be configured return error with error code session not created. Otherwise set the has
        // proxy configuration flag to true.
        return Error::from_string_literal("Proxy configuration is not yet supported");
    }
    // -> Otherwise
    else {
        // Set a property of capabilities with name "proxy" and a value that is a new JSON Object.
        capabilities.set("proxy"sv, JsonObject {});
    }

    // FIXME: 4. If capabilites has a property named "acceptInsecureCerts", set the endpoint node's accept insecure TLS flag
    //           to the result of getting a property named "acceptInsecureCerts" from capabilities.

    // 5. Let user prompt handler capability be the result of getting property "unhandledPromptBehavior" from capabilities.
    auto user_prompt_handler_capability = capabilities.get_object("unhandledPromptBehavior"sv);

    // 6. If user prompt handler capability is not undefined, update the user prompt handler with user prompt handler capability.
    if (user_prompt_handler_capability.has_value())
        Web::WebDriver::update_the_user_prompt_handler(*user_prompt_handler_capability);

    session->web_content_connection().async_set_user_prompt_handler(Web::WebDriver::user_prompt_handler());

    // 7. Let serialized user prompt handler be serialize the user prompt handler.
    auto serialized_user_prompt_handler = Web::WebDriver::serialize_the_user_prompt_handler();

    // 8. Set a property on capabilities with the name "unhandledPromptBehavior", and the value serialized user prompt handler.
    capabilities.set("unhandledPromptBehavior"sv, move(serialized_user_prompt_handler));

    // 9. If flags contains "http":
    if (has_flag(flags, Web::WebDriver::SessionFlags::Http)) {
        // 1. Let strategy be the result of getting property "pageLoadStrategy" from capabilities. If strategy is a
        //    string, set the session's page loading strategy to strategy. Otherwise, set the page loading strategy to
        //    normal and set a property of capabilities with name "pageLoadStrategy" and value "normal".
        if (auto strategy = capabilities.get_string("pageLoadStrategy"sv); strategy.has_value()) {
            session->m_page_load_strategy = Web::WebDriver::page_load_strategy_from_string(*strategy);
            session->web_content_connection().async_set_page_load_strategy(session->m_page_load_strategy);
        } else {
            capabilities.set("pageLoadStrategy"sv, "normal"sv);
        }

        // 3. Let strictFileInteractability be the result of getting property "strictFileInteractability" from .
        //    capabilities. If strictFileInteractability is a boolean, set session's strict file interactability to
        //    strictFileInteractability.
        if (auto strict_file_interactiblity = capabilities.get_bool("strictFileInteractability"sv); strict_file_interactiblity.has_value()) {
            session->m_strict_file_interactiblity = *strict_file_interactiblity;
            session->web_content_connection().async_set_strict_file_interactability(session->m_strict_file_interactiblity);
        }

        // 4. Let timeouts be the result of getting a property "timeouts" from capabilities. If timeouts is not
        //    undefined, set session's session timeouts to timeouts.
        if (auto timeouts = capabilities.get_object("timeouts"sv); timeouts.has_value()) {
            MUST(session->set_timeouts(*timeouts));
        }

        // 5. Set a property on capabilities with name "timeouts" and value serialize the timeouts configuration with
        //    session's session timeouts.
        capabilities.set("timeouts"sv, session->m_timeouts_configuration.value_or_lazy_evaluated([]() {
            return Web::WebDriver::timeouts_object({});
        }));
    }

    // FIXME: 10. Process any extension capabilities in capabilities in an implementation-defined manner.

    // FIXME: 11. Run any WebDriver new session algorithm defined in external specifications, with arguments session, capabilities, and flags.

    // 12. Append session to active sessions.
    s_sessions.set(session->session_id(), session);

    // 13. If flags contains "http", append session to active HTTP sessions.
    if (has_flag(flags, Web::WebDriver::SessionFlags::Http))
        s_http_sessions.set(session->session_id(), session);

    // 14. Set the webdriver-active flag to true.
    session->web_content_connection().async_set_is_webdriver_active(true);

    return session;
}

Session::Session(NonnullRefPtr<Client> client, JsonObject const& capabilities, String session_id, Web::WebDriver::SessionFlags flags)
    : m_client(move(client))
    , m_options(capabilities)
    , m_session_id(move(session_id))
    , m_session_flags(flags)
    , m_event_loop(Core::EventLoop::current_weak())
{
}

Session::~Session() = default;

ErrorOr<NonnullRefPtr<Session>, Web::WebDriver::Error> Session::find_session(StringView session_id, Web::WebDriver::SessionFlags session_flags, AllowInvalidWindowHandle allow_invalid_window_handle)
{
    auto& sessions = has_flag(session_flags, Web::WebDriver::SessionFlags::Http) ? s_http_sessions : s_sessions;

    if (auto session = sessions.get(session_id); session.has_value()) {
        if (allow_invalid_window_handle == AllowInvalidWindowHandle::No)
            TRY(session.value()->ensure_current_window_handle_is_valid());

        return *session.release_value();
    }

    return Web::WebDriver::Error::from_code(Web::WebDriver::ErrorCode::InvalidSessionId, "Invalid session id"sv);
}

size_t Session::session_count(Web::WebDriver::SessionFlags session_flags)
{
    if (has_flag(session_flags, Web::WebDriver::SessionFlags::Http))
        return s_http_sessions.size();
    return s_sessions.size();
}

// https://w3c.github.io/webdriver/#dfn-close-the-session
void Session::close()
{
    // 1. If session's HTTP flag is set, remove session from active HTTP sessions.
    if (has_flag(session_flags(), Web::WebDriver::SessionFlags::Http))
        s_http_sessions.remove(m_session_id);

    // 2. Remove session from active sessions.
    s_sessions.remove(m_session_id);

    // 3. Perform the following substeps based on the remote end's type:
    // -> Remote end is an endpoint node
    //     1. If the list of active sessions is empty:
    if (s_sessions.is_empty()) {
        // 1. Set the webdriver-active flag to false
        // NOTE: This is handled by the WebContent process.

        // 2. Set the user prompt handler to null.
        Web::WebDriver::set_user_prompt_handler({});

        // FIXME: 3. Unset the accept insecure TLS flag.

        // 4. Reset the has proxy configuration flag to its default value.
        Web::WebDriver::reset_has_proxy_configuration();

        // 5. Optionally, close all top-level browsing contexts, without prompting to unload.
        for (auto& it : m_windows)
            it.value.web_content_connection->close_session();
    }
    // -> Remote end is an intermediary node
    //     1. Close the associated session. If this causes an error to occur, complete the remainder of this algorithm
    //        before returning the error.

    // 4. Perform any implementation-specific cleanup steps.
    for (auto& [_, connection] : m_pending_connections) {
        connection->on_close = nullptr;
        connection->on_did_set_window_handle = nullptr;
    }
    m_pending_connections.clear();

    if (m_browser_process.has_value())
        MUST(Core::System::kill(m_browser_process->pid(), SIGTERM));

#if defined(AK_OS_MACOS)
    m_web_content_mach_port_server = nullptr;
#else
    if (!m_web_content_endpoint.is_empty())
        MUST(Core::System::unlink(m_web_content_endpoint));
#endif
    m_web_content_endpoint = {};

    // 5. If an error has occurred in any of the steps above, return the error, otherwise return success with data null.
}

ErrorOr<void> Session::accept_web_content_transport(NonnullOwnPtr<IPC::Transport> transport, NonnullRefPtr<ServerPromise> promise)
{
    auto web_content_connection = TRY(adopt_nonnull_ref_or_enomem(new (nothrow) WebContentConnection(move(transport))));
    dbgln("WebDriver is connected to WebContent");

    auto connection_id = m_next_pending_connection_id++;
    // Publish the connection before the initial did_set_window_handle message can race in.
    m_pending_connections.set(connection_id, web_content_connection);

    web_content_connection->on_close = [this, promise, connection_id]() {
        if (m_pending_connections.remove(connection_id)) {
            dbgln_if(WEBDRIVER_DEBUG, "Pending connection {} closed before sending its handle", connection_id);
            promise->reject(Error::from_string_literal("Window was closed before sending its handle"));
        }
    };

    web_content_connection->on_did_set_window_handle = [this, promise, connection_id](String window_handle) {
        auto maybe_pending_connection = m_pending_connections.take(connection_id);
        if (!maybe_pending_connection.has_value())
            return;

        auto pending_connection = maybe_pending_connection.value();
        pending_connection->on_did_set_window_handle = nullptr;

        dbgln_if(WEBDRIVER_DEBUG, "Window {} registered with WebDriver.", window_handle);

        pending_connection->on_close = [this, window_handle]() {
            dbgln_if(WEBDRIVER_DEBUG, "Window {} was closed remotely.", window_handle);
            m_windows.remove(window_handle);
            if (m_windows.is_empty())
                close();
        };

        pending_connection->async_set_page_load_strategy(m_page_load_strategy);
        pending_connection->async_set_strict_file_interactability(m_strict_file_interactiblity);
        pending_connection->async_set_user_prompt_handler(Web::WebDriver::user_prompt_handler());
        if (m_timeouts_configuration.has_value())
            pending_connection->async_set_timeouts(*m_timeouts_configuration);

        m_windows.set(window_handle, Session::Window { window_handle, move(pending_connection) });

        if (m_current_window_handle.is_empty())
            m_current_window_handle = window_handle;

        promise->resolve({});
    };
    return {};
}

ErrorOr<void> Session::create_server(NonnullRefPtr<ServerPromise> promise)
{
#if defined(AK_OS_WINDOWS)
    static_assert(IsSame<IPC::Transport, IPC::TransportSocketWindows>, "Need to handle other IPC transports here");
#elif defined(AK_OS_MACOS)
    static_assert(IsSame<IPC::Transport, IPC::TransportMachPort>, "Need to handle other IPC transports here");
#else
    static_assert(IsSame<IPC::Transport, IPC::TransportSocket>, "Need to handle other IPC transports here");
#endif

    dbgln("Listening for WebDriver connection on {}", m_web_content_endpoint);

#if defined(AK_OS_MACOS)
    m_web_content_mach_port_server = make<IPC::MachBootstrapListener>(m_web_content_endpoint);
    if (!m_web_content_mach_port_server->is_initialized())
        return Error::from_string_literal("Failed to initialize Mach port server for WebDriver");

    m_web_content_mach_port_server->on_bootstrap_request = [this, promise](auto request) {
        auto result = m_transport_bootstrap_server.handle_bootstrap_request(request.pid, move(request.reply_port));
        if (result.is_error()) {
            auto event_loop = m_event_loop->take();
            VERIFY(event_loop);
            event_loop->deferred_invoke([promise, error = result.release_error()]() mutable {
                promise->resolve(move(error));
            });
            return;
        }

        result.release_value().visit(
            [](IPC::TransportBootstrapMachServer::ChildTransportHandled) {
                VERIFY_NOT_REACHED();
            },
            [this, promise](IPC::TransportBootstrapMachServer::OnDemandTransport& transport) {
                auto event_loop = m_event_loop->take();
                VERIFY(event_loop);
                event_loop->deferred_invoke([this, promise, transport = move(transport.ports)]() mutable {
                    if (auto result = accept_web_content_transport(make<IPC::Transport>(move(transport.receive_right), move(transport.send_right)), promise); result.is_error())
                        promise->resolve(result.release_error());
                });
            });
    };

    return {};
#else
    (void)Core::System::unlink(m_web_content_endpoint);

    auto server = Core::LocalServer::construct();
    server->listen(m_web_content_endpoint);

    server->on_accept = [this, promise](auto client_socket) {
        auto maybe_transport = IPC::Transport::from_socket(move(client_socket));
        if (maybe_transport.is_error()) {
            promise->resolve(maybe_transport.release_error());
            return;
        }
        if (auto result = accept_web_content_transport(maybe_transport.release_value(), promise); result.is_error())
            promise->resolve(result.release_error());
    };

    server->on_accept_error = [promise](auto error) {
        promise->resolve(move(error));
    };

    m_web_content_server = server;
    return {};
#endif
}

ErrorOr<void> Session::start(LaunchBrowserCallback const& launch_browser_callback)
{
    auto promise = ServerPromise::construct();

#if defined(AK_OS_MACOS)
    m_web_content_endpoint = ByteString::formatted("{}.{}", WebView::mach_server_name_for_process("WebDriver"sv, Core::System::getpid()), m_session_id);
#else
    m_web_content_endpoint = ByteString::formatted("{}/webdriver/session_{}_{}", TRY(Core::StandardPaths::runtime_directory()), Core::System::getpid(), m_session_id);
#endif
    TRY(create_server(promise));

    m_browser_process = TRY(launch_browser_callback(m_web_content_endpoint, m_options.headless));

    // FIXME: Allow this to be more asynchronous. For now, this at least allows us to propagate
    //        errors received while accepting the Browser and WebContent sockets.
    TRY(TRY(promise->await()));

    return {};
}

Web::WebDriver::Response Session::set_timeouts(JsonValue payload)
{
    m_timeouts_configuration = TRY(web_content_connection().set_timeouts(move(payload)));
    return JsonValue {};
}

// 11.2 Close Window, https://w3c.github.io/webdriver/#dfn-close-window
Web::WebDriver::Response Session::close_window()
{
    // 3. Close the current top-level browsing context.
    TRY(perform_async_action([&](auto& connection) {
        return connection.close_window();
    }));

    {
        // Defer removing the window handle from this session until after we know we are done with its connection.
        ScopeGuard guard { [this] { m_windows.remove(m_current_window_handle); m_current_window_handle = "NoSuchWindowPleaseSelectANewOne"_string; } };

        // 4. If there are no more open top-level browsing contexts, then close the session.
        if (m_windows.size() == 1)
            close();
    }

    // 5. Return the result of running the remote end steps for the Get Window Handles command.
    return get_window_handles();
}

// 11.3 Switch to Window, https://w3c.github.io/webdriver/#dfn-switch-to-window
Web::WebDriver::Response Session::switch_to_window(StringView handle)
{
    // 4. If handle is equal to the associated window handle for some top-level browsing context, let context be the that
    //    browsing context, and set the current top-level browsing context with session and context.
    //    Otherwise, return error with error code no such window.
    if (auto it = m_windows.find(handle); it != m_windows.end())
        m_current_window_handle = it->key;
    else
        return Web::WebDriver::Error::from_code(Web::WebDriver::ErrorCode::NoSuchWindow, "Window not found"sv);

    // 5. Update any implementation-specific state that would result from the user selecting the current
    //    browsing context for interaction, without altering OS-level focus.
    TRY(web_content_connection().switch_to_window(m_current_window_handle));

    // 6. Return success with data null.
    return JsonValue {};
}

// 11.4 Get Window Handles, https://w3c.github.io/webdriver/#dfn-get-window-handles
Web::WebDriver::Response Session::get_window_handles() const
{
    // 1. Let handles be a JSON List.
    JsonArray handles {};

    // 2. For each top-level browsing context in the remote end, push the associated window handle onto handles.
    for (auto const& window_handle : m_windows.keys()) {
        handles.must_append(JsonValue(window_handle));
    }

    // 3. Return success with data handles.
    return JsonValue { move(handles) };
}

ErrorOr<void, Web::WebDriver::Error> Session::ensure_current_window_handle_is_valid() const
{
    if (auto current_window = m_windows.get(m_current_window_handle); current_window.has_value())
        return {};
    return Web::WebDriver::Error::from_code(Web::WebDriver::ErrorCode::NoSuchWindow, "Window not found"sv);
}

}