#pragma once

// Include after import mean_field. These helpers transfer an initial guess;
// they do not certify that it solves the destination model. The caller binds
// model, mesh, field orders and surface-basis provenance before installation.
#include <algorithm>
#include <charconv>
#include <cmath>
#include <cstring>
#include <exception>
#include <filesystem>
#include <fstream>
#include <limits>
#include <map>
#include <set>
#include <span>
#include <stdexcept>
#include <string>
#include <string_view>
#include <utility>
#include <vector>

namespace experiment::continuation {
using SnapshotMetadata = std::map<std::string, std::string>;

struct SnapshotBlock final {
    std::string name;
    int offset;
    int size;
};

namespace detail {
inline constexpr std::string_view centralBorder = "fixed_central_density.border";

inline bool Finite(const mfem::Vector &vector) {
    for (int i = 0; i < vector.Size(); ++i)
        if (!std::isfinite(vector(i))) return false;
    return true;
}

inline bool Identical(const mfem::Vector &left, const mfem::Vector &right) {
    return left.Size() == right.Size() &&
        (left.Size() == 0 || std::memcmp(left.GetData(), right.GetData(),
            sizeof(mfem::real_t) * static_cast<std::size_t>(left.Size())) == 0);
}

inline void RequireSerial(MPI_Comm communicator) {
    int ranks = 0;
    if (communicator == MPI_COMM_NULL || MPI_Comm_size(communicator, &ranks) != MPI_SUCCESS || ranks != 1)
        throw std::invalid_argument("Continuation snapshots currently require one MPI rank.");
}

inline int NonnegativeInteger(std::string_view text) {
    int value = -1;
    const auto result = std::from_chars(text.data(), text.data() + text.size(), value);
    if (text.empty() || result.ec != std::errc{} || result.ptr != text.data() + text.size() || value < 0)
        throw std::invalid_argument("A snapshot layout offset or size is not a nonnegative integer.");
    return value;
}

inline std::vector<std::string> Fields(std::string line) {
    if (!line.empty() && line.back() == '\r') line.pop_back();
    if (line.find('"') != std::string::npos)
        throw std::invalid_argument("Snapshot layout requires the producer's unquoted CSV format.");
    std::vector<std::string> result;
    std::size_t first = 0;
    while (true) {
        const auto comma = line.find(',', first);
        result.push_back(line.substr(first, comma == std::string::npos ? comma : comma - first));
        if (comma == std::string::npos) return result;
        first = comma + 1;
    }
}

// Contiguity establishes an exact coordinate partition independently of file
// order. Equal total sizes alone cannot detect permuted or missing physics.
inline int ValidatePartition(std::span<const SnapshotBlock> blocks) {
    if (blocks.empty()) throw std::invalid_argument("Snapshot layout has no physical value blocks.");
    std::vector<SnapshotBlock> ordered(blocks.begin(), blocks.end());
    std::set<std::string> names;
    for (const auto &block : ordered) {
        if (block.name.empty() || !names.insert(block.name).second || block.offset < 0 || block.size < 0 ||
            block.offset > std::numeric_limits<int>::max() - block.size)
            throw std::invalid_argument("Snapshot layout has invalid or repeated block coordinates.");
        if (block.name == centralBorder && block.size != 1)
            throw std::invalid_argument("A central-density border must contain exactly one scalar.");
    }
    std::sort(ordered.begin(), ordered.end(), [](const auto &a, const auto &b) {
        return a.offset < b.offset || (a.offset == b.offset && a.size < b.size);
    });
    int size = 0;
    for (const auto &block : ordered) {
        if (block.offset != size)
            throw std::invalid_argument("Snapshot layout contains a gap or overlapping coordinates.");
        size += block.size;
    }
    if (size == 0) throw std::invalid_argument("Snapshot state is empty.");
    return size;
}

inline std::vector<SnapshotBlock> ReadValueBlocks(const std::filesystem::path &path) {
    std::ifstream input(path);
    std::string line;
    if (!input || !std::getline(input, line))
        throw std::invalid_argument("Cannot read snapshot layout: " + path.string());
    const auto header = Fields(line);
    std::map<std::string, std::size_t> columns;
    for (std::size_t i = 0; i < header.size(); ++i)
        if (header[i].empty() || !columns.emplace(header[i], i).second)
            throw std::invalid_argument("Snapshot layout has an invalid CSV header.");
    for (const auto *name : {"kind", "block", "offset", "size"})
        if (!columns.contains(name)) throw std::invalid_argument("Snapshot layout lacks a required column.");
    std::vector<SnapshotBlock> values, residuals;
    while (std::getline(input, line)) {
        const auto fields = Fields(line);
        if (fields.size() != header.size())
            throw std::invalid_argument("Snapshot layout has an incomplete CSV row.");
        const auto &kind = fields[columns.at("kind")];
        SnapshotBlock block{fields[columns.at("block")], NonnegativeInteger(fields[columns.at("offset")]),
                            NonnegativeInteger(fields[columns.at("size")])};
        if (kind == "value") values.push_back(std::move(block));
        else if (kind == "residual") residuals.push_back(std::move(block));
        else throw std::invalid_argument("Snapshot layout has an unknown coordinate kind.");
    }
    if (input.bad()) throw std::runtime_error("Snapshot layout read failed.");
    ValidatePartition(values);
    if (!residuals.empty()) ValidatePartition(residuals);
    return values;
}

inline mfem::Vector ReadVector(const std::filesystem::path &path, int expectedSize) {
    std::ifstream input(path);
    long long count = -1;
    if (!(input >> count) || count != expectedSize)
        throw std::invalid_argument("Snapshot vector length disagrees with its named layout: " + path.string());
    mfem::Vector values(expectedSize);
    for (int i = 0; i < expectedSize; ++i)
        if (!(input >> values(i)) || !std::isfinite(values(i)))
            throw std::invalid_argument("Snapshot vector is incomplete or nonfinite: " + path.string());
    std::string trailing;
    if (input >> trailing || input.bad())
        throw std::invalid_argument("Snapshot vector has trailing data or a read error: " + path.string());
    return values;
}
} // namespace detail

inline SnapshotMetadata ParseSnapshotMetadata(const std::filesystem::path &directory) {
    const auto path = directory / "metadata.txt";
    std::ifstream input(path);
    if (!input) throw std::invalid_argument("Continuation requires snapshot metadata: " + path.string());
    SnapshotMetadata metadata;
    std::string line;
    while (std::getline(input, line)) {
        if (!line.empty() && line.back() == '\r') line.pop_back();
        if (line.empty()) continue;
        const auto separator = line.find('=');
        if (separator == std::string::npos || separator == 0 ||
            !metadata.emplace(line.substr(0, separator), line.substr(separator + 1)).second)
            throw std::invalid_argument("Snapshot metadata has a malformed or repeated key.");
    }
    if (input.bad() || metadata.empty()) throw std::invalid_argument("Snapshot metadata is empty or unreadable.");
    return metadata;
}

inline mfem::Vector MapNamedPhysicalBlocks(
    const mfem::Vector &source, std::span<const SnapshotBlock> sourceBlocks,
    std::span<const SnapshotBlock> targetBlocks, bool allowCentralBorderChange = false
) {
    if (source.Size() != detail::ValidatePartition(sourceBlocks) || !detail::Finite(source))
        throw std::invalid_argument("Snapshot source does not match its finite physical coordinate layout.");
    mfem::Vector target(detail::ValidatePartition(targetBlocks));
    target = 0.0;
    std::map<std::string, const SnapshotBlock *> sourceByName, targetByName;
    for (const auto &block : sourceBlocks) sourceByName.emplace(block.name, &block);
    for (const auto &block : targetBlocks) targetByName.emplace(block.name, &block);
    for (const auto &block : sourceBlocks)
        if (!targetByName.contains(block.name) && !(allowCentralBorderChange && block.name == detail::centralBorder))
            throw std::invalid_argument("Snapshot contains an unknown destination block: " + block.name);
    for (const auto &block : targetBlocks) {
        const auto found = sourceByName.find(block.name);
        if (found == sourceByName.end()) {
            if (allowCentralBorderChange && block.name == detail::centralBorder) continue;
            throw std::invalid_argument("Snapshot is missing a required destination block: " + block.name);
        }
        const auto &from = *found->second;
        if (from.size != block.size)
            throw std::invalid_argument("Snapshot block has a different coordinate count: " + block.name);
        for (int i = 0; i < block.size; ++i) target(block.offset + i) = source(from.offset + i);
    }
    return target;
}

template <class Problem>
mfem::Vector LoadMappedSnapshot(const Problem &problem, const std::filesystem::path &path,
                               bool allowCentralBorderChange = false, bool allowUnconvergedRestart = false) {
    detail::RequireSerial(problem.GetCommunicator());
    const auto statePath = std::filesystem::is_directory(path) ? path / "accepted_state.txt" : path;
    const auto directory = std::filesystem::exists(statePath.parent_path() / "metadata.txt") ?
        statePath.parent_path() : statePath.parent_path().parent_path();
    const bool isAcceptedEndpoint = statePath.filename() == "accepted_state.txt" && statePath.parent_path() == directory;
    if (!isAcceptedEndpoint && !allowUnconvergedRestart)
        throw std::invalid_argument("Captured iterates require the explicit unconverged-restart override.");
    const auto metadata = ParseSnapshotMetadata(directory);
    const auto converged = metadata.find("converged");
    if (converged != metadata.end() && converged->second != "0" && converged->second != "1")
        throw std::invalid_argument("Snapshot metadata has an invalid convergence flag.");
    if (!allowUnconvergedRestart && (converged == metadata.end() || converged->second != "1"))
        throw std::invalid_argument("Continuation refuses an unverified or unconverged predecessor; an explicit restart override is required.");
    const auto sourceBlocks = detail::ReadValueBlocks(directory / "layout.csv");
    const auto source = detail::ReadVector(statePath, detail::ValidatePartition(sourceBlocks));
    std::vector<SnapshotBlock> targetBlocks;
    for (const auto &block : problem.GetManifest().valueBlocks())
        targetBlocks.push_back({std::string(block.stableId), block.offset, block.size});
    auto mapped = MapNamedPhysicalBlocks(source, sourceBlocks, targetBlocks, allowCentralBorderChange);
    if (mapped.Size() != problem.StateSize())
        throw std::logic_error("Destination manifest does not cover its physical state.");
    return mapped;
}

// No mesh or FEM preparation is needed: these cases catch the silent block
// reorder/border mistakes that a length-only snapshot reader would permit.
inline int RunSnapshotCoordinateSelfChecks() {
    int checks = 0;
    const auto require = [&](bool condition) {
        if (!condition) throw std::logic_error("Continuation coordinate self-check failed.");
        ++checks;
    };
    const auto rejects = [&](auto action) {
        bool rejected = false;
        try { action(); } catch (const std::invalid_argument &) { rejected = true; }
        require(rejected);
    };
    const std::vector<SnapshotBlock> source{{"density", 0, 2}, {"enthalpy", 2, 1}};
    const std::vector<SnapshotBlock> reordered{{"enthalpy", 0, 1}, {"density", 1, 2}};
    const std::vector<SnapshotBlock> bordered{{"enthalpy", 0, 1}, {std::string(detail::centralBorder), 1, 1}, {"density", 2, 2}};
    mfem::Vector values(3);
    values(0) = 2.; values(1) = 3.; values(2) = 7.;
    const auto mapped = MapNamedPhysicalBlocks(values, source, reordered);
    require(mapped(0) == 7. && mapped(1) == 2. && mapped(2) == 3.);
    const auto added = MapNamedPhysicalBlocks(values, source, bordered, true);
    require(added.Size() == 4 && added(0) == 7. && added(1) == 0. && added(2) == 2. && added(3) == 3.);
    require(detail::Identical(MapNamedPhysicalBlocks(added, bordered, source, true), values));
    rejects([&] { MapNamedPhysicalBlocks(values, source, bordered); });
    rejects([&] { MapNamedPhysicalBlocks(added, bordered, source); });
    auto bad = source; bad[1].name = "unknown";
    rejects([&] { MapNamedPhysicalBlocks(values, bad, source); });
    bad = source; bad[1].offset = 1;
    rejects([&] { MapNamedPhysicalBlocks(values, bad, source); });
    bad = source; bad[1].offset = 3;
    rejects([&] { MapNamedPhysicalBlocks(values, bad, source); });
    bad = source; bad[1].name = "density";
    rejects([&] { MapNamedPhysicalBlocks(values, bad, source); });
    bad = source; bad[0].size = 1; bad[1].offset = 1; bad[1].size = 2;
    rejects([&] { MapNamedPhysicalBlocks(values, source, bad); });
    values(0) = std::numeric_limits<double>::quiet_NaN();
    rejects([&] { MapNamedPhysicalBlocks(values, source, reordered); });
    return checks;
}
} // namespace experiment::continuation