// Daemon shutdown behavior integration tests
// Validates graceful/force shutdown, timing, idempotency, and in-flight operations

#define CATCH_CONFIG_MAIN
#include <spdlog/spdlog.h>
#include <atomic>
#include <cstring>
#include <future>
#include <latch>
#include <thread>
#include <vector>
#include "test_async_helpers.h"
#include "test_daemon_harness.h"
#include <boost/asio/post.hpp>
#include <catch2/catch_session.hpp>
#include <catch2/catch_test_macros.hpp>
#include <yams/compat/unistd.h>
#include <yams/daemon/client/daemon_client.h>
#include <yams/daemon/components/IOCoordinator.h>
#include <yams/daemon/components/ServiceManager.h>
#include <yams/daemon/components/SocketServer.h>

#ifndef _WIN32
#include <sys/socket.h>
#include <sys/un.h>
#endif

using namespace yams::daemon;
using namespace yams::test;
using namespace std::chrono_literals;

namespace {
DaemonClient createClient(const std::filesystem::path& socketPath,
                          std::chrono::milliseconds connectTimeout = 2s) {
    ClientConfig config;
    config.socketPath = socketPath;
    config.connectTimeout = connectTimeout;
    config.autoStart = false;
    config.requestTimeout = 5s;
    return DaemonClient(config);
}

// Helper to connect with retry logic for GlobalIOContext stabilization
bool connectWithRetry(DaemonClient& client, int maxRetries = 3,
                      std::chrono::milliseconds retryDelay = 200ms) {
    // Initial delay for GlobalIOContext threads to stabilize
    std::this_thread::sleep_for(200ms);

    for (int attempt = 0; attempt < maxRetries; ++attempt) {
        auto result = yams::cli::run_sync(client.connect(), 5s);
        if (result.has_value()) {
            return true;
        }
        spdlog::warn("[TEST] Connect attempt {} failed: {}, retrying...", attempt + 1,
                     result.error().message);
        std::this_thread::sleep_for(retryDelay);
    }
    return false;
}

void startHarnessWithRetry(DaemonHarness& harness, std::size_t maxRetries = 3) {
    if (!harness.startWithRetry(30s, maxRetries)) {
        SKIP("Skipping shutdown integration section due to daemon startup instability");
    }
}

bool waitForLifecycleState(const YamsDaemon& daemon, LifecycleState state,
                           std::chrono::milliseconds timeout) {
    const auto deadline = std::chrono::steady_clock::now() + timeout;
    while (std::chrono::steady_clock::now() < deadline) {
        if (daemon.getLifecycle().snapshot().state == state) {
            return true;
        }
        if (state == LifecycleState::Ready) {
            if (const auto* serviceManager = daemon.getServiceManager();
                serviceManager && serviceManager->getServiceManagerFsmSnapshot().state ==
                                      ServiceManagerState::Ready) {
                return true;
            }
        }
        std::this_thread::sleep_for(10ms);
    }
    if (daemon.getLifecycle().snapshot().state == state) {
        return true;
    }
    if (state == LifecycleState::Ready) {
        if (const auto* serviceManager = daemon.getServiceManager();
            serviceManager &&
            serviceManager->getServiceManagerFsmSnapshot().state == ServiceManagerState::Ready) {
            return true;
        }
    }
    return false;
}

bool waitForDaemonToStop(const YamsDaemon& daemon, std::chrono::milliseconds timeout) {
    const auto deadline = std::chrono::steady_clock::now() + timeout;
    while (std::chrono::steady_clock::now() < deadline) {
        if (!daemon.isRunning()) {
            return true;
        }
        std::this_thread::sleep_for(10ms);
    }
    return !daemon.isRunning();
}
} // namespace

TEST_CASE("Daemon shutdown timing", "[daemon][shutdown][timing]") {
    SKIP_DAEMON_TEST_ON_WINDOWS();

    DaemonHarness harness;
    startHarnessWithRetry(harness);

    auto client = createClient(harness.socketPath());
    REQUIRE(connectWithRetry(client));

    SECTION("shutdown request returns immediately") {
        auto start = std::chrono::steady_clock::now();
        auto result = yams::cli::run_sync(client.shutdown(true), 5s);
        auto elapsed = std::chrono::steady_clock::now() - start;

        REQUIRE(result.has_value());
        REQUIRE((elapsed < 500ms));
    }

    SECTION("shutdown request still works after restarting the same daemon instance") {
        namespace fs = std::filesystem;

        harness.stop();
        yams::daemon::AsioConnectionPool::shutdown_all(500ms);
#ifndef _WIN32
        yams::daemon::GlobalIOContext::safe_restart();
#endif

        const auto unique =
            std::to_string(std::chrono::steady_clock::now().time_since_epoch().count());
        const auto root = fs::temp_directory_path() / ("yams_shutdown_restart_" + unique);
        const auto dataDir = root / "data";
        const auto pidFile = root / "daemon.pid";
        const auto logFile = root / "daemon.log";
#ifdef _WIN32
        const auto socketPath =
            fs::temp_directory_path() / ("yams_shutdown_restart_" + unique + ".sock");
#else
        const auto socketPath = fs::path("/tmp") / ("yams_shutdown_restart_" + unique + ".sock");
#endif
        fs::create_directories(dataDir);

        struct Cleanup {
            fs::path root;
            fs::path socketPath;
            ~Cleanup() {
                std::error_code ec;
                fs::remove(socketPath, ec);
                fs::remove_all(root, ec);
            }
        } cleanup{root, socketPath};

        DaemonConfig cfg;
        cfg.dataDir = dataDir;
        cfg.socketPath = socketPath;
        cfg.pidFile = pidFile;
        cfg.logFile = logFile;
        cfg.enableModelProvider = false;
        cfg.useMockModelProvider = false;
        cfg.autoLoadPlugins = false;

        YamsDaemon daemon(cfg);
        std::thread runLoopThread;
        auto cleanupCycle = [&]() {
            if (runLoopThread.joinable()) {
                runLoopThread.join();
            }
        };
        auto startCycle = [&]() {
            auto started = daemon.start();
            INFO("daemon.start() error: " << (started ? std::string{} : started.error().message));
            REQUIRE(started.has_value());
            runLoopThread = std::thread([&daemon]() {
                daemon.runLoop();
                (void)daemon.stop();
            });
            const bool reachedUsableState =
                waitForLifecycleState(daemon, LifecycleState::Ready, 15s) ||
                waitForLifecycleState(daemon, LifecycleState::Degraded, 1s);
            INFO("restart cycle reached lifecycle-usable state: " << reachedUsableState);
        };

        startCycle();
        {
            auto cycleClient = createClient(socketPath);
            REQUIRE(connectWithRetry(cycleClient));
            auto result = yams::cli::run_sync(cycleClient.shutdown(true), 5s);
            REQUIRE(result.has_value());
            REQUIRE(waitForDaemonToStop(daemon, 10s));
        }
        cleanupCycle();

        yams::daemon::AsioConnectionPool::shutdown_all(500ms);
#ifndef _WIN32
        yams::daemon::GlobalIOContext::safe_restart();
#endif
        std::this_thread::sleep_for(200ms);

        startCycle();
        {
            auto cycleClient = createClient(socketPath);
            REQUIRE(connectWithRetry(cycleClient));
            auto result = yams::cli::run_sync(cycleClient.shutdown(true), 5s);
            REQUIRE(result.has_value());
            CHECK(waitForDaemonToStop(daemon, 5s));
        }

        if (daemon.isRunning()) {
            (void)daemon.stop();
        }
        cleanupCycle();
    }
}

TEST_CASE("Daemon shutdown with in-flight operations", "[daemon][shutdown][operations]") {
    SKIP_DAEMON_TEST_ON_WINDOWS();

    DaemonHarness::Options opts;
    opts.enableModelProvider = false;
    opts.useMockModelProvider = false;
    DaemonHarness harness(opts);
    startHarnessWithRetry(harness);

    auto client = createClient(harness.socketPath());
    REQUIRE(connectWithRetry(client));

    SECTION("shutdown accepts requests to completion") {
        std::atomic<bool> operationFinished{false};
        std::atomic<bool> operationComplete{false};
        std::atomic<bool> operationCancelledByShutdown{false};
        std::atomic<bool> operationStarted{false};

        std::thread bgThread([&]() {
            AddDocumentRequest req;
            req.name = "test_doc.txt";
            req.content = "Test content";

            operationStarted = true;
            auto result = yams::cli::run_sync(client.streamingAddDocument(req), 10s);
            operationComplete = result.has_value();
            if (!result.has_value()) {
                const std::string message = result.error().message;
                operationCancelledByShutdown = (message.find("shutdown") != std::string::npos) ||
                                               (message.find("not running") != std::string::npos) ||
                                               (message.find("cancel") != std::string::npos);
            }
            operationFinished = true;
        });

        while (!operationStarted) {
            std::this_thread::sleep_for(10ms);
        }
        std::this_thread::sleep_for(50ms);

        auto shutdownResult = yams::cli::run_sync(client.shutdown(true), 5s);
        REQUIRE(shutdownResult.has_value());

        bgThread.join();

        REQUIRE(operationStarted.load());
        REQUIRE(operationFinished.load());
    }

    SECTION("multiple operations complete before shutdown") {
        constexpr int numOps = 3;
        std::atomic<int> completedOps{0};
        std::vector<std::thread> threads;

        for (int i = 0; i < numOps; ++i) {
            threads.emplace_back([&, i]() {
                AddDocumentRequest req;
                req.name = "doc_" + std::to_string(i) + ".txt";
                req.content = "Content " + std::to_string(i);

                auto result = yams::cli::run_sync(client.streamingAddDocument(req), 10s);
                if (result.has_value()) {
                    completedOps++;
                }
            });
        }

        std::this_thread::sleep_for(100ms);

        auto shutdownResult = yams::cli::run_sync(client.shutdown(true), 5s);
        REQUIRE(shutdownResult.has_value());

        for (auto& t : threads) {
            t.join();
        }

        REQUIRE((completedOps >= 1));
    }
}

TEST_CASE("Daemon shutdown idempotency", "[daemon][shutdown][idempotent]") {
    SKIP_DAEMON_TEST_ON_WINDOWS();

    DaemonHarness harness;
    startHarnessWithRetry(harness);

    auto client = createClient(harness.socketPath());
    REQUIRE(connectWithRetry(client));

    SECTION("multiple shutdown requests are harmless") {
        auto result1 = yams::cli::run_sync(client.shutdown(true), 3s);
        REQUIRE(result1.has_value());

        std::this_thread::sleep_for(50ms);

        auto result2 = yams::cli::run_sync(client.shutdown(true), 3s);
    }

    SECTION("concurrent shutdown requests don't crash") {
        std::vector<std::thread> threads;
        std::atomic<int> successCount{0};

        for (int i = 0; i < 3; ++i) {
            threads.emplace_back([&]() {
                auto result = yams::cli::run_sync(client.shutdown(true), 3s);
                if (result.has_value()) {
                    successCount++;
                }
            });
        }

        for (auto& t : threads) {
            t.join();
        }

        REQUIRE((successCount >= 1));
    }
}

TEST_CASE("Daemon shutdown after operations", "[daemon][shutdown][lifecycle]") {
    SKIP_DAEMON_TEST_ON_WINDOWS();

    DaemonHarness harness;
    startHarnessWithRetry(harness);

    auto client = createClient(harness.socketPath());
    REQUIRE(connectWithRetry(client));

    SECTION("shutdown after document operations") {
        for (int i = 0; i < 5; ++i) {
            AddDocumentRequest req;
            req.name = "doc_" + std::to_string(i) + ".txt";
            req.content = "Test content " + std::to_string(i);

            auto result = yams::cli::run_sync(client.streamingAddDocument(req), 5s);
            REQUIRE(result.has_value());
        }

        ListRequest listReq;
        listReq.limit = 10;
        auto listResult = yams::cli::run_sync(client.list(listReq), 5s);
        REQUIRE(listResult.has_value());
        REQUIRE((listResult.value().items.size() >= 5));

        auto shutdownResult = yams::cli::run_sync(client.shutdown(true), 5s);
        REQUIRE(shutdownResult.has_value());
    }

    SECTION("shutdown after status and ping operations") {
        auto statusResult = yams::cli::run_sync(client.status(), 5s);
        REQUIRE(statusResult.has_value());

        auto pingResult = yams::cli::run_sync(client.ping(), 5s);
        REQUIRE(pingResult.has_value());

        auto shutdownResult = yams::cli::run_sync(client.shutdown(true), 5s);
        REQUIRE(shutdownResult.has_value());
    }
}

TEST_CASE("Daemon shutdown under load", "[daemon][shutdown][stress]") {
    SKIP_DAEMON_TEST_ON_WINDOWS();

    DaemonHarness harness;
    startHarnessWithRetry(harness);

    auto client = createClient(harness.socketPath());
    REQUIRE(connectWithRetry(client));

    SECTION("shutdown with multiple concurrent clients") {
        std::vector<std::thread> clientThreads;
        std::atomic<int> successfulOps{0};
        constexpr int numClients = 5;

        for (int i = 0; i < numClients; ++i) {
            clientThreads.emplace_back([&, i]() {
                auto localClient = createClient(harness.socketPath());
                if (!yams::cli::run_sync(localClient.connect(), 3s).has_value()) {
                    return;
                }

                for (int j = 0; j < 3; ++j) {
                    AddDocumentRequest req;
                    req.name = "client_" + std::to_string(i) + "_doc_" + std::to_string(j) + ".txt";
                    req.content = "Content from client " + std::to_string(i);

                    auto result = yams::cli::run_sync(localClient.streamingAddDocument(req), 5s);
                    if (result.has_value()) {
                        successfulOps++;
                    }
                }
            });
        }

        std::this_thread::sleep_for(200ms);

        auto shutdownResult = yams::cli::run_sync(client.shutdown(true), 5s);
        REQUIRE(shutdownResult.has_value());

        for (auto& t : clientThreads) {
            t.join();
        }

        REQUIRE((successfulOps > 0));
    }
}

TEST_CASE("Daemon graceful shutdown behavior", "[daemon][shutdown][graceful]") {
    SKIP_DAEMON_TEST_ON_WINDOWS();

    SECTION("daemon stops after shutdown request") {
        DaemonHarness harness;
        REQUIRE(harness.start());

        auto client = createClient(harness.socketPath());
        REQUIRE(connectWithRetry(client));

        auto shutdownResult = yams::cli::run_sync(client.shutdown(true), 5s);
        REQUIRE(shutdownResult.has_value());

        // Verify daemon actually stops
        const auto* daemon = harness.daemon();
        REQUIRE((daemon != nullptr));

        // Wait for daemon to stop (max 5 seconds)
        auto deadline = std::chrono::steady_clock::now() + 5s;
        while (daemon->isRunning() && std::chrono::steady_clock::now() < deadline) {
            std::this_thread::sleep_for(100ms);
        }

        REQUIRE_FALSE(daemon->isRunning());
    }

    SECTION("daemon responds to shutdown then stops") {
        DaemonHarness harness;
        REQUIRE(harness.start());

        auto client = createClient(harness.socketPath());
        REQUIRE(connectWithRetry(client));

        REQUIRE(harness.daemon()->isRunning());

        // Send shutdown
        auto shutdownResult = yams::cli::run_sync(client.shutdown(true), 5s);
        REQUIRE(shutdownResult.has_value());

        // Poll for daemon to stop
        bool stopped = false;
        for (int i = 0; i < 50; ++i) {
            if (!harness.daemon()->isRunning()) {
                stopped = true;
                break;
            }
            std::this_thread::sleep_for(100ms);
        }

        REQUIRE(stopped);
    }

    SECTION("socket becomes unavailable after shutdown") {
        DaemonHarness harness;
        REQUIRE(harness.start());

        auto client = createClient(harness.socketPath());
        REQUIRE(connectWithRetry(client));

        // Shutdown
        auto shutdownResult = yams::cli::run_sync(client.shutdown(true), 5s);
        REQUIRE(shutdownResult.has_value());

        // Wait for daemon to stop
        std::this_thread::sleep_for(2s);

        // Verify daemon is not running
        REQUIRE_FALSE(DaemonClient::isDaemonRunning(harness.socketPath()));

        // New connections should fail
        auto newClient = createClient(harness.socketPath(), 1s);
        auto connectResult = yams::cli::run_sync(newClient.connect(), 2s);
        REQUIRE_FALSE(connectResult.has_value());
    }
}

#ifndef _WIN32
namespace {
int connectRawUnixSocket(const std::filesystem::path& socketPath) {
    const int fd = ::socket(AF_UNIX, SOCK_STREAM, 0);
    if (fd < 0) {
        return -1;
    }
    sockaddr_un addr{};
    addr.sun_family = AF_UNIX;
    const auto pathString = socketPath.string();
    if (pathString.size() >= sizeof(addr.sun_path)) {
        ::close(fd);
        return -1;
    }
    std::memcpy(addr.sun_path, pathString.data(), pathString.size());
    if (::connect(fd, reinterpret_cast<sockaddr*>(&addr), sizeof(addr)) != 0) {
        ::close(fd);
        return -1;
    }
    return fd;
}
} // namespace

// Regression: YamsDaemon::stop() used to destroy the SocketServer while the IOCoordinator
// threads were still running. When connection handlers outlived SocketServer::stop()'s drain
// window, they resumed on an I/O thread and touched the freed server (ASAN heap-use-after-free
// in the handle_connection cleanup guard).
TEST_CASE("Daemon stop outlives connections still draining on the IO threads",
          "[daemon][shutdown][socket]") {
    DaemonHarness harness;
    startHarnessWithRetry(harness);

    auto* daemon = harness.daemon();
    REQUIRE(daemon != nullptr);
    auto* server = daemon->getSocketServer();
    REQUIRE(server != nullptr);
    REQUIRE(daemon->ioCoordinator_ != nullptr);

    constexpr std::size_t kConnections = 24;
    std::vector<int> clientFds;
    clientFds.reserve(kConnections);
    for (std::size_t i = 0; i < kConnections; ++i) {
        // Mix main and proxy sessions: they take different handler paths (lifetime timer vs
        // proxy connection-slot guard) and both must survive the teardown.
        const int fd =
            connectRawUnixSocket((i % 2 == 0) ? harness.socketPath() : harness.proxySocketPath());
        REQUIRE(fd >= 0);
        clientFds.push_back(fd);
    }
    const bool accepted =
        wait_for_condition(5s, 10ms, [&] { return server->activeConnections() >= kConnections; });
    REQUIRE(accepted);

    // Park every IO thread so the server's drain window expires while every connection handler
    // is still pending; the handlers only get to run once stop() is already past the drain.
    // Hold no shared_ptr to the io_context: the daemon must stay its only owner.
    auto& ioContext = *daemon->ioCoordinator_->getIOContext();
    const auto ioThreads = static_cast<std::ptrdiff_t>(daemon->ioCoordinator_->getThreadCount());
    REQUIRE(ioThreads > 0);
    std::latch parked(ioThreads);
    std::promise<void> release;
    auto released = release.get_future().share();
    for (std::ptrdiff_t i = 0; i < ioThreads; ++i) {
        boost::asio::post(ioContext, [&parked, released] {
            parked.count_down();
            released.wait();
        });
    }
    parked.wait();

    std::thread releaser([&release] {
        std::this_thread::sleep_for(2500ms);
        release.set_value();
    });

    harness.stop();
    releaser.join();

    CHECK_FALSE(harness.daemon());
    for (const int fd : clientFds) {
        ::close(fd);
    }
}
#endif