yams/include/yams/daemon/daemon.h

#pragma once

#include <yams/core/types.h>
#include <yams/daemon/components/DaemonLifecycleFsm.h>
#include <yams/daemon/components/StateComponent.h>
#include <yams/daemon/components/TuningConfig.h>
#include <yams/daemon/resource/onnx_model_pool.h> // For DaemonConfig

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <cstdio>
#include <filesystem>
#include <functional>
#include <future>
#include <map>
#include <memory>
#include <mutex>
#include <thread>
#include <yams/compat/thread_stop_compat.h>
//

namespace yams::daemon {

// Forward declarations for components
class LifecycleComponent;
class DaemonLifecycleAdapter;
class ServiceManager;
class RequestDispatcher;
class DaemonMetrics;
class SocketServer;
class IOCoordinator;
class TuningManager; // centralized tuner (moved from ServiceManager)
// Forward decls for GTEST-only accessors are below guarded by YAMS_TESTING

struct DaemonConfig {
    std::filesystem::path dataDir;
    std::filesystem::path socketPath;
    std::filesystem::path pidFile;
    std::filesystem::path logFile;
    size_t maxMemoryGb = 4;
    std::chrono::milliseconds requestTimeout{5000};
    bool healthMonitoring = true;
    std::string logLevel = "info";
    size_t maxLogFiles = 5;
    size_t maxLogSizeMb = 10;
    bool enableModelProvider = false;
    bool modelProviderRequired = false;
    ModelPoolConfig modelPoolConfig;
    std::filesystem::path pluginDir;
    std::string pluginNamePolicy{"relaxed"};
    bool pluginDirStrict{false};
    std::vector<std::filesystem::path> trustedPluginPaths;
    // Default to disabled to avoid unintended plugin loading in plain daemon
    // invocations and tests that don't explicitly opt in. Tests that rely on
    // plugins (e.g., ONNX embeddings) set this to true in their config.
    bool autoLoadPlugins = false;
    bool enableAutoRepair = true;
    bool useMockModelProvider = false;
    size_t autoRepairBatchSize = 32;
    bool autoRebuildOnDimMismatch = false;
    size_t maxPendingRepairs = 1000;
    // Embedded one-shot clients (sandboxed CLIs, Codex) should avoid startup maintenance that can
    // outlive the command and make teardown expensive. Request handlers may still perform explicit
    // writes, but opportunistic reconcile/rebuild work is suppressed.
    bool embeddedOneShot = false;

    // Typed runtime instrumentation profile. The "memory" profile suppresses
    // opportunistic startup maintenance that creates large transient allocator
    // churn, which is especially expensive under macOS MallocStackLogging.
    struct InstrumentationPolicy {
        // "auto" enables the memory profile when MallocStackLogging is present;
        // "memory" forces it; "normal"/"off" disables it.
        std::string profile{"auto"};
        bool memoryProfileActive{false};
        bool suppressAutoRepair{false};
        bool suppressSimeonLexicalBuild{false};
        bool suppressVectorIndexBuild{false};
        std::uint64_t mslStackLogWarnBytes{2ULL * 1024ULL * 1024ULL * 1024ULL};
    } instrumentation;

    // Per-plugin configuration: plugin name -> JSON config string
    // Parsed from [plugins.<name>] sections in config.toml
    std::map<std::string, std::string> pluginConfigs;

    // Streaming keepalive configuration
    std::chrono::milliseconds heartbeatInterval{500}; // default 500ms
    std::chrono::milliseconds heartbeatJitter{50};    // default +/-50ms applied per tick

    // Adaptive tuning configuration (no envs)
    TuningConfig tuning{};
    // Path to loaded config file for reloads
    std::filesystem::path configFilePath;

    struct GraphPrunePolicy {
        bool enabled{false};
        std::size_t keepLatestPerCanonical{3};
        std::chrono::minutes interval{60};
        std::chrono::minutes initialDelay{10};
    } graphPrune;

    // Forward decls for GTEST-only accessors are below guarded by YAMS_TESTING
    struct DownloadPolicy {
        bool enable{false};                               // feature gate
        std::vector<std::string> allowedHosts{};          // exact or wildcard patterns
        std::vector<std::string> allowedSchemes{"https"}; // schemes to allow (default https)
        bool requireChecksum{true};                       // require algo:hex on requests
        bool storeOnly{true};                     // only write into CAS (no arbitrary paths)
        std::chrono::milliseconds timeout{60000}; // per-request timeout
        std::uint64_t maxFileBytes{0};            // 0 = unlimited
        size_t rateLimitRps{0};                   // 0 = unlimited
        std::string sandbox{"subprocess"};        // future: platform-specific isolation
    } downloadPolicy;
};

class YamsDaemon {
public:
    explicit YamsDaemon(const DaemonConfig& config = {});
    ~YamsDaemon();

    // Lifecycle management
    Result<void> start();
    Result<void> stop();
    bool waitForStopCompletion(std::chrono::milliseconds timeout);
    /// Run the daemon main loop on the calling thread. Call after start().
    /// Returns when stopRequested_ becomes true or requestStop() is called.
    void runLoop();
    /// Test helper: trigger deferred async initialization and the init waiter
    /// without entering the long-lived runLoop wait loop.
    void testingStartAsyncInitWithoutRunLoop();
    void requestStop() {
        stopRequested_.store(true, std::memory_order_release);
        stop_cv_.notify_all();
    }
    bool isRunning() const { return running_.load(); }
    bool isStopRequested() const { return stopRequested_.load(); }

    const StateComponent& getState() const { return state_; }
    // Lifecycle FSM accessor
    const DaemonLifecycleFsm& getLifecycle() const { return lifecycleFsm_; }
    // Mark or clear subsystem degradation (sticky until cleared by operator action)
    void setSubsystemDegraded(const std::string& name, bool degraded,
                              const std::string& reason = std::string()) {
        lifecycleFsm_.setSubsystemDegraded(name, degraded, reason);
    }

    // Socket server accessor (for testing)
    SocketServer* getSocketServer() const { return socketServer_.get(); }

    // Document event notifications - forwarded to RepairService via ServiceManager
    void onDocumentAdded(const std::string& hash, const std::string& path);
    void onDocumentRemoved(const std::string& hash);

    // Path resolution helpers
    enum class PathType { Socket, PidFile, LogFile };
    static std::filesystem::path resolveSystemPath(PathType type);
    static bool canWriteToDirectory(const std::filesystem::path& dir);
    static std::filesystem::path getXDGRuntimeDir();
    DaemonConfig config_;
    StateComponent state_;
    std::unique_ptr<LifecycleComponent> lifecycleManager_;
    std::shared_ptr<ServiceManager> serviceManager_;
    std::unique_ptr<DaemonLifecycleAdapter> lifecycleAdapter_;
    std::unique_ptr<RequestDispatcher> requestDispatcher_;
    // Accessed from multiple threads (runLoop + request handling/shutdown). Guard with mutex.
    mutable std::mutex metricsMutex_;
    // Protects mutable config fields updated after startup (currently tuning reload state).
    mutable std::mutex configMutex_;
    std::shared_ptr<DaemonMetrics> metrics_;
    // Integrated socket server (replaces external yams-socket-server)
    std::unique_ptr<SocketServer> socketServer_;
    // Dedicated I/O thread pool for IPC sockets
    std::unique_ptr<IOCoordinator> ioCoordinator_;
    std::unique_ptr<TuningManager> tuningManager_;

    // Lifecycle FSM (authoritative lifecycle state)
    DaemonLifecycleFsm lifecycleFsm_{};

    // Threading and state
    std::atomic<bool> running_{false};
    std::atomic<bool> stopRequested_{false};
    std::atomic<bool> reloadRequested_{false};
    std::atomic<bool> runLoopStarted_{false};
    // Main loop runs on the calling thread via runLoop().
    std::mutex stop_mutex_;
    std::condition_variable stop_cv_;
    bool stopCompleted_{true};

    std::promise<void> asyncInitStartedPromise_;
    std::shared_future<void> asyncInitStartedFuture_;
    std::atomic<bool> asyncInitBarrierSet_{false};

    std::thread initWaiterThread_;
    std::atomic<bool> initHandled_{false};

    std::mutex shutdownThreadMutex_;
    std::thread shutdownThread_;
    std::atomic<bool> shutdownThreadActive_{false};

    bool modelPreloadSkipped_{false};
    std::chrono::steady_clock::time_point repairBusySince_{};
    std::chrono::steady_clock::time_point repairReadySince_{};

    // Signal check hook for integration with main loop (avoids separate thread)
    std::function<bool()> signalCheckHook_;

    // External shutdown flag reference for CV predicate wake-up (yams-qe6r fix)
    // When set, the CV predicate in runLoop() also checks this flag, enabling
    // prompt response to SIGTERM even when blocked waiting on the CV.
    // Atomic pointer to avoid data race between setter and runLoop reader.
    std::atomic<std::atomic<bool>*> externalShutdownFlag_{nullptr};

    // On-demand plugin autoload bridge
    Result<size_t> autoloadPluginsNow();
    ServiceManager* getServiceManager() const { return serviceManager_.get(); }
    void requestReload() { reloadRequested_.store(true, std::memory_order_relaxed); }
    void reloadTuningConfig();

    // Set a hook that will be called each iteration of runLoop() to check for signals
    // Returns true if shutdown was requested
    void setSignalCheckHook(std::function<bool()> hook) { signalCheckHook_ = std::move(hook); }

    // Set an external shutdown flag that the CV predicate will also check (yams-qe6r)
    // This enables the runLoop to respond promptly to SIGTERM even when blocked on CV
    void setExternalShutdownFlag(std::atomic<bool>* flag) {
        externalShutdownFlag_.store(flag, std::memory_order_release);
    }

    void spawnShutdownThread(std::function<void()> shutdownFn);

    void reapCompletedShutdownThread();
};

} // namespace yams::daemon