#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("Cannot open: " + path);
    json j; ifs >> j;

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

    std::vector<double> x = j.at("dof_vector").get<std::vector<double>>();

    if (res) {
        res->x             = x;
        res->converged     = j["solver"]["converged"].get<bool>();
        res->iterations    = j["solver"]["iterations"].get<int>();
        res->grad_inf_norm = j["solver"]["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;
}

// ════════════════════════════════════════════════════════════════════════════
// 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.
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;

    while (std::getline(ifs, line)) {
        // Root element
        if (line.find("<ConformalResult") != std::string::npos) {
            if (geom) *geom = detail_xml::xml_get_attr(line, "geometry");
        }
        // Solver metadata
        else if (line.find("<Solver") != std::string::npos) {
            if (res) {
                res->converged     = (detail_xml::xml_get_attr(line, "converged") == "true");
                res->iterations    = std::stoi(detail_xml::xml_get_attr(line, "iterations"));
                res->grad_inf_norm = std::stod(detail_xml::xml_get_attr(line, "grad_inf_norm"));
            }
        }
        // DOF vector
        else if (line.find("<DOFVector") != std::string::npos) {
            // Text may be on same line: <DOFVector n="...">0 1 2...</DOFVector>
            auto open_end = line.find('>');
            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;
        }
    }
    return x;
}

} // namespace conformallab