#include "source.h"
#include <algorithm>
#include <chrono>
#include <cstdio>
#include <unordered_map>
#include <unordered_set>

// ── Colour palette (matches WebUI) ───────────────────────────────────

static const float kPalette[][4] = {
    {0.122f, 0.467f, 0.706f, 1.f},  // #1f77b4
    {1.000f, 0.498f, 0.055f, 1.f},  // #ff7f0e
    {0.173f, 0.627f, 0.173f, 1.f},  // #2ca02c
    {0.839f, 0.153f, 0.157f, 1.f},  // #d62728
    {0.580f, 0.404f, 0.741f, 1.f},  // #9467bd
    {0.549f, 0.337f, 0.294f, 1.f},  // #8c564b
    {0.890f, 0.467f, 0.761f, 1.f},  // #e377c2
    {0.498f, 0.498f, 0.498f, 1.f},  // #7f7f7f
    {0.737f, 0.741f, 0.133f, 1.f},  // #bcbd22
    {0.090f, 0.745f, 0.812f, 1.f},  // #17becf
};
static constexpr size_t kPaletteSize =
    sizeof(kPalette) / sizeof(kPalette[0]);

// ── Source ────────────────────────────────────────────────────────────

Source::Source(std::string label, std::string host, uint16_t port,
               std::string multicast_group, uint16_t data_port)
    : label_(std::move(label))
    , host_(std::move(host))
    , port_(port)
    , client_(std::make_shared<UDPClient>(host_, port_,
                                          std::move(multicast_group), data_port))
{
    client_->on_config = [this](const std::vector<SignalDef>& defs) {
        handle_config(defs);
    };
    client_->on_data = [this](const DataPacket& dp, const std::vector<SignalDef>& defs) {
        handle_data(dp, defs);
    };
}

Source::~Source() {
    stop();
}

void Source::start()  { client_->start(); }
void Source::stop()   { client_->stop(); connected_ = false; }

Source::Stats Source::stats() const {
    Stats s;
    s.packets_rx   = stat_rx_.load(std::memory_order_relaxed);
    s.packets_lost = stat_lost_.load(std::memory_order_relaxed);
    s.bytes_rx     = stat_bytes_.load(std::memory_order_relaxed);
    return s;
}

// ── Ingest callbacks ──────────────────────────────────────────────────

void Source::handle_config(const std::vector<SignalDef>& defs) {
    std::lock_guard<std::mutex> lock(signals_mutex);

    // Reset time calibrations (indexed by protocol signal index)
    time_calibs_.assign(defs.size(), TimeCalib{});

    // Preserve render style for signal_states with the same name
    std::unordered_map<std::string, SignalState> old_by_name;
    for (auto& ss : signal_states)
        old_by_name[ss.def.name] = ss;
    signal_states.clear();
    sig_expand_.clear();
    signal_states.reserve(defs.size() * 2);  // may expand spatial arrays
    sig_expand_.reserve(defs.size());

    static size_t g_color_counter = 0;

    auto make_state = [&](SignalDef def, size_t cap) -> SignalState {
        if (cap < 4096)            cap = 4096;
        if (cap > MAX_SIGNAL_SAMPLES) cap = MAX_SIGNAL_SAMPLES;

        SignalState ss;
        ss.def    = std::move(def);
        ss.buffer  = std::make_shared<SignalBuffer>(cap);
        ss.archive = std::make_shared<SignalBuffer>(MAX_ARCHIVE_SAMPLES);
        ss.archive_last_t = -1e18;

        auto it = old_by_name.find(ss.def.name);
        if (it != old_by_name.end()) {
            const auto& old = it->second;
            ss.visible    = old.visible;
            ss.plot_mask  = old.plot_mask;
            ss.line_width = old.line_width;
            std::copy(old.color, old.color + 4, ss.color);
        } else {
            size_t ci = g_color_counter++ % kPaletteSize;
            std::copy(kPalette[ci], kPalette[ci] + 4, ss.color);
        }
        return ss;
    };

    // Collect protocol signal indices that are referenced as time sources.
    // These must NOT be expanded even if they are TM_PACKET arrays, because
    // other signals reference them by protocol index (time_signal_idx) and
    // their dp.values[] slot must remain a single "flat" array.
    std::unordered_set<uint32_t> time_refs;
    for (const auto& sd : defs)
        if (sd.time_signal_idx != NO_TIME_SIGNAL)
            time_refs.insert(sd.time_signal_idx);

    for (size_t si = 0; si < defs.size(); ++si) {
        const auto& sd = defs[si];

        // Spatial array (TM_PACKET with N>1 elements): expand into N scalar entries.
        // Each element gets its own ring buffer and appears as an independent signal.
        // Exception: do NOT expand time-reference signals — they must stay as one
        // entry so that time_signal_idx look-ups remain valid.
        bool is_time_ref = time_refs.count(static_cast<uint32_t>(si)) > 0;
        if (sd.num_elements > 1 && sd.time_mode == TM_PACKET && !is_time_ref) {
            int start_idx = static_cast<int>(signal_states.size());
            sig_expand_.push_back({start_idx, sd.num_elements});

            // Buffer sized for sampling_rate × 60 s; fall back to packet rate ~4096
            size_t elem_cap = 4096;
            if (sd.sampling_rate > 0)
                elem_cap = static_cast<size_t>(sd.sampling_rate * 60.0);

            for (uint32_t ei = 0; ei < sd.num_elements; ++ei) {
                SignalDef elem_def  = sd;
                elem_def.name       = sd.name + "[" + std::to_string(ei) + "]";
                elem_def.num_elements = 1;
                elem_def.num_rows     = 1;
                elem_def.num_cols     = 1;
                signal_states.push_back(make_state(std::move(elem_def), elem_cap));
            }
        } else {
            // Scalar or temporal array: one SignalState, buffer sized for full-rate data
            int start_idx = static_cast<int>(signal_states.size());
            sig_expand_.push_back({start_idx, 1});

            size_t cap = static_cast<size_t>(
                sd.sampling_rate * sd.num_elements * 60.0);
            signal_states.push_back(make_state(sd, cap));
        }
    }

    connected_ = true;
    config_dirty = true;
    std::fprintf(stderr, "[source %s] configured: %zu protocol signals → %zu states\n",
                 label_.c_str(), defs.size(), signal_states.size());
}

// Helper: write one sample to both the main buffer and the archive (time-decimated).
static inline void push_sample(SignalState& sig, double t, double v)
{
    sig.buffer->push(t, v);
    if (t - sig.archive_last_t >= Source::ARCHIVE_DT) {
        sig.archive->push(t, v);
        sig.archive_last_t = t;
    }
}

// Helper: write samples with packet-arrival timing (evenly spaced)
static void write_packet_timed(SignalState& sig, const std::vector<double>& vals,
                                double now_wall, double sample_rate)
{
    const size_t n = vals.size();
    if (n == 0) return;
    double dt = (sample_rate > 0 && n > 1) ? 1.0 / sample_rate : 0.0;
    double t0 = now_wall - dt * static_cast<double>(n > 0 ? n - 1 : 0);
    for (size_t i = 0; i < n; ++i)
        push_sample(sig, t0 + static_cast<double>(i) * dt, vals[i]);
}

void Source::handle_data(const DataPacket& dp, const std::vector<SignalDef>& defs) {
    // Use Unix epoch wall time — matches the WebUI Go client's time.Now().UnixNano()/1e9
    using clock = std::chrono::system_clock;
    double now_wall = std::chrono::duration<double>(
        clock::now().time_since_epoch()).count();

    // First pass: calibrate time signal_states on the very first packet.
    // Mirrors the Go client exactly:
    //   timer_to_sec = 1e-9 for uint64, 1e-6 for everything else
    //   calib_offset = wallT - timerVal * timer_to_sec
    for (size_t si = 0; si < defs.size() && si < dp.values.size(); ++si) {
        auto& calib = time_calibs_[si];
        if (calib.valid) continue;
        const auto& vals = dp.values[si];
        if (vals.empty()) continue;
        calib.timer_to_sec = (defs[si].type_code == TC_UINT64) ? 1e-9 : 1e-6;
        double timerS = vals[0] * calib.timer_to_sec;
        calib.calib_offset = now_wall - timerS;
        calib.valid        = true;
    }

    // Second pass: write signal data using sig_expand_ to map protocol signal
    // index → SignalState(s).
    std::lock_guard<std::mutex> lock(signals_mutex);

    for (size_t si = 0; si < defs.size() && si < sig_expand_.size()
                        && si < dp.values.size(); ++si) {
        const auto& def   = defs[si];
        const auto& expnd = sig_expand_[si];
        const auto& vals  = dp.values[si];
        if (vals.empty()) continue;

        if (expnd.num_elements > 1) {
            // Spatial array: write element[i] as a scalar sample into state[start+i]
            for (uint32_t ei = 0; ei < expnd.num_elements && ei < vals.size(); ++ei) {
                int idx = expnd.start_state_idx + static_cast<int>(ei);
                if (idx < static_cast<int>(signal_states.size()) &&
                    signal_states[idx].buffer)
                    push_sample(signal_states[idx], now_wall, vals[ei]);
            }
        } else {
            // Scalar or temporal array: existing time-mode logic
            int idx = expnd.start_state_idx;
            if (idx >= static_cast<int>(signal_states.size())) continue;
            auto& sig = signal_states[idx];
            if (!sig.buffer) continue;

            const size_t n = vals.size();
            bool handled = false;

            if (def.time_mode == TM_FULL_ARRAY &&
                def.time_signal_idx != NO_TIME_SIGNAL &&
                def.time_signal_idx < defs.size() &&
                def.time_signal_idx < dp.values.size()) {

                const auto& tvals = dp.values[def.time_signal_idx];
                const auto& calib = time_calibs_[def.time_signal_idx];
                size_t count = std::min(n, tvals.size());

                if (calib.valid) {
                    for (size_t i = 0; i < count; ++i) {
                        double wall = calib.calib_offset + tvals[i] * calib.timer_to_sec;
                        push_sample(sig, wall, vals[i]);
                    }
                    handled = true;
                }
            } else if (def.time_mode == TM_FIRST_SAMPLE &&
                       def.time_signal_idx != NO_TIME_SIGNAL &&
                       def.time_signal_idx < defs.size() &&
                       def.time_signal_idx < dp.values.size()) {

                const auto& tvals = dp.values[def.time_signal_idx];
                const auto& calib = time_calibs_[def.time_signal_idx];

                if (!tvals.empty() && calib.valid) {
                    double t0 = calib.calib_offset + tvals[0] * calib.timer_to_sec;
                    double dt = def.sampling_rate > 0 ? 1.0 / def.sampling_rate : 0.0;
                    for (size_t i = 0; i < n; ++i)
                        push_sample(sig, t0 + static_cast<double>(i) * dt, vals[i]);
                    handled = true;
                }
            } else if (def.time_mode == TM_LAST_SAMPLE &&
                       def.time_signal_idx != NO_TIME_SIGNAL &&
                       def.time_signal_idx < defs.size() &&
                       def.time_signal_idx < dp.values.size()) {

                const auto& tvals = dp.values[def.time_signal_idx];
                const auto& calib = time_calibs_[def.time_signal_idx];

                if (!tvals.empty() && calib.valid) {
                    double t_last = calib.calib_offset + tvals[0] * calib.timer_to_sec;
                    double dt = def.sampling_rate > 0 ? 1.0 / def.sampling_rate : 0.0;
                    double t0 = t_last - dt * static_cast<double>(n > 0 ? n - 1 : 0);
                    for (size_t i = 0; i < n; ++i)
                        push_sample(sig, t0 + static_cast<double>(i) * dt, vals[i]);
                    handled = true;
                }
            }

            if (!handled) {
                // TM_PACKET scalar or fallback: evenly space samples at packet arrival
                write_packet_timed(sig, vals, now_wall, def.sampling_rate);
            }
        }
    }

    stat_rx_++;
}
