#pragma once
#include "diagnostics.hpp"
#include "inode.hpp"
#include "port.hpp"

#ifdef KPN_WEB_DEBUG
#include "web_debug.hpp"
#include <memory>
#endif

#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <set>
#include <sstream>
#include <stdexcept>
#include <string>
#include <string_view>
#include <thread>
#include <vector>

namespace kpn {

// ── Exceptions ────────────────────────────────────────────────────────────────

class NetworkCycleError : public std::runtime_error {
public:
    NetworkCycleError() : std::runtime_error("network graph contains a directed cycle") {}
};

class NetworkBuildError : public std::runtime_error {
public:
    explicit NetworkBuildError(std::string msg)
        : std::runtime_error(std::move(msg)) {}
};

// ── Network ───────────────────────────────────────────────────────────────────

class Network : public INode {
public:
    using ErrorHandler =
        std::function<void(std::string_view node_name, std::exception_ptr)>;
    using DiagnosticsHandler =
        std::function<void(const std::vector<NodeSnapshot>&,
                           const std::vector<ChannelSnapshot>&)>;

    // ── Builder API ───────────────────────────────────────────────────────────

    template<typename NodeT>
    Network& add(std::string name, NodeT& node) {
        if (nodes_.count(name))
            throw NetworkBuildError("duplicate node name: " + name);
        node.set_name(name);
        nodes_.emplace(name, &node);
        adj_[name];
        return *this;
    }

    template<typename SrcNode, std::size_t SrcIdx,
             typename DstNode, std::size_t DstIdx>
    Network& connect(const std::string& src_name,
                     OutputPort<SrcNode, SrcIdx>,
                     const std::string& dst_name,
                     InputPort<DstNode,  DstIdx>) {
        using out_t    = std::tuple_element_t<SrcIdx, typename SrcNode::return_tuple>;
        using dst_in_t = std::tuple_element_t<DstIdx, typename DstNode::args_tuple>;
        static_assert(std::is_same_v<out_t, dst_in_t>,
            "connect: output type does not match input type");

        auto* src = dynamic_cast<SrcNode*>(nodes_.at(src_name));
        auto* dst = dynamic_cast<DstNode*>(nodes_.at(dst_name));
        if (!src) throw NetworkBuildError("node '" + src_name + "' type mismatch");
        if (!dst) throw NetworkBuildError("node '" + dst_name + "' type mismatch");

        auto port_key = std::make_pair(src_name, SrcIdx);
        if (connected_outputs_.count(port_key))
            throw NetworkBuildError(
                "connect: output port '" + src_name + "':" + std::to_string(SrcIdx)
                + " is already connected — use make_fanout<T, N> for fan-out");
        connected_outputs_.insert(port_key);

        auto& in_ch = dst->template input_channel<DstIdx>();
        src->template set_output_channel<SrcIdx>(&in_ch);

        // Register channel probe for diagnostics
        std::string ch_name = src_name + ":" + std::to_string(SrcIdx)
                            + " → " + dst_name + ":" + std::to_string(DstIdx);
        channel_probes_.push_back(
            std::make_unique<ChannelProbe<out_t>>(in_ch, ch_name));

        adj_[src_name].push_back(dst_name);
        return *this;
    }

    template<typename NodeT, std::size_t Idx>
    Network& expose_input(std::string boundary_name, InputPort<NodeT, Idx>) {
        exposed_inputs_[boundary_name] = boundary_name;
        return *this;
    }

    template<typename NodeT, std::size_t Idx>
    Network& expose_output(std::string boundary_name, OutputPort<NodeT, Idx>) {
        exposed_outputs_[boundary_name] = boundary_name;
        return *this;
    }

    Network& build() {
        topo_.clear();
        std::map<std::string, int> color;
        for (auto& [name, _] : nodes_)
            if (color[name] == 0)
                dfs(name, color);
        return *this;
    }

    // ── INode ─────────────────────────────────────────────────────────────────

    void start() override {
        start_time_ = clock_t::now();
        for (auto& name : topo_)
            nodes_.at(name)->start();
        start_watchdog();
#ifdef KPN_WEB_DEBUG
        web_server_ = std::make_unique<web_debug::WebDebugServer>(
            web_debug_port_,
            [this]() {
                auto s = collect_snapshots();
                return web_debug::to_json(s.nodes, s.channels, {}, s.elapsed_s, s.pools);
            });
        web_server_->start();
        std::cerr << "[kpn] web debug UI: http://localhost:" << web_debug_port_ << "\n";
#endif
    }

    void stop() override { halt(); }

    // halt(): immediate stop — broadcasts disable to all channels and joins threads.
    void halt() override {
#ifdef KPN_WEB_DEBUG
        if (web_server_) web_server_->stop();
#endif
        stop_watchdog();
        for (auto it = topo_.rbegin(); it != topo_.rend(); ++it)
            nodes_.at(*it)->stop();
    }

    // shutdown(): graceful drain in topological order.
    // Stops source nodes first, polls until their output channels drain to zero,
    // then stops the next layer, and so on.
    void shutdown() override {
#ifdef KPN_WEB_DEBUG
        if (web_server_) web_server_->stop();
#endif
        stop_watchdog();

        // Identify which nodes have no incoming edges (sources).
        std::map<std::string, std::size_t> in_degree;
        for (auto& [name, _] : nodes_) in_degree[name] = 0;
        for (auto& [src, dsts] : adj_)
            for (auto& dst : dsts) in_degree[dst]++;

        // Walk topo order: stop each source layer, wait for its output channels
        // to drain, then proceed to the next layer.
        std::set<std::string> stopped;
        for (auto& name : topo_) {
            if (in_degree[name] == 0 || all_predecessors_stopped(name, stopped)) {
                nodes_.at(name)->stop();
                stopped.insert(name);
                // Wait for output channels of this node to drain.
                drain_output_channels(name);
            }
        }
        // Stop any remaining nodes (sinks / nodes not yet stopped).
        for (auto it = topo_.rbegin(); it != topo_.rend(); ++it)
            if (!stopped.count(*it)) nodes_.at(*it)->stop();
    }

    bool running() const override { return watchdog_.joinable(); }
    void set_name(std::string) override {}

    const NodeStats& stats() const override {
        static NodeStats dummy;
        return dummy;
    }

    NodeSnapshot node_snapshot(const std::string& name, double) const override {
        return {name, 0, 0, 0, 0, 0, 0, 0};
    }

    // ── Configuration ─────────────────────────────────────────────────────────

    void set_watchdog_interval(std::chrono::milliseconds interval) {
        watchdog_interval_ = interval;
    }

    void set_error_handler(ErrorHandler h)       { error_handler_       = std::move(h); }
    void set_diagnostics_handler(DiagnosticsHandler h) { diag_handler_  = std::move(h); }

    void register_pool(const std::string& name, IPoolProbe* probe) {
        pool_probes_.emplace_back(name, probe);
    }

#ifdef KPN_WEB_DEBUG
    void set_web_debug_port(uint16_t port) { web_debug_port_ = port; }
#endif

    // Print a diagnostics report to a stream (default: stderr).
    // Can be called at any time; thread-safe (reads atomics with relaxed ordering).
    void print_diagnostics(std::ostream& os = std::cerr) const {
        auto s = collect_snapshots();
        os << format_report(s.nodes, s.channels, s.pools, s.elapsed_s);
    }

private:
    // ── Diagnostics collection ────────────────────────────────────────────────

    struct Snapshots {
        std::vector<NodeSnapshot>    nodes;
        std::vector<ChannelSnapshot> channels;
        std::vector<PoolSnapshot>    pools;
        double                       elapsed_s;
    };

    Snapshots collect_snapshots() const {
        double elapsed_s = std::chrono::duration<double>(
            clock_t::now() - start_time_).count();

        std::vector<NodeSnapshot> nodes;
        for (auto& name : topo_)
            nodes.push_back(nodes_.at(name)->node_snapshot(name, elapsed_s));

        std::vector<ChannelSnapshot> channels;
        for (auto& probe : channel_probes_)
            channels.push_back(probe->snapshot());

        std::vector<PoolSnapshot> pools;
        for (auto& [name, probe] : pool_probes_)
            pools.push_back(probe->snapshot(name));

        return {std::move(nodes), std::move(channels), std::move(pools), elapsed_s};
    }

    static std::string format_report(const std::vector<NodeSnapshot>& nodes,
                                     const std::vector<ChannelSnapshot>& channels,
                                     const std::vector<PoolSnapshot>& pools = {},
                                     double elapsed_s = 0.0) {
        std::ostringstream os;
        os << std::fixed << std::setprecision(1);

        os << "\n┌─ KPN++ Diagnostics ────────────────────────────────────────────────────────\n";

        // Node table
        os << "│ Nodes:\n";
        os << "│  " << std::left
           << std::setw(16) << "name"
           << std::setw(10) << "frames"
           << std::setw(12) << "exec ms"
           << std::setw(12) << "max ms"
           << std::setw(14) << "blocked ms"
           << std::setw(10) << "fps"
           << std::setw(12) << "cpu ms"
           << std::setw(10) << "util%"
           << "\n│  " << std::string(92, '-') << "\n";
        for (auto& n : nodes) {
            os << "│  " << std::left
               << std::setw(16) << n.name
               << std::setw(10) << n.frames_processed
               << std::setw(12) << n.ema_exec_ms
               << std::setw(12) << n.max_exec_ms
               << std::setw(14) << n.total_blocked_ms
               << std::setw(10) << n.throughput_fps
               << std::setw(12) << n.total_cpu_ms
               << std::setw(10) << n.cpu_util_pct
               << "\n";
        }

        // Channel table
        os << "│\n│ Channels:\n";
        os << "│  " << std::left
           << std::setw(40) << "edge"
           << std::setw(8)  << "fill%"
           << std::setw(8)  << "peak%"
           << std::setw(8)  << "pushes"
           << std::setw(8)  << "drops"
           << std::setw(8)  << "oflow"
           << std::setw(12) << "MB/s"
           << std::setw(10) << "item B"
           << "\n│  " << std::string(102, '-') << "\n";
        for (auto& c : channels) {
            std::string flag = c.fill_pct() >= 80.0 ? " <<<" :
                               c.peak_pct() >= 80.0 ? " (peak)" : "";
            os << "│  " << std::left
               << std::setw(40) << c.name
               << std::setw(8)  << c.fill_pct()
               << std::setw(8)  << c.peak_pct()
               << std::setw(8)  << c.pushes
               << std::setw(8)  << c.drops
               << std::setw(8)  << c.overflows
               << std::setw(12) << c.bandwidth_mbs(elapsed_s)
               << std::setw(10) << c.item_bytes
               << flag << "\n";
        }

        // Pool table
        if (!pools.empty()) {
            os << "│\n│ Thread Pools:\n";
            os << "│  " << std::left
               << std::setw(16) << "name"
               << std::setw(10) << "threads"
               << std::setw(12) << "queued"
               << std::setw(12) << "active"
               << std::setw(14) << "in/s"
               << std::setw(14) << "out/s"
               << "\n│  " << std::string(78, '-') << "\n";
            for (auto& p : pools) {
                double in_rate  = elapsed_s > 0.0 ? p.tasks_submitted / elapsed_s : 0.0;
                double out_rate = elapsed_s > 0.0 ? p.tasks_completed / elapsed_s : 0.0;
                os << "│  " << std::left
                   << std::setw(16) << p.name
                   << std::setw(10) << p.thread_count
                   << std::setw(12) << p.queue_depth
                   << std::setw(12) << p.active_count
                   << std::setw(14) << in_rate
                   << std::setw(14) << out_rate
                   << "\n";
            }
        }

        // Bottleneck hint: node with highest ema_exec_ms
        if (!nodes.empty()) {
            auto it = std::max_element(nodes.begin(), nodes.end(),
                [](const NodeSnapshot& a, const NodeSnapshot& b) {
                    return a.ema_exec_ms < b.ema_exec_ms;
                });
            os << "│\n│ Bottleneck hint: '" << it->name
               << "' (avg exec " << it->ema_exec_ms << " ms)\n";
        }
        os << "└────────────────────────────────────────────────────────────────\n";
        return os.str();
    }

    // ── Shutdown helpers ──────────────────────────────────────────────────────

    bool all_predecessors_stopped(const std::string& name,
                                  const std::set<std::string>& stopped) const {
        for (auto& [src, dsts] : adj_)
            for (auto& dst : dsts)
                if (dst == name && !stopped.count(src)) return false;
        return true;
    }

    void drain_output_channels(const std::string& /*name*/) const {
        // Poll all channel probes until none report non-zero fill.
        // A short sleep prevents busy-spin; 1 ms is fine for drain purposes.
        bool any_full = true;
        while (any_full) {
            any_full = false;
            for (auto& probe : channel_probes_) {
                auto snap = probe->snapshot();
                if (snap.current_fill > 0) { any_full = true; break; }
            }
            if (any_full)
                std::this_thread::sleep_for(std::chrono::milliseconds(1));
        }
    }

    // ── Cycle detection / topological sort ───────────────────────────────────

    void dfs(const std::string& name, std::map<std::string, int>& color) {
        color[name] = 1;
        for (auto& nbr : adj_[name]) {
            if (color[nbr] == 1) throw NetworkCycleError{};
            if (color[nbr] == 0) dfs(nbr, color);
        }
        color[name] = 2;
        topo_.insert(topo_.begin(), name);
    }

    // ── Watchdog ──────────────────────────────────────────────────────────────

    void start_watchdog() {
        watchdog_ = std::jthread([this](std::stop_token tok) {
            while (!tok.stop_requested()) {
                std::this_thread::sleep_for(watchdog_interval_);
                if (tok.stop_requested()) break;

                auto s = collect_snapshots();
                check_hung_nodes();

                if (diag_handler_) {
                    diag_handler_(s.nodes, s.channels);
                } else {
                    std::cerr << format_report(s.nodes, s.channels, s.pools, s.elapsed_s);
                }
            }
        });
    }

    void check_hung_nodes() const {
        auto now_us = std::chrono::duration_cast<std::chrono::microseconds>(
            clock_t::now().time_since_epoch()).count();
        for (auto& [name, node] : nodes_) {
            int64_t start = node->stats().exec_start_us.load(std::memory_order_relaxed);
            if (start == 0) continue;
            int64_t elapsed_ms = (now_us - start) / 1000;
            // Warn if a node has been executing for > 5 s with no max_exec_time set,
            // or if it exceeds its configured max. Threshold: 5000 ms default.
            if (elapsed_ms > 5000) {
                std::cerr << "[kpn] WARNING: node '" << name
                          << "' has been executing for " << elapsed_ms << " ms\n";
            }
        }
    }

    void stop_watchdog() {
        if (watchdog_.joinable())
            watchdog_.request_stop(), watchdog_.join();
    }

    // ── State ─────────────────────────────────────────────────────────────────

    std::map<std::string, INode*>                          nodes_;
    std::map<std::string, std::vector<std::string>>        adj_;
    std::vector<std::string>                               topo_;
    std::map<std::string, std::string>                     exposed_inputs_;
    std::map<std::string, std::string>                     exposed_outputs_;
    std::set<std::pair<std::string, std::size_t>>          connected_outputs_;
    std::vector<std::unique_ptr<IChannelProbe>>            channel_probes_;
    std::vector<std::pair<std::string, IPoolProbe*>>       pool_probes_;
    ErrorHandler                                           error_handler_;
    DiagnosticsHandler                                     diag_handler_;
    std::chrono::milliseconds                              watchdog_interval_{3000};
    std::jthread                                           watchdog_;
    clock_t::time_point                                    start_time_;
#ifdef KPN_WEB_DEBUG
    uint16_t                                               web_debug_port_{9090};
    std::unique_ptr<web_debug::WebDebugServer>             web_server_;
#endif
};

} // namespace kpn