#include <chrono>
#include <filesystem>
#include <fstream>
#include <string>
#include <vector>
#include <catch2/catch_test_macros.hpp>
#include <mfem.hpp>
#include <mpi.h>
import mean_field;
#include "../../experiments/continuation_snapshot.hpp"

namespace {
struct Manifest {
    struct Block { std::string stableId; int offset, size; };
    std::vector<Block> blocks{{"specific_enthalpy", 0, 1}, {"density", 1, 2}};
    const auto &valueBlocks() const { return blocks; }
};
struct Problem {
    Manifest manifest;
    MPI_Comm GetCommunicator() const { return MPI_COMM_SELF; }
    const auto &GetManifest() const { return manifest; }
    int StateSize() const { return 3; }
};
struct SnapshotFiles {
    const std::filesystem::path directory = std::filesystem::temp_directory_path() /
        ("meanfield-continuation-" + std::to_string(std::chrono::steady_clock::now().time_since_epoch().count()));
    SnapshotFiles() {
        std::filesystem::create_directories(directory / "capture_0");
        write("metadata.txt", "converged=1\n");
        write("layout.csv", "kind,block,offset,size,scale\nvalue,density,0,2,1\nvalue,specific_enthalpy,2,1,1\n");
        write("accepted_state.txt", "3\n2\n3\n7\n");
        write("capture_0/physical_base.txt", "3\n2\n3\n7\n");
    }
    ~SnapshotFiles() { std::error_code error; std::filesystem::remove_all(directory, error); }
    void write(const char *name, const std::string &contents) { std::ofstream(directory / name) << contents; }
};
}

TEST_CASE("Continuation file loading checks provenance and distinguishes captures from endpoints",
          "[solver][continuation_initial_state]") {
    using experiment::continuation::LoadMappedSnapshot;
    SnapshotFiles files;
    const Problem problem;
    auto mapped = LoadMappedSnapshot(problem, files.directory);
    REQUIRE(mapped.Size() == 3);
    CHECK(mapped(0) == 7.); CHECK(mapped(1) == 2.); CHECK(mapped(2) == 3.);
    CHECK_NOTHROW(LoadMappedSnapshot(problem, files.directory / "accepted_state.txt"));
    CHECK_THROWS_AS(LoadMappedSnapshot(problem, files.directory / "capture_0/physical_base.txt"), std::invalid_argument);
    CHECK_NOTHROW(LoadMappedSnapshot(problem, files.directory / "capture_0/physical_base.txt", false, true));
    files.write("metadata.txt", "converged=0\n");
    CHECK_THROWS_AS(LoadMappedSnapshot(problem, files.directory), std::invalid_argument);
    CHECK_NOTHROW(LoadMappedSnapshot(problem, files.directory, false, true));
    files.write("metadata.txt", "n=2\n");
    CHECK_THROWS_AS(LoadMappedSnapshot(problem, files.directory), std::invalid_argument);
    files.write("metadata.txt", "converged=1\nconverged=1\n");
    CHECK_THROWS_AS(LoadMappedSnapshot(problem, files.directory), std::invalid_argument);
    files.write("metadata.txt", "converged=1\n");
    for (const auto contents : {"2\n2\n3\n", "3\n2\n3\n7\nextra", "3\n2\n3\nnan\n"}) {
        files.write("accepted_state.txt", contents);
        CHECK_THROWS_AS(LoadMappedSnapshot(problem, files.directory), std::invalid_argument);
    }
    files.write("accepted_state.txt", "3\n2\n3\n7\n");
    files.write("layout.csv", "kind,block,offset,size,scale\nvalue,density,0,3,1\n");
    CHECK_THROWS_AS(LoadMappedSnapshot(problem, files.directory), std::invalid_argument);
}