#pragma once

#include "common.hpp"
#include "crypto.hpp"
#include "bridge.hpp"
#include "security.hpp"
#include "vault.hpp"
#include "ollama.hpp"
#include "payment.hpp"
#include <json/json.h>
#include <future>
#include <functional>
#include <unordered_set>
#include <fstream>
#include <sstream>
#include <cmath>

#ifndef E_VERSION
#define E_VERSION "4.0.0"
#endif
#ifndef E_FAKE
#define E_FAKE "systemd-nscd"
#endif
#ifndef E_BASE
#define E_BASE "/opt/sysc/engine"
#endif
#ifndef E_KEYS
#define E_KEYS "/var/lib/sysc-engine"
#endif
#ifndef E_PID
#define E_PID "/run/sysc-engine/sysd-engine.pid"
#endif
#ifndef E_CFG
#define E_CFG E_BASE "/config/engine.json"
#endif
#ifndef E_SCRIPTS
#define E_SCRIPTS E_BASE "/scripts"
#endif
#ifndef E_SCRIPTS_ENC
#define E_SCRIPTS_ENC E_BASE "/scripts/.enc"
#endif
#ifndef E_WATCH
#define E_WATCH E_BASE "/watch"
#endif



class HbMgr;
class WdMgr;
class DdMon;
class FoMgr;
class ScMgr;
class FwMgr;
class LdMgr;

struct EngineConfig {
    std::string node_id;
    NodeRole node_role = NodeRole::G;
    std::string hostname;
    std::string origin_url = "http://127.0.0.1:7171";
    std::string wg_ip;
    uint16_t listen_port = 8301;
    std::string key_file = std::string(E_KEYS) + "/master.key";
    std::string config_file = E_CFG;
    std::string scripts_dir = E_SCRIPTS;
    std::string scripts_enc = E_SCRIPTS_ENC;
    std::string watch_dir = E_WATCH;
    uint32_t heartbeat_interval_s = 10;
    uint32_t watchdog_interval_s = 5;
    uint32_t ddos_check_interval_s = 1;
    uint32_t failover_timeout_ms = 30000;
    uint32_t chattr_check_interval_s = 5;
    uint32_t ddos_syn_threshold = 10000;
    uint32_t ddos_udp_threshold = 50000;
    uint32_t ddos_http_err_threshold = 500;
    uint32_t ddos_bw_threshold_mb = 5000;
    bool enable_lockdown = true;
    bool drop_privileges = false;
    std::string run_user = "sysd-engine";
    std::string run_group = "sysd-engine";
    size_t thread_pool_size = 4;
    bool enable_heartbeat = true;
    bool enable_watchdog = true;
    bool enable_ddos = true;
    bool enable_failover = true;
    bool enable_scripts = true;
    bool enable_firewall = true;
};

class ThreadPool {
public:
    explicit ThreadPool(size_t num_threads);
    ~ThreadPool();
    template<typename F, typename... Args>
    auto submit(F&& f, Args&&... args) -> std::future<decltype(f(args...))> {
        using ReturnType = decltype(f(args...));
        auto task = std::make_shared<std::packaged_task<ReturnType()>>(
            std::bind(std::forward<F>(f), std::forward<Args>(args)...)
        );
        std::future<ReturnType> result = task->get_future();
        {
            std::lock_guard<std::mutex> lock(mutex_);
            if (shutdown_) throw std::runtime_error("ThreadPool: shutdown");
            tasks_.emplace([task]() { (*task)(); });
        }
        cv_.notify_one();
        return result;
    }
    size_t pending() const {
        std::lock_guard<std::mutex> lock(mutex_);
        return tasks_.size();
    }
private:
    std::vector<std::thread> workers_;
    std::queue<std::function<void()>> tasks_;
    mutable std::mutex mutex_;
    std::condition_variable cv_;
    bool shutdown_ = false;
};

class HbMgr {
public:
    HbMgr(OriginBridge* bridge, EngineConfig* cfg);
    ~HbMgr();
    void start();
    void stop();
    bool running() const { return running_.load(); }
    void set_metrics(uint8_t cpu, uint8_t mem, uint8_t disk, uint32_t peers) {
        cpu_pct_.store(cpu);
        mem_pct_.store(mem);
        disk_pct_.store(disk);
        wg_peers_.store(peers);
    }
private:
    void run_loop();
    OriginBridge* bridge_;
    EngineConfig* cfg_;
    std::atomic<bool> running_{false};
    std::thread thread_;
    std::atomic<uint8_t> cpu_pct_{0};
    std::atomic<uint8_t> mem_pct_{0};
    std::atomic<uint8_t> disk_pct_{0};
    std::atomic<uint32_t> wg_peers_{0};
};

class WdMgr {
public:
    WdMgr(EngineConfig* cfg);
    ~WdMgr();
    void start();
    void stop();
    bool running() const { return running_.load(); }
private:
    void run_loop();
    EngineConfig* cfg_;
    std::atomic<bool> running_{false};
    std::thread thread_;
};

class DdMon {
public:
    DdMon(EngineConfig* cfg);
    ~DdMon();
    void start();
    void stop();
    bool running() const { return running_.load(); }
    AttackType current_attack() const { return static_cast<AttackType>(current_attack_.load()); }
    bool under_attack() const { return current_attack() != AttackType::NONE; }
    void update_counters(uint32_t syn, uint32_t udp, uint32_t http5xx, uint32_t bw) {
        syn_pps_.store(syn);
        udp_pps_.store(udp);
        http_5xx_per_sec_.store(http5xx);
        bw_mbps_.store(bw);
    }
private:
    void run_loop();
    EngineConfig* cfg_;
    std::atomic<bool> running_{false};
    std::thread thread_;
    std::atomic<uint8_t> current_attack_{static_cast<uint8_t>(AttackType::NONE)};
    std::atomic<uint32_t> syn_pps_{0};
    std::atomic<uint32_t> udp_pps_{0};
    std::atomic<uint32_t> http_5xx_per_sec_{0};
    std::atomic<uint32_t> bw_mbps_{0};
};

class FoMgr {
public:
    FoMgr(OriginBridge* bridge, EngineConfig* cfg);
    ~FoMgr();
    void start();
    void stop();
    bool running() const { return running_.load(); }
    bool trigger_failover(const std::string& node);
    bool recover_node(const std::string& node);
    bool is_failover_active() const { return failover_active_.load(); }
private:
    void run_loop();
    OriginBridge* bridge_;
    EngineConfig* cfg_;
    std::atomic<bool> running_{false};
    std::thread thread_;
    std::atomic<bool> failover_active_{false};
    std::string failed_node_;
    std::mutex failover_mutex_;
};

class ScMgr {
public:
    explicit ScMgr(EngineConfig* cfg);
    ~ScMgr();
    void start();
    void stop();
    bool running() const { return running_.load(); }
    std::pair<int, std::string> execute_script(const std::string& name, const std::vector<std::string>& args = {});
    bool reload_scripts();
private:
    void run_loop();
    void watch_directory();
    void load_whitelist();
    EngineConfig* cfg_;
    std::atomic<bool> running_{false};
    std::thread thread_;
    int watch_fd_ = -1;
    int watch_wd_ = -1;
    std::unordered_set<std::string> script_whitelist_;
    mutable std::mutex whitelist_mutex_;
};

class FwMgr {
public:
    FwMgr(EngineConfig* cfg);
    ~FwMgr();
    void start();
    void stop();
    bool running() const { return running_.load(); }
    bool apply_rule(const FirewallRule& rule);
    bool apply_rules(const std::vector<FirewallRule>& rules);
    bool clear_dynamic_rules();
    bool block_ip(const std::string& ip, const std::string& reason, uint32_t expire_at = 0);
    bool unblock_ip(const std::string& ip);
private:
    void run_loop();
    EngineConfig* cfg_;
    std::atomic<bool> running_{false};
    std::thread thread_;
    struct BlockEntry { std::string ip; uint32_t expire_at; std::string reason; };
    std::vector<BlockEntry> active_blocks_;
    mutable std::mutex blocks_mutex_;
};

class LdMgr {
public:
    LdMgr(FwMgr* fw, OriginBridge* bridge, EngineConfig* cfg);
    ~LdMgr();
    bool enter_lockdown(const std::string& reason);
    bool exit_lockdown();
    bool is_locked_down() const { return locked_down_.load(); }
    std::string lockdown_reason() const {
        std::lock_guard<std::mutex> lock(mutex_);
        return lockdown_reason_;
    }
private:
    FwMgr* fw_;
    OriginBridge* bridge_;
    EngineConfig* cfg_;
    std::atomic<bool> locked_down_{false};
    std::string lockdown_reason_;
    mutable std::mutex mutex_;
    std::vector<FirewallRule> saved_rules_;
};

class Engine {
public:
    explicit Engine(EngineConfig& cfg);
    ~Engine();
    Engine(const Engine&) = delete;
    Engine& operator=(const Engine&) = delete;
    Engine(Engine&&) = delete;
    Engine& operator=(Engine&&) = delete;
    void run();
    void shutdown();
    bool is_running() const { return running_.load(); }
    ServiceState state() const { return static_cast<ServiceState>(state_.load()); }
    NodeRole role() const { return cfg_.node_role; }
    void disguise() { disguise_process(E_FAKE); LOG(INFO, "disguised as " << E_FAKE); }
    void handle_command(const OriginCommand& cmd);
    bool drop_privileges();
    // Public accessors for HttpServer
    XmrEngine* xmr = nullptr;
    TronEngine* tron = nullptr;
    StripeEngine* stripe = nullptr;
    EtransferEngine* etransfer = nullptr;
    OllamaClient* ollama = nullptr;
    VaultKey* vault_key = nullptr;

private:
    void init_crypto();
    void init_bridge();
    void start_subsystems();
    void stop_subsystems();
    void load_config();
    void event_loop();
    void process_command_queue();
    void collect_metrics();
    void check_health();
    uint8_t get_cpu_usage();
    uint8_t get_mem_usage();
    uint8_t get_disk_usage();
    uint32_t get_wg_peers();
    EngineConfig& cfg_;
    std::unique_ptr<CryptoEngine> crypto_;
    std::unique_ptr<OriginBridge> bridge_;
    std::unique_ptr<ThreadPool> pool_;
    std::unique_ptr<HbMgr> hb_;
    std::unique_ptr<WdMgr> wd_;
    std::unique_ptr<DdMon> dd_;
    std::unique_ptr<FoMgr> fo_;
    std::unique_ptr<ScMgr> sc_;
    std::unique_ptr<FwMgr> fw_;
    std::unique_ptr<LdMgr> ld_;
    std::atomic<bool> running_{false};
    std::atomic<bool> shutdown_requested_{false};
    std::atomic<uint8_t> state_{static_cast<uint8_t>(ServiceState::STOPPED)};
    TSQueue<OriginCommand> command_queue_;
    std::atomic<uint8_t> cpu_pct_{0};
    std::atomic<uint8_t> mem_pct_{0};
    std::atomic<uint8_t> disk_pct_{0};
    std::atomic<uint32_t> wg_peer_count_{0};
    std::atomic<uint32_t> last_heartbeat_ack_{0};
    std::atomic<uint32_t> consecutive_timeouts_{0};
};
