#include <cassert>
#include <cmath>
#include <filesystem>
#include <iostream>
#include <memory>
#include <sstream>
#include <string>
#include <vector>

#include "json.hpp"

#include "NAM/get_dsp.h"
#include "NAM/registry.h"

namespace test_get_dsp
{
// Config
const std::string basicConfigStr =
  R"({"version": "0.5.4", "metadata": {"date": {"year": 2024, "month": 10, "day": 9, "hour": 18, "minute": 44, "second": 41}, "loudness": -37.8406867980957, "gain": 0.13508800804658277, "name": "Test LSTM", "modeled_by": "Steve", "gear_type": "amp", "gear_make": "Darkglass Electronics", "gear_model": "Microtubes 900 v2", "tone_type": "clean", "input_level_dbu": 18.3, "output_level_dbu": 12.3, "training": {"settings": {"ignore_checks": false}, "data": {"latency": {"manual": null, "calibration": {"algorithm_version": 1, "delays": [-16], "safety_factor": 1, "recommended": -17, "warnings": {"matches_lookahead": false, "disagreement_too_high": false}}}, "checks": {"version": 3, "passed": true}}, "validation_esr": null}}, "architecture": "LSTM", "config": {"input_size": 1, "hidden_size": 3, "num_layers": 1}, "weights": [-0.21677088737487793, -0.6683622002601624, -0.2560940980911255, -0.3588429093360901, 0.17952610552310944, 0.19445613026618958, -0.01662646047770977, 0.5353694558143616, -0.2536540627479553, -0.5132213234901428, -0.020476307719945908, 0.08592455089092255, -0.6891753673553467, 0.3627359867095947, 0.008421811275184155, 0.3113192617893219, 0.14251480996608734, 0.07989779114723206, -0.18211324512958527, 0.7118963003158569, 0.41084015369415283, -0.6571938395500183, -0.13214066624641418, -0.2698603868484497, 0.49387243390083313, -0.3491725027561188, 0.6353667974472046, -0.5005152225494385, 0.2052856683731079, -0.4301638901233673, -0.15770092606544495, -0.7181791067123413, 0.056290093809366226, -0.49049463868141174, 0.6623441576957703, 0.09029324352741241, 0.34005245566368103, 0.16416560113430023, 0.15520110726356506, -0.4155678153038025, -0.36928507685661316, 0.3211132884025574, -0.6769840121269226, -0.1575538069009781, 0.05268515646457672, -0.4191459119319916, 0.599330484867096, 0.21518059074878693, -4.246325492858887, -3.315647840499878, -4.328850746154785, 4.496089458465576, 5.015639305114746, 3.6492037773132324, 0.14431169629096985, -0.6633821725845337, 0.11673200130462646, -0.1418764889240265, -0.4897872805595398, -0.8689419031143188, -0.06714004278182983, -0.4450395107269287, -0.02142983116209507, -0.15136894583702087, -2.775207996368408, -0.08681213855743408, 0.05702732503414154, 0.670292317867279, 0.31442636251449585, 0.30793967843055725], "sample_rate": 48000})";

// Copied over but shouldn't be publicly-exposed.
std::vector<float> GetWeights(nlohmann::json const& j)
{
  auto it = j.find("weights");
  if (it != j.end())
  {
    return *it;
  }
  else
  {
    throw std::runtime_error("Corrupted model file is missing weights.");
  }
}

nam::dspData _GetConfig(const std::string& configStr = basicConfigStr)
{
  nlohmann::json j = nlohmann::json::parse(configStr);

  std::vector<float> weights = GetWeights(j);
  nam::dspData returnedConfig;
  returnedConfig.version = j["version"];
  returnedConfig.architecture = j["architecture"];
  returnedConfig.config = j["config"];
  returnedConfig.metadata = j["metadata"];
  returnedConfig.weights = weights;

  return returnedConfig;
}

namespace
{

constexpr const char* kConstructorResetArchitecture = "ConstructorResetPrewarmPolicyTest";
int gConstructorResetConstructCount = 0;

class ConstructorResetDSP : public nam::DSP
{
public:
  explicit ConstructorResetDSP(const double expected_sample_rate)
  : nam::DSP(1, 1, expected_sample_rate)
  {
    gConstructorResetConstructCount++;
    Reset(expected_sample_rate, 16);
  }

  void prewarm() override { prewarm_count++; }

  int prewarm_count = 0;
};

std::unique_ptr<nam::DSP> ConstructorResetFactory(const nlohmann::json& config, std::vector<float>& weights,
                                                  const double expectedSampleRate)
{
  (void)config;
  (void)weights;
  return std::make_unique<ConstructorResetDSP>(expectedSampleRate);
}

static nam::factory::Helper _register_ConstructorResetPrewarmPolicyTest(kConstructorResetArchitecture,
                                                                        ConstructorResetFactory);

nlohmann::json build_constructor_reset_config()
{
  return nlohmann::json{{"version", "0.7.0"},
                        {"metadata", nlohmann::json::object()},
                        {"architecture", kConstructorResetArchitecture},
                        {"config", nlohmann::json::object()},
                        {"weights", nlohmann::json::array()},
                        {"sample_rate", 48000}};
}

} // namespace

void test_gets_input_level()
{
  nam::dspData config = _GetConfig();
  std::unique_ptr<nam::DSP> dsp = get_dsp(config);
  assert(dsp->HasInputLevel());
}
void test_gets_output_level()
{
  nam::dspData config = _GetConfig();
  std::unique_ptr<nam::DSP> dsp = get_dsp(config);
  assert(dsp->HasOutputLevel());
}

void test_null_input_level()
{
  // Issue 129
  const std::string configStr =
    R"({"version": "0.5.4", "metadata": {"date": {"year": 2024, "month": 10, "day": 9, "hour": 18, "minute": 44, "second": 41}, "loudness": -37.8406867980957, "gain": 0.13508800804658277, "name": "Test LSTM", "modeled_by": "Steve", "gear_type": "amp", "gear_make": "Darkglass Electronics", "gear_model": "Microtubes 900 v2", "tone_type": "clean", "input_level_dbu": null, "output_level_dbu": 12.3, "training": {"settings": {"ignore_checks": false}, "data": {"latency": {"manual": null, "calibration": {"algorithm_version": 1, "delays": [-16], "safety_factor": 1, "recommended": -17, "warnings": {"matches_lookahead": false, "disagreement_too_high": false}}}, "checks": {"version": 3, "passed": true}}, "validation_esr": null}}, "architecture": "LSTM", "config": {"input_size": 1, "hidden_size": 3, "num_layers": 1}, "weights": [-0.21677088737487793, -0.6683622002601624, -0.2560940980911255, -0.3588429093360901, 0.17952610552310944, 0.19445613026618958, -0.01662646047770977, 0.5353694558143616, -0.2536540627479553, -0.5132213234901428, -0.020476307719945908, 0.08592455089092255, -0.6891753673553467, 0.3627359867095947, 0.008421811275184155, 0.3113192617893219, 0.14251480996608734, 0.07989779114723206, -0.18211324512958527, 0.7118963003158569, 0.41084015369415283, -0.6571938395500183, -0.13214066624641418, -0.2698603868484497, 0.49387243390083313, -0.3491725027561188, 0.6353667974472046, -0.5005152225494385, 0.2052856683731079, -0.4301638901233673, -0.15770092606544495, -0.7181791067123413, 0.056290093809366226, -0.49049463868141174, 0.6623441576957703, 0.09029324352741241, 0.34005245566368103, 0.16416560113430023, 0.15520110726356506, -0.4155678153038025, -0.36928507685661316, 0.3211132884025574, -0.6769840121269226, -0.1575538069009781, 0.05268515646457672, -0.4191459119319916, 0.599330484867096, 0.21518059074878693, -4.246325492858887, -3.315647840499878, -4.328850746154785, 4.496089458465576, 5.015639305114746, 3.6492037773132324, 0.14431169629096985, -0.6633821725845337, 0.11673200130462646, -0.1418764889240265, -0.4897872805595398, -0.8689419031143188, -0.06714004278182983, -0.4450395107269287, -0.02142983116209507, -0.15136894583702087, -2.775207996368408, -0.08681213855743408, 0.05702732503414154, 0.670292317867279, 0.31442636251449585, 0.30793967843055725], "sample_rate": 48000})";
  nam::dspData config = _GetConfig(configStr);
  // The first part of this is that the following line doesn't fail:
  std::unique_ptr<nam::DSP> dsp = get_dsp(config);

  assert(!dsp->HasInputLevel());
  assert(dsp->HasOutputLevel());
}

void test_null_output_level()
{
  // Issue 129
  const std::string configStr =
    R"({"version": "0.5.4", "metadata": {"date": {"year": 2024, "month": 10, "day": 9, "hour": 18, "minute": 44, "second": 41}, "loudness": -37.8406867980957, "gain": 0.13508800804658277, "name": "Test LSTM", "modeled_by": "Steve", "gear_type": "amp", "gear_make": "Darkglass Electronics", "gear_model": "Microtubes 900 v2", "tone_type": "clean", "input_level_dbu": 19.0, "output_level_dbu": null, "training": {"settings": {"ignore_checks": false}, "data": {"latency": {"manual": null, "calibration": {"algorithm_version": 1, "delays": [-16], "safety_factor": 1, "recommended": -17, "warnings": {"matches_lookahead": false, "disagreement_too_high": false}}}, "checks": {"version": 3, "passed": true}}, "validation_esr": null}}, "architecture": "LSTM", "config": {"input_size": 1, "hidden_size": 3, "num_layers": 1}, "weights": [-0.21677088737487793, -0.6683622002601624, -0.2560940980911255, -0.3588429093360901, 0.17952610552310944, 0.19445613026618958, -0.01662646047770977, 0.5353694558143616, -0.2536540627479553, -0.5132213234901428, -0.020476307719945908, 0.08592455089092255, -0.6891753673553467, 0.3627359867095947, 0.008421811275184155, 0.3113192617893219, 0.14251480996608734, 0.07989779114723206, -0.18211324512958527, 0.7118963003158569, 0.41084015369415283, -0.6571938395500183, -0.13214066624641418, -0.2698603868484497, 0.49387243390083313, -0.3491725027561188, 0.6353667974472046, -0.5005152225494385, 0.2052856683731079, -0.4301638901233673, -0.15770092606544495, -0.7181791067123413, 0.056290093809366226, -0.49049463868141174, 0.6623441576957703, 0.09029324352741241, 0.34005245566368103, 0.16416560113430023, 0.15520110726356506, -0.4155678153038025, -0.36928507685661316, 0.3211132884025574, -0.6769840121269226, -0.1575538069009781, 0.05268515646457672, -0.4191459119319916, 0.599330484867096, 0.21518059074878693, -4.246325492858887, -3.315647840499878, -4.328850746154785, 4.496089458465576, 5.015639305114746, 3.6492037773132324, 0.14431169629096985, -0.6633821725845337, 0.11673200130462646, -0.1418764889240265, -0.4897872805595398, -0.8689419031143188, -0.06714004278182983, -0.4450395107269287, -0.02142983116209507, -0.15136894583702087, -2.775207996368408, -0.08681213855743408, 0.05702732503414154, 0.670292317867279, 0.31442636251449585, 0.30793967843055725], "sample_rate": 48000})";
  nam::dspData config = _GetConfig(configStr);
  // The first part of this is that the following line doesn't fail:
  std::unique_ptr<nam::DSP> dsp = get_dsp(config);
  assert(dsp->HasInputLevel());
  assert(!dsp->HasOutputLevel());
}

// Helper function to process buffers through a DSP model
void process_buffers(nam::DSP* dsp, int num_buffers, int buffer_size)
{
  const int in_channels = dsp->NumInputChannels();
  const int out_channels = dsp->NumOutputChannels();
  const double sample_rate = dsp->GetExpectedSampleRate() > 0 ? dsp->GetExpectedSampleRate() : 48000.0;

  // Reset the model
  dsp->Reset(sample_rate, buffer_size);

  // Allocate buffers for all channels
  std::vector<std::vector<NAM_SAMPLE>> inputBuffers(in_channels);
  std::vector<std::vector<NAM_SAMPLE>> outputBuffers(out_channels);
  std::vector<NAM_SAMPLE*> inputPtrs(in_channels);
  std::vector<NAM_SAMPLE*> outputPtrs(out_channels);

  for (int ch = 0; ch < in_channels; ch++)
  {
    inputBuffers[ch].resize(buffer_size, (NAM_SAMPLE)0.0);
    inputPtrs[ch] = inputBuffers[ch].data();
  }
  for (int ch = 0; ch < out_channels; ch++)
  {
    outputBuffers[ch].resize(buffer_size, (NAM_SAMPLE)0.0);
    outputPtrs[ch] = outputBuffers[ch].data();
  }

  // Process num_buffers buffers
  for (int buf = 0; buf < num_buffers; buf++)
  {
    // Fill input with some test data (simple sine-like pattern)
    for (int ch = 0; ch < in_channels; ch++)
    {
      for (int i = 0; i < buffer_size; i++)
      {
        inputBuffers[ch][i] = (NAM_SAMPLE)(0.1 * (ch + 1) * ((buf * buffer_size + i) % 100) / 100.0);
      }
    }

    // Process the buffer
    dsp->process(inputPtrs.data(), outputPtrs.data(), buffer_size);

    // Verify output is finite
    for (int ch = 0; ch < out_channels; ch++)
    {
      for (int i = 0; i < buffer_size; i++)
      {
        assert(std::isfinite(outputBuffers[ch][i]));
      }
    }
  }
}

void test_load_and_process_nam_files()
{
  // Test loading and processing three different .nam files
  // Paths are relative to root directory where tests run (./build/tools/run_tests)
  const std::vector<std::string> nam_files = {"example_models/wavenet.nam", "example_models/lstm.nam",
                                              "example_models/wavenet_condition_dsp.nam",
                                              "example_models/wavenet_a2_max.nam"};

  const int num_buffers = 3;
  const int buffer_size = 64;

  for (const auto& nam_file : nam_files)
  {
    std::filesystem::path model_path(nam_file);

    // Load the model
    std::unique_ptr<nam::DSP> dsp = nam::get_dsp(model_path);
    assert(dsp != nullptr);

    // Process buffers through the model
    process_buffers(dsp.get(), num_buffers, buffer_size);
  }
}

// Helper function to create a config string with a specific version
std::string createConfigWithVersion(const std::string& version)
{
  nlohmann::json j = nlohmann::json::parse(basicConfigStr);
  j["version"] = version;
  return j.dump();
}

class DemoVersionSupportChecker : public nam::IVersionSupportChecker
{
public:
  nam::Supported support(const std::string& version) const override
  {
    const std::string prefix = "DEMO::";
    if (version.rfind(prefix, 0) != 0)
      return nam::Supported::NO;

    const std::string scopedVersion = version.substr(prefix.size());
    if (scopedVersion == "1.0.0")
      return nam::Supported::YES;
    if (scopedVersion.rfind("1.0.", 0) == 0)
      return nam::Supported::PARTIAL;
    return nam::Supported::NO;
  }
};

void test_version_patch_one_beyond_supported()
{
  // Test that a .nam file with version one patch beyond the latest fully supported
  // can still be loaded
  nam::Version latestVersion = nam::ParseVersion(nam::LATEST_FULLY_SUPPORTED_NAM_FILE_VERSION);
  latestVersion.patch++;
  const std::string configStr = createConfigWithVersion(latestVersion.toString());
  nam::dspData config = _GetConfig(configStr);

  // Suppress the warning message that gets printed to std::cerr
  std::streambuf* originalCerr = std::cerr.rdbuf();
  std::ostringstream nullStream;
  std::cerr.rdbuf(nullStream.rdbuf());

  // Should succeed (with a warning, but we suppress it)
  std::unique_ptr<nam::DSP> dsp = get_dsp(config);
  assert(dsp != nullptr);

  // Restore original cerr
  std::cerr.rdbuf(originalCerr);
}

void test_version_minor_one_beyond_supported()
{
  // Test that a .nam file with version one minor beyond the latest fully supported
  // cannot be loaded
  nam::Version latestVersion = nam::ParseVersion(nam::LATEST_FULLY_SUPPORTED_NAM_FILE_VERSION);
  latestVersion.minor++;
  latestVersion.patch = 0; // Reset patch when incrementing minor
  const std::string configStr = createConfigWithVersion(latestVersion.toString());
  nam::dspData config = _GetConfig(configStr);

  bool threw = false;
  try
  {
    std::unique_ptr<nam::DSP> dsp = get_dsp(config);
  }
  catch (const std::runtime_error&)
  {
    threw = true;
  }
  assert(threw);
}

void test_version_too_early()
{
  // Test that a .nam file with version too early (before earliest supported) cannot be loaded
  nam::Version earliestVersion = nam::ParseVersion(nam::EARLIEST_SUPPORTED_NAM_FILE_VERSION);
  earliestVersion.minor--; // Decrement minor to get a version before earliest supported
  const std::string configStr = createConfigWithVersion(earliestVersion.toString());
  nam::dspData config = _GetConfig(configStr);

  bool threw = false;
  try
  {
    std::unique_ptr<nam::DSP> dsp = get_dsp(config);
  }
  catch (const std::runtime_error&)
  {
    threw = true;
  }
  assert(threw);
}

void test_is_version_supported_core_behavior()
{
  assert(nam::is_version_supported(nam::LATEST_FULLY_SUPPORTED_NAM_FILE_VERSION) == nam::Supported::YES);

  nam::Version patchBeyondLatest = nam::ParseVersion(nam::LATEST_FULLY_SUPPORTED_NAM_FILE_VERSION);
  patchBeyondLatest.patch++;
  assert(nam::is_version_supported(patchBeyondLatest.toString()) == nam::Supported::PARTIAL);

  nam::Version minorBeyondLatest = nam::ParseVersion(nam::LATEST_FULLY_SUPPORTED_NAM_FILE_VERSION);
  minorBeyondLatest.minor++;
  minorBeyondLatest.patch = 0;
  assert(nam::is_version_supported(minorBeyondLatest.toString()) == nam::Supported::NO);
}

void test_register_custom_version_support_checker()
{
  nam::register_version_support_checker(std::make_shared<DemoVersionSupportChecker>());

  assert(nam::is_version_supported("DEMO::1.0.0") == nam::Supported::YES);
  assert(nam::is_version_supported("DEMO::1.0.3") == nam::Supported::PARTIAL);
  assert(nam::is_version_supported("DEMO::2.0.0") == nam::Supported::NO);
}

void test_get_dsp_default_allows_constructor_reset_prewarm()
{
  gConstructorResetConstructCount = 0;

  auto dsp = nam::get_dsp(build_constructor_reset_config());
  auto* typed = dynamic_cast<ConstructorResetDSP*>(dsp.get());

  assert(typed != nullptr);
  assert(gConstructorResetConstructCount == 1);
  assert(typed->prewarm_count == 1);
  assert(typed->GetPrewarmOnReset());
}

void test_get_dsp_default_inherits_scoped_prewarm_default()
{
  gConstructorResetConstructCount = 0;

  nam::ScopedPrewarmOnResetDefault scoped_default(false);
  auto dsp = nam::get_dsp(build_constructor_reset_config());
  auto* typed = dynamic_cast<ConstructorResetDSP*>(dsp.get());

  assert(typed != nullptr);
  assert(gConstructorResetConstructCount == 1);
  assert(typed->prewarm_count == 0);
  assert(!typed->GetPrewarmOnReset());
}

void test_get_dsp_prewarm_option_suppresses_constructor_reset_prewarm()
{
  gConstructorResetConstructCount = 0;

  nam::DspLoadOptions options;
  options.prewarm = false;
  auto dsp = nam::get_dsp(build_constructor_reset_config(), options);
  auto* typed = dynamic_cast<ConstructorResetDSP*>(dsp.get());

  assert(typed != nullptr);
  assert(gConstructorResetConstructCount == 1);
  assert(typed->prewarm_count == 0);
  assert(typed->GetPrewarmOnReset());

  typed->Reset(48000.0, 16);
  assert(typed->prewarm_count == 1);
}

void test_get_dsp_prewarm_option_forces_constructor_reset_prewarm()
{
  gConstructorResetConstructCount = 0;

  nam::ScopedPrewarmOnResetDefault scoped_default(false);
  nam::DspLoadOptions options;
  options.prewarm = true;
  auto dsp = nam::get_dsp(build_constructor_reset_config(), options);
  auto* typed = dynamic_cast<ConstructorResetDSP*>(dsp.get());

  assert(typed != nullptr);
  assert(gConstructorResetConstructCount == 1);
  assert(typed->prewarm_count == 1);
  assert(!typed->GetPrewarmOnReset());

  typed->Reset(48000.0, 16);
  assert(typed->prewarm_count == 1);
}

void test_get_dsp_with_returned_config_constructs_once()
{
  gConstructorResetConstructCount = 0;

  nam::dspData returned_config;
  auto dsp = nam::get_dsp(build_constructor_reset_config(), returned_config);

  assert(dsp != nullptr);
  assert(gConstructorResetConstructCount == 1);
  assert(returned_config.architecture == kConstructorResetArchitecture);
}
}; // namespace test_get_dsp
