// SPDX-License-Identifier: MIT
// SPDX-FileCopyrightText: 2026 Martino Ferrari

#include "waveform.hpp"

#include <algorithm>
#include <cerrno>
#include <cmath>
#include <cstdlib>
#include <cstring>
#include <sys/stat.h>

#include "layout_store.hpp"

namespace udpscope {
namespace {

/// Reads `'<f4'`, `'<f8'` … out of one dtype entry and returns the bytes per element, 0 for a type
/// this loader does not know. The writer emits little-endian only, and the check that the host is
/// little-endian is the writer's own (`static_assert` there).
std::size_t dtype_bytes(const std::string& type) {
  if (type == "<f8" || type == "<i8" || type == "<u8") return 8;
  if (type == "<f4" || type == "<i4" || type == "<u4") return 4;
  if (type == "<i2" || type == "<u2") return 2;
  if (type == "|i1" || type == "|u1" || type == "|b1") return 1;
  return 0;
}

/// The field list of a structured dtype: `[('time', '<f8'), ('ch1', '<f4')]` -> {name, type} pairs.
/// Deliberately a scanner rather than a parser: a `.npy` header is a Python dict literal, and the
/// three keys this reader needs are the only thing it may trust.
bool parse_fields(const std::string& descr, std::vector<std::pair<std::string, std::string>>* out) {
  std::size_t i = 0;
  while (true) {
    const std::size_t open = descr.find("('", i);
    if (open == std::string::npos) break;
    std::size_t cursor = open + 2;
    std::string name;
    while (cursor < descr.size() && descr[cursor] != '\'') name.push_back(descr[cursor++]);
    if (cursor >= descr.size()) return false;
    ++cursor;  // past the name's closing quote, which a comma should follow
    const std::size_t type_open = descr.find('\'', cursor);
    if (type_open == std::string::npos) return false;
    std::string type;
    cursor = type_open + 1;
    while (cursor < descr.size() && descr[cursor] != '\'') type.push_back(descr[cursor++]);
    if (cursor >= descr.size()) return false;
    out->emplace_back(name, type);
    i = cursor + 1;
  }
  return !out->empty();
}

}  // namespace

WaveformReader::~WaveformReader() { close(); }

void WaveformReader::close() noexcept {
  if (f_ != nullptr) {
    std::fclose(f_);
    f_ = nullptr;
  }
}

bool WaveformReader::open(const std::string& path, std::string* error) {
  close();
  f_ = std::fopen(path.c_str(), "rb");
  if (f_ == nullptr) {
    if (error != nullptr) *error = "cannot open '" + path + "': " + std::strerror(errno);
    return false;
  }
  char magic[6] = {};
  char version[2] = {};
  if (std::fread(magic, 1, 6, f_) != 6 || std::fread(version, 1, 2, f_) != 2 ||
      std::memcmp(magic, "\x93NUMPY", 6) != 0) {
    if (error != nullptr) *error = "'" + path + "' is not a NumPy file";
    close();
    return false;
  }
  std::uint8_t len_bytes[4] = {};
  std::size_t header_len = 0;
  if (version[0] == 1) {
    if (std::fread(len_bytes, 1, 2, f_) != 2) {
      if (error != nullptr) *error = "'" + path + "' has a truncated header";
      close();
      return false;
    }
    header_len = static_cast<std::size_t>(len_bytes[0]) | (std::size_t(len_bytes[1]) << 8);
  } else {
    // v2.0 (and v3.0, which is UTF-8 and the same layout): a 4-byte little-endian length.
    if (std::fread(len_bytes, 1, 4, f_) != 4) {
      if (error != nullptr) *error = "'" + path + "' has a truncated header";
      close();
      return false;
    }
    header_len = std::size_t(len_bytes[0]) | (std::size_t(len_bytes[1]) << 8) |
                 (std::size_t(len_bytes[2]) << 16) | (std::size_t(len_bytes[3]) << 24);
  }
  std::string header(header_len, '\0');
  if (header_len == 0 || std::fread(header.data(), 1, header_len, f_) != header_len) {
    if (error != nullptr) *error = "'" + path + "' has a truncated header";
    close();
    return false;
  }

  // `descr` is either the field list (a structured array, which is what this project writes) or a
  // single type string (a plain array, which carries no names and so cannot say what its columns
  // are).
  std::vector<std::pair<std::string, std::string>> fields;
  const std::size_t descr_at = header.find("'descr':");
  if (descr_at == std::string::npos) {
    if (error != nullptr) *error = "'" + path + "' has no dtype";
    close();
    return false;
  }
  const std::size_t descr_end = header.find('\'', descr_at + 9);
  const std::size_t bracket = header.find('(', descr_at);
  if (bracket != std::string::npos && (descr_end == std::string::npos || bracket < descr_end + 1)) {
    if (!parse_fields(header.substr(descr_at), &fields)) {
      if (error != nullptr) *error = "'" + path + "': could not read the column names";
      close();
      return false;
    }
  } else {
    if (error != nullptr) {
      *error = "'" + path + "' is a plain array, not a udpscope export: its columns have no names";
    }
    close();
    return false;
  }

  // The two time columns by name, the rest as signals -- and every one of them must be a type this
  // loader can widen.
  bool have_time = false;
  bool have_abs = false;
  record_bytes_ = 0;
  rows_read_ = 0;
  names_.clear();
  field_bytes_.clear();
  for (const auto& [name, type] : fields) {
    const std::size_t bytes = dtype_bytes(type);
    if (bytes == 0) {
      if (error != nullptr) *error = "'" + path + "' has a column of an unsupported type '" + type + "'";
      close();
      return false;
    }
    record_bytes_ += bytes;
    field_bytes_.push_back(bytes);
    if (name == "time") {
      have_time = true;
    } else if (name == "abs_seconds") {
      have_abs = true;
    } else {
      names_.push_back(name);
    }
  }
  if (!have_time || !have_abs) {
    if (error != nullptr) {
      *error = "'" + path + "' is not a udpscope export: no 'time' and 'abs_seconds' columns";
    }
    close();
    return false;
  }

  const std::size_t shape_at = header.find("'shape':");
  if (shape_at == std::string::npos) {
    if (error != nullptr) *error = "'" + path + "' has no shape";
    close();
    return false;
  }
  const std::size_t open_paren = header.find('(', shape_at);
  if (open_paren == std::string::npos) {
    if (error != nullptr) *error = "'" + path + "' has an unreadable shape";
    close();
    return false;
  }
  rows_ = static_cast<std::size_t>(std::strtoull(header.c_str() + open_paren + 1, nullptr, 10));

  // The rate comes from the first two stamps, and reading them needs the file positioned: peek, then
  // rewind. An export's time column is the ring's own stamps, which are uniform -- one dt describes
  // all of it.
  const long data_at = std::ftell(f_);
  if (rows_ >= 2) {
    std::vector<char> first(record_bytes_);
    if (std::fread(first.data(), 1, record_bytes_, f_) == record_bytes_) {
      double t0 = 0.0;
      std::memcpy(&t0, first.data(), 8);
      if (std::fread(first.data(), 1, record_bytes_, f_) == record_bytes_) {
        double t1 = 0.0;
        std::memcpy(&t1, first.data(), 8);
        sample_dt_ = std::abs(t1 - t0);
      }
    }
    std::fseek(f_, data_at, SEEK_SET);
  }
  return true;
}

bool WaveformReader::next(double* time, double* abs_seconds, float* values) {
  if (f_ == nullptr || record_bytes_ == 0) return false;
  if (rows_read_ >= rows_) return false;
  std::vector<char> record(record_bytes_);
  if (std::fread(record.data(), 1, record_bytes_, f_) != record_bytes_) return false;
  const char* p = record.data();
  std::memcpy(time, p, 8);
  p += 8;
  std::memcpy(abs_seconds, p, 8);
  p += 8;
  // The signals are the fields after the two time columns, each at its own width: `write_npy()`
  // emits `<f4`, and a file an external tool widened to `<f8` is read as one rather than mis-sliced.
  for (std::size_t i = 0; i < names_.size(); ++i) {
    const std::size_t bytes = field_bytes_[2 + i];
    switch (bytes) {
      case 8: {
        double v = 0.0;
        std::memcpy(&v, p, 8);
        values[i] = static_cast<float>(v);
        break;
      }
      case 4: {
        float v = 0.0f;
        std::memcpy(&v, p, 4);
        values[i] = v;
        break;
      }
      default: {
        // A narrower integer column: widen through the integer types rather than reinterpreting.
        long long v = 0;
        switch (bytes) {
          case 2: {
            short s = 0;
            std::memcpy(&s, p, 2);
            v = s;
            break;
          }
          default:
            v = static_cast<signed char>(*p);
            break;
        }
        values[i] = static_cast<float>(v);
        break;
      }
    }
    p += bytes;
  }
  ++rows_read_;
  return true;
}

bool load_waveform_into_model(ScopeModel& model, const std::string& path,
                              WaveformLoadReport* report, std::string* error) {
  WaveformReader reader;
  if (!reader.open(path, error)) return false;
  if (reader.names().empty() || reader.rows() == 0) {
    if (error != nullptr) *error = "'" + path + "' holds no samples";
    return false;
  }
  const double dt = reader.sample_dt();
  const double rate = dt > 0.0 ? 1.0 / dt : 0.0;

  // No lock held across this: every model call below takes the model's own, and they are the ones
  // that must be atomic with each other -- holding it here deadlocked on the first of them.
  // The span first: the plan that the enables below build has to cover the whole file, or the oldest
  // samples are written past the ring's end and the file arrives as its own tail.
  std::vector<SignalInfo> signals;
  signals.reserve(reader.names().size());
  for (const std::string& name : reader.names()) {
    SignalInfo s;
    s.name = name;
    s.sampling_rate = rate > 0.0 ? rate : 1.0;
    signals.push_back(std::move(s));
  }
  model.set_target_span(reader.rows() * (dt > 0.0 ? dt : 0.001));
  model.set_signals(std::move(signals));
  for (std::size_t i = 0; i < reader.names().size(); ++i) model.set_signal_enabled(i, true);

  const std::size_t channels = model.channels().size();
  if (channels == 0) {
    if (error != nullptr) *error = "'" + path + "': the model built no channels for its columns";
    return false;
  }
  std::vector<float> values(reader.names().size(), 0.0f);
  std::size_t pushed = 0;
  double stamp = 0.0;
  double abs_stamp = 0.0;
  double first = 0.0;
  double last = 0.0;
  while (reader.next(&stamp, &abs_stamp, values.data())) {
    if (pushed == 0) first = stamp;
    last = stamp;
    model.push_loaded_slot(values, stamp, abs_stamp);
    ++pushed;
  }
  reader.close();

  model.set_epoch_if_unset(first);
  // The newest stamp is "now": the axis then reads `[-window, 0]`, which is the window the export was
  // taken from.
  model.set_display_epoch(pushed > 0 ? last : 0.0);

  if (report != nullptr) {
    report->signals = channels;
    report->samples = pushed;
    report->span_s = pushed > 0 ? std::abs(last - first) : 0.0;
    report->rate_hz = rate;
    report->first_time = first;
    report->last_time = last;
  }
  if (pushed == 0) {
    if (error != nullptr) *error = "'" + path + "' could not be read to the end";
    return false;
  }
  return true;
}

namespace {

std::string waveforms_path() {
  const std::string dir = config_dir();
  ::mkdir(dir.c_str(), 0755);  // EEXIST is the common case and not an error
  return dir + "/waveforms.txt";
}

}  // namespace

std::vector<std::string> recent_waveforms() {
  std::vector<std::string> out;
  std::FILE* f = std::fopen(waveforms_path().c_str(), "rb");
  if (f == nullptr) return out;
  char line[4096];
  while (std::fgets(line, sizeof(line), f) != nullptr) {
    std::string path(line);
    while (!path.empty() && (path.back() == '\n' || path.back() == '\r')) path.pop_back();
    if (!path.empty()) out.push_back(path);
  }
  std::fclose(f);
  return out;
}

void remember_waveform(const std::string& path) {
  std::vector<std::string> list = recent_waveforms();
  list.erase(std::remove(list.begin(), list.end(), path), list.end());
  list.insert(list.begin(), path);
  // A list long enough to be a history, short enough to be a menu: twenty entries is more than a
  // session's worth of exports and far less than a directory's.
  if (list.size() > 20) list.resize(20);
  std::FILE* f = std::fopen(waveforms_path().c_str(), "wb");
  if (f == nullptr) return;
  for (const std::string& p : list) std::fprintf(f, "%s\n", p.c_str());
  std::fclose(f);
}

}  // namespace udpscope