#pragma once
// Copyright (c) 2024-2026 Tarik Moussa.
// SPDX-License-Identifier: MIT

// serialization.hpp
//
// Phase 5 — Save and load conformal map results in JSON and XML formats.
//
// JSON (nlohmann/json, bundled in deps/single_includes/json.hpp):
//   save_result_json / load_result_json
//
// XML (hand-written, no external parser dependency):
//   save_result_xml / load_result_xml
//
// Both formats store:
//   - geometry type string ("euclidean" / "spherical" / "hyper_ideal")
//   - mesh statistics (vertex/face count)
//   - solver metadata (converged, iterations, grad_inf_norm)
//   - DOF vector x
//   - optional 2D or 3D layout positions
//
// XML schema (ConformalResult):
//
//   <?xml version="1.0" encoding="UTF-8"?>
//   <ConformalResult geometry="euclidean" vertices="4" faces="2">
//     <Solver converged="true" iterations="3" grad_inf_norm="1.43e-13"/>
//     <DOFVector n="3">0.0 1.23e-13 2.87e-13</DOFVector>
//     <Layout dim="2" n="4">
//       0.0 0.0  1.0 0.0  0.5 0.866  1.5 0.866
//     </Layout>
//   </ConformalResult>

#include "newton_solver.hpp"
#include "layout.hpp"
#include <json.hpp>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#include <stdexcept>
#include <iomanip>

namespace conformallab {

// ════════════════════════════════════════════════════════════════════════════
// JSON
// ════════════════════════════════════════════════════════════════════════════

/// Save the Newton-solver result (+ optional 2-D layout) to a JSON file.
inline void save_result_json(
    const std::string& path,
    const NewtonResult& res,
    const std::string&  geometry,
    int                 n_vertices,
    int                 n_faces,
    const Layout2D*     layout2d = nullptr,
    const Layout3D*     layout3d = nullptr)
{
    using json = nlohmann::json;
    json j;
    j["geometry"]  = geometry;
    j["mesh"]      = { {"vertices", n_vertices}, {"faces", n_faces} };
    j["solver"]    = {
        {"converged",      res.converged},
        {"iterations",     res.iterations},
        {"grad_inf_norm",  res.grad_inf_norm}
    };
    j["dof_vector"] = res.x;

    if (layout2d && layout2d->success) {
        json uv = json::array();
        for (auto& p : layout2d->uv) uv.push_back({p.x(), p.y()});
        j["layout"] = { {"dim", 2}, {"uv", uv}, {"has_seam", layout2d->has_seam} };
    }
    if (layout3d && layout3d->success) {
        json pos = json::array();
        for (auto& p : layout3d->pos) pos.push_back({p.x(), p.y(), p.z()});
        j["layout"] = { {"dim", 3}, {"pos", pos}, {"has_seam", layout3d->has_seam} };
    }

    std::ofstream ofs(path);
    if (!ofs) throw std::runtime_error("Cannot write: " + path);
    ofs << std::setw(2) << j << "\n";
}

/// Load a DOF vector from a JSON result file written by
/// `save_result_json`.  If `res` is non-null its fields are filled too.
inline std::vector<double> load_result_json(
    const std::string& path,
    NewtonResult*       res     = nullptr,
    std::string*        geom    = nullptr,
    Layout2D*           layout2d = nullptr)
{
    using json = nlohmann::json;
    std::ifstream ifs(path);
    if (!ifs) throw std::runtime_error("conformallab: cannot open: " + path);

    // V1: wrap parse + field extraction so nlohmann exceptions (parse_error,
    // type_error, out_of_range) surface as std::runtime_error with the path.
    json j;
    try {
        ifs >> j;
    } catch (const json::exception& e) {
        throw std::runtime_error(
            "conformallab: malformed JSON in " + path + ": " + e.what());
    }

    try {
        if (geom && j.contains("geometry"))
            *geom = j["geometry"].get<std::string>();

        // V4: validate required top-level key before accessing it.
        if (!j.contains("dof_vector"))
            throw std::runtime_error(
                "conformallab: result JSON missing field 'dof_vector' in " + path);
        std::vector<double> x = j.at("dof_vector").get<std::vector<double>>();

        if (res) {
            // V4: validate nested solver keys before accessing.
            if (!j.contains("solver"))
                throw std::runtime_error(
                    "conformallab: result JSON missing field 'solver' in " + path);
            const auto& s = j.at("solver");
            for (const char* key : {"converged", "iterations", "grad_inf_norm"}) {
                if (!s.contains(key))
                    throw std::runtime_error(
                        std::string("conformallab: result JSON missing field 'solver.")
                        + key + "' in " + path);
            }
            res->x             = x;
            res->converged     = s.at("converged").get<bool>();
            res->iterations    = s.at("iterations").get<int>();
            res->grad_inf_norm = s.at("grad_inf_norm").get<double>();
        }

        if (layout2d && j.contains("layout") && j["layout"]["dim"] == 2) {
            auto uv = j["layout"]["uv"];
            layout2d->uv.resize(uv.size());
            for (std::size_t i = 0; i < uv.size(); ++i)
                layout2d->uv[i] = { uv[i][0].get<double>(), uv[i][1].get<double>() };
            layout2d->has_seam = j["layout"].value("has_seam", false);
            layout2d->success  = true;
        }

        return x;
    } catch (const json::exception& e) {
        throw std::runtime_error(
            "conformallab: malformed JSON in " + path + ": " + e.what());
    }
}

// ════════════════════════════════════════════════════════════════════════════
// XML
// ════════════════════════════════════════════════════════════════════════════

namespace detail_xml {

// Minimal XML attribute escaping
inline std::string xml_attr(double v)
{
    std::ostringstream s;
    s << std::setprecision(15) << v;
    return s.str();
}

// Write a flat vector of doubles as space-separated values
inline std::string flat_doubles(const std::vector<double>& v)
{
    std::ostringstream s;
    s << std::setprecision(15);
    for (std::size_t i = 0; i < v.size(); ++i) {
        if (i) s << ' ';
        s << v[i];
    }
    return s.str();
}

// Simple attribute parser: find value of key= in a tag string
inline std::string xml_get_attr(const std::string& tag, const std::string& key)
{
    auto pos = tag.find(key + "=\"");
    if (pos == std::string::npos) return {};
    pos += key.size() + 2;
    auto end = tag.find('"', pos);
    return tag.substr(pos, end - pos);
}

// Read all text content between the current position and </tag>
inline std::string xml_read_text(std::istream& is, const std::string& close_tag)
{
    std::string buf, line;
    std::string ctag = "</" + close_tag + ">";
    while (std::getline(is, line)) {
        auto pos = line.find(ctag);
        if (pos != std::string::npos) {
            buf += line.substr(0, pos);
            return buf;
        }
        buf += line + " ";
    }
    return buf;
}

// Parse space-separated doubles
inline std::vector<double> parse_doubles(const std::string& s)
{
    std::vector<double> v;
    std::istringstream ss(s);
    double d;
    while (ss >> d) v.push_back(d);
    return v;
}

} // namespace detail_xml

/// Save the Newton-solver result (+ optional layout) to an XML file.
inline void save_result_xml(
    const std::string& path,
    const NewtonResult& res,
    const std::string&  geometry,
    int                 n_vertices,
    int                 n_faces,
    const Layout2D*     layout2d = nullptr,
    const Layout3D*     layout3d = nullptr)
{
    std::ofstream ofs(path);
    if (!ofs) throw std::runtime_error("Cannot write: " + path);
    ofs << std::setprecision(15);

    ofs << "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n";
    ofs << "<ConformalResult"
        << " geometry=\"" << geometry << "\""
        << " vertices=\""  << n_vertices << "\""
        << " faces=\""     << n_faces << "\">\n";

    ofs << "  <Solver"
        << " converged=\""     << (res.converged ? "true" : "false") << "\""
        << " iterations=\""    << res.iterations << "\""
        << " grad_inf_norm=\"" << detail_xml::xml_attr(res.grad_inf_norm) << "\""
        << "/>\n";

    ofs << "  <DOFVector n=\"" << res.x.size() << "\">"
        << detail_xml::flat_doubles(res.x)
        << "</DOFVector>\n";

    if (layout2d && layout2d->success) {
        ofs << "  <Layout dim=\"2\" n=\"" << layout2d->uv.size()
            << "\" has_seam=\"" << (layout2d->has_seam ? "true" : "false") << "\">\n";
        for (auto& p : layout2d->uv)
            ofs << "    " << p.x() << " " << p.y() << "\n";
        ofs << "  </Layout>\n";
    }
    if (layout3d && layout3d->success) {
        ofs << "  <Layout dim=\"3\" n=\"" << layout3d->pos.size()
            << "\" has_seam=\"" << (layout3d->has_seam ? "true" : "false") << "\">\n";
        for (auto& p : layout3d->pos)
            ofs << "    " << p.x() << " " << p.y() << " " << p.z() << "\n";
        ofs << "  </Layout>\n";
    }

    ofs << "</ConformalResult>\n";
}

/// Load a DOF vector from an XML result file written by
/// `save_result_xml`.  If `res`, `geom`, `layout2d` are non-null they
/// are filled as well.
///
/// V5 (input-validation audit, 2026-06-01): this reader implements a
/// **strict internal-only XML subset** — not a general XML parser.  It
/// expects the exact one-element-per-line layout written by
/// `save_result_xml`.  Files that are semantically equivalent XML but
/// formatted differently (attributes split across lines, extra
/// whitespace, XML declaration on its own line, etc.) are explicitly
/// *rejected* with `std::runtime_error` rather than silently mis-read
/// into zeros.  Interoperability with other XML producers is out of
/// scope; use the JSON format for that.
///
/// Strict-subset requirements that are validated:
///  1. A line containing `<ConformalResult` must also carry a `geometry=`
///     attribute on the same line.
///  2. A line containing `<Solver` must carry `iterations=` and
///     `grad_inf_norm=` on the same line (when `res` is non-null).
///  3. A line containing `<DOFVector` must carry the `>` character (tag
///     open) on the same line.
///  4. The `<DOFVector` element must be present and must produce a
///     non-empty doubles list (a missing DOFVector silently returns an
///     empty x, which is incorrect for any mesh with at least one DOF).
inline std::vector<double> load_result_xml(
    const std::string& path,
    NewtonResult*       res      = nullptr,
    std::string*        geom     = nullptr,
    Layout2D*           layout2d = nullptr)
{
    std::ifstream ifs(path);
    if (!ifs) throw std::runtime_error("Cannot open: " + path);

    std::vector<double> x;
    std::string line;
    bool found_root      = false;
    bool found_dofvector = false;

    while (std::getline(ifs, line)) {
        // Root element — V5: geometry attribute must be on the same line.
        if (line.find("<ConformalResult") != std::string::npos) {
            found_root = true;
            // V5: reject if the required geometry= attribute is absent on this line.
            // (Would be present if written by save_result_xml; absent if reformatted.)
            std::string g = detail_xml::xml_get_attr(line, "geometry");
            if (g.empty())
                throw std::runtime_error(
                    "conformallab: XML strict-subset violation in " + path
                    + ": <ConformalResult geometry=...> attribute not found on its"
                    " opening line. Only the format written by save_result_xml is"
                    " supported — reformatted XML is rejected to prevent silent"
                    " misreads. Use the JSON format for interoperability.");
            if (geom) *geom = g;
        }
        // Solver metadata — V5: required attributes must be on the same line.
        else if (line.find("<Solver") != std::string::npos) {
            if (res) {
                res->converged = (detail_xml::xml_get_attr(line, "converged") == "true");
                // V2: stoi/stod throw std::invalid_argument on empty or non-numeric
                // attribute values; wrap and rethrow as runtime_error with context.
                try {
                    auto iter_str = detail_xml::xml_get_attr(line, "iterations");
                    auto grad_str = detail_xml::xml_get_attr(line, "grad_inf_norm");
                    if (iter_str.empty())
                        throw std::runtime_error("missing attribute 'iterations'");
                    if (grad_str.empty())
                        throw std::runtime_error("missing attribute 'grad_inf_norm'");
                    res->iterations    = std::stoi(iter_str);
                    res->grad_inf_norm = std::stod(grad_str);
                } catch (const std::exception& e) {
                    throw std::runtime_error(
                        "conformallab: malformed XML Solver element in "
                        + path + ": " + e.what());
                }
            }
        }
        // DOF vector — V5: the '>' tag-open must be on the same line.
        else if (line.find("<DOFVector") != std::string::npos) {
            found_dofvector = true;
            // V5: require the tag to be closed ('>') on the same line so the
            // content-extraction below works correctly.
            auto open_end = line.find('>');
            if (open_end == std::string::npos)
                throw std::runtime_error(
                    "conformallab: XML strict-subset violation in " + path
                    + ": <DOFVector> opening '>' not on same line as tag."
                    " Only the format written by save_result_xml is supported.");
            auto close    = line.find("</DOFVector>");
            std::string text;
            if (close != std::string::npos) {
                text = line.substr(open_end + 1, close - open_end - 1);
            } else {
                text = line.substr(open_end + 1);
                text += detail_xml::xml_read_text(ifs, "DOFVector");
            }
            x = detail_xml::parse_doubles(text);
            if (res) res->x = x;
        }
        // Layout
        else if (layout2d && line.find("<Layout") != std::string::npos
                          && detail_xml::xml_get_attr(line, "dim") == "2") {
            layout2d->has_seam = (detail_xml::xml_get_attr(line, "has_seam") == "true");
            std::string text   = detail_xml::xml_read_text(ifs, "Layout");
            auto vals          = detail_xml::parse_doubles(text);
            layout2d->uv.clear();
            for (std::size_t i = 0; i + 1 < vals.size(); i += 2)
                layout2d->uv.push_back({vals[i], vals[i+1]});
            layout2d->success = true;
        }
    }

    // V5: if the file was non-empty but never produced a <ConformalResult> root
    // element, the file is likely reformatted or not a ConformalResult XML at all.
    if (!found_root) {
        // Distinguish "empty file" (ifs.peek() == EOF at open) from wrong format.
        // We re-open to check file size — if it had content but no root element
        // was found on a single line, it was reformatted.
        std::ifstream probe(path, std::ios::ate);
        if (probe && probe.tellg() > 0)
            throw std::runtime_error(
                "conformallab: XML strict-subset violation in " + path
                + ": <ConformalResult> root element not found on its own line."
                " Only the format written by save_result_xml is supported.");
    }

    return x;
}

/// Validate that a loaded DOF vector has the expected number of DOFs.
///
/// V6 (input-validation audit, 2026-06-01): a result file from a *different*
/// mesh loads happily; the size mismatch only surfaces later (out-of-bounds
/// or wrong-answer) when `x` is indexed against the new mesh.  This helper
/// provides a clear early check at the call-site where the loaded vector is
/// paired with the mesh.
///
/// Throws `std::runtime_error` if `x.size() != expected_dofs`.
inline void check_dof_vector_size(
    const std::vector<double>& x,
    int                         expected_dofs,
    const std::string&          context = "")
{
    if (static_cast<int>(x.size()) != expected_dofs) {
        std::ostringstream msg;
        msg << "conformallab: DOF-vector size mismatch";
        if (!context.empty()) msg << " in " << context;
        msg << ": loaded " << x.size()
            << " values but mesh has " << expected_dofs << " DOFs.";
        throw std::runtime_error(msg.str());
    }
}

} // namespace conformallab
