/********************************************************************************
 *                                                                              *
 * This file is part of IfcOpenShell.                                           *
 *                                                                              *
 * IfcOpenShell is free software: you can redistribute it and/or modify         *
 * it under the terms of the Lesser GNU General Public License as published by  *
 * the Free Software Foundation, either version 3.0 of the License, or          *
 * (at your option) any later version.                                          *
 *                                                                              *
 * IfcOpenShell is distributed in the hope that it will be useful,              *
 * but WITHOUT ANY WARRANTY; without even the implied warranty of               *
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the                 *
 * Lesser GNU General Public License for more details.                          *
 *                                                                              *
 * You should have received a copy of the Lesser GNU General Public License     *
 * along with this program. If not, see <http://www.gnu.org/licenses/>.         *
 *                                                                              *
 ********************************************************************************/

#include "GeometryStreamer.h"
#include "AppSettings.h"
#include "../ifcgeom/hybrid_kernel.h"
#include "../ifcgeom/taxonomy.h"

#include <Eigen/Dense>

#include <thread>
#include <unordered_map>
#include <cmath>
#include <cstring>
#include <algorithm>
#include <limits>

#include <QDebug>
#include <QElapsedTimer>

struct MaterialInfo {
    float r = 0.75f, g = 0.75f, b = 0.78f, a = 1.0f;
};

static MaterialInfo materialFromStyle(const ifcopenshell::geometry::taxonomy::style::ptr& style) {
    MaterialInfo m;
    if (!style) return m;
    const auto& color = style->get_color();
    if (color) {
        m.r = static_cast<float>(color.r());
        m.g = static_cast<float>(color.g());
        m.b = static_cast<float>(color.b());
    }
    if (!std::isnan(style->transparency)) {
        m.a = 1.0f - static_cast<float>(style->transparency);
    }
    return m;
}

static inline uint32_t packRGBA8(const MaterialInfo& m) {
    auto to_byte = [](float v) -> uint32_t {
        float c = std::clamp(v, 0.0f, 1.0f);
        return static_cast<uint32_t>(c * 255.0f + 0.5f);
    };
    uint32_t r = to_byte(m.r);
    uint32_t g = to_byte(m.g);
    uint32_t b = to_byte(m.b);
    uint32_t a = to_byte(m.a);
    // Little-endian byte layout [r,g,b,a] for GL_UNSIGNED_BYTE * 4 normalized.
    return r | (g << 8) | (b << 16) | (a << 24);
}

GeometryStreamer::GeometryStreamer(QObject* parent)
    : QObject(parent)
{
}

GeometryStreamer::~GeometryStreamer() {
    cancel();
    if (worker_thread_ && worker_thread_->isRunning()) {
        worker_thread_->quit();
        worker_thread_->wait();
    }
}

void GeometryStreamer::setIfcFile(std::unique_ptr<ifcopenshell::file> file) {
    ifc_file_ = std::move(file);
}

void GeometryStreamer::loadFile(const std::string& path, uint32_t start_object_id, uint32_t model_id, int num_threads) {
    if (running_.load()) {
        cancel();
        if (worker_thread_ && worker_thread_->isRunning()) {
            worker_thread_->quit();
            worker_thread_->wait();
        }
    }

    cancel_requested_ = false;
    succeeded_ = false;
    running_ = true;
    progress_ = 0;
    next_object_id_ = start_object_id;
    model_id_ = model_id;

    {
        std::lock_guard<std::mutex> lock(elements_mutex_);
        pending_elements_.clear();
    }

    if (num_threads <= 0) {
        num_threads = std::max(1u, std::thread::hardware_concurrency());
    }

    worker_thread_ = std::make_unique<QThread>();
    QObject* context = new QObject();
    context->moveToThread(worker_thread_.get());

    connect(worker_thread_.get(), &QThread::started, context, [this, path, num_threads, context]() {
        run(path, num_threads);
        context->deleteLater();
        worker_thread_->quit();
    });

    connect(worker_thread_.get(), &QThread::finished, this, [this]() {
        running_ = false;
        if (succeeded_.load()) {
            emit finished();
        } else if (cancel_requested_.load()) {
            emit cancelled();
        }
    });

    worker_thread_->start();
}

void GeometryStreamer::cancel() {
    cancel_requested_ = true;
}

std::vector<ElementInfo> GeometryStreamer::drainElements() {
    std::lock_guard<std::mutex> lock(elements_mutex_);
    std::vector<ElementInfo> result;
    result.swap(pending_elements_);
    return result;
}

// Build a mesh chunk (local coords, 28-byte interleaved vertices) from a
// TriangulationElement. Per-vertex color is baked from material_ids so that
// triangulations with per-face materials still render correctly.
static MeshChunk buildMeshChunk(uint32_t model_id,
                                uint32_t local_mesh_id,
                                const IfcGeom::TriangulationElement* elem) {
    MeshChunk chunk;
    chunk.model_id = model_id;
    chunk.local_mesh_id = local_mesh_id;

    const auto& geom = elem->geometry();
    const auto& verts = geom.verts();
    const auto& faces = geom.faces();
    const auto& normals = geom.normals();
    const auto& materials = geom.materials();
    const auto& material_ids = geom.material_ids();

    if (verts.empty() || faces.empty()) return chunk;

    const size_t num_verts_src = verts.size() / 3;
    const size_t num_tris = faces.size() / 3;
    const bool have_per_tri_material = (material_ids.size() == num_tris);

    // Dedupe (original vertex index, material id) so vertices shared across
    // triangles of the same material stay shared; vertices spanning multiple
    // materials are split (per-face color demands it).
    auto make_key = [](uint32_t orig_idx, int mat_id) -> uint64_t {
        return (static_cast<uint64_t>(orig_idx) << 32) | static_cast<uint32_t>(mat_id);
    };

    std::unordered_map<uint64_t, uint32_t> remap;
    remap.reserve(num_verts_src);

    chunk.vertices.reserve(num_verts_src * INSTANCED_VERTEX_STRIDE_FLOATS);
    chunk.indices.reserve(faces.size());

    // Track local AABB as we emit vertices.
    float amin[3] = { std::numeric_limits<float>::max(),
                      std::numeric_limits<float>::max(),
                      std::numeric_limits<float>::max() };
    float amax[3] = { -std::numeric_limits<float>::max(),
                      -std::numeric_limits<float>::max(),
                      -std::numeric_limits<float>::max() };

    auto emit_vertex = [&](uint32_t orig_idx, int mat_id) -> uint32_t {
        const uint64_t key = make_key(orig_idx, mat_id);
        auto it = remap.find(key);
        if (it != remap.end()) return it->second;

        const uint32_t new_idx = static_cast<uint32_t>(
            chunk.vertices.size() / INSTANCED_VERTEX_STRIDE_FLOATS);

        float px = static_cast<float>(verts[orig_idx * 3 + 0]);
        float py = static_cast<float>(verts[orig_idx * 3 + 1]);
        float pz = static_cast<float>(verts[orig_idx * 3 + 2]);
        chunk.vertices.push_back(px);
        chunk.vertices.push_back(py);
        chunk.vertices.push_back(pz);
        if (px < amin[0]) amin[0] = px; if (px > amax[0]) amax[0] = px;
        if (py < amin[1]) amin[1] = py; if (py > amax[1]) amax[1] = py;
        if (pz < amin[2]) amin[2] = pz; if (pz > amax[2]) amax[2] = pz;

        if (orig_idx * 3 + 2 < normals.size()) {
            chunk.vertices.push_back(static_cast<float>(normals[orig_idx * 3 + 0]));
            chunk.vertices.push_back(static_cast<float>(normals[orig_idx * 3 + 1]));
            chunk.vertices.push_back(static_cast<float>(normals[orig_idx * 3 + 2]));
        } else {
            chunk.vertices.push_back(0.0f);
            chunk.vertices.push_back(1.0f);
            chunk.vertices.push_back(0.0f);
        }

        MaterialInfo m;
        if (mat_id >= 0 && mat_id < static_cast<int>(materials.size())) {
            m = materialFromStyle(materials[mat_id]);
        }
        uint32_t packed = packRGBA8(m);
        float packed_as_float;
        std::memcpy(&packed_as_float, &packed, sizeof(float));
        chunk.vertices.push_back(packed_as_float);

        remap.emplace(key, new_idx);
        return new_idx;
    };

    for (size_t t = 0; t < num_tris; ++t) {
        const int mat_id = have_per_tri_material ? material_ids[t] : -1;
        chunk.indices.push_back(emit_vertex(static_cast<uint32_t>(faces[t * 3 + 0]), mat_id));
        chunk.indices.push_back(emit_vertex(static_cast<uint32_t>(faces[t * 3 + 1]), mat_id));
        chunk.indices.push_back(emit_vertex(static_cast<uint32_t>(faces[t * 3 + 2]), mat_id));
    }

    if (chunk.vertices.empty()) {
        for (int a = 0; a < 3; ++a) amin[a] = amax[a] = 0.0f;
    }
    for (int a = 0; a < 3; ++a) {
        chunk.local_aabb_min[a] = amin[a];
        chunk.local_aabb_max[a] = amax[a];
    }
    return chunk;
}

// Compute the world-space AABB by transforming the 8 corners of the local
// AABB through the column-major 4x4 transform.
static void worldAabbFromLocal(const float local_min[3],
                               const float local_max[3],
                               const float M[16],
                               float out_min[3], float out_max[3]) {
    out_min[0] = out_min[1] = out_min[2] =  std::numeric_limits<float>::max();
    out_max[0] = out_max[1] = out_max[2] = -std::numeric_limits<float>::max();
    for (int c = 0; c < 8; ++c) {
        float x = (c & 1) ? local_max[0] : local_min[0];
        float y = (c & 2) ? local_max[1] : local_min[1];
        float z = (c & 4) ? local_max[2] : local_min[2];
        // Column-major: world = M * [x,y,z,1].
        float wx = M[0]*x + M[4]*y + M[8]*z  + M[12];
        float wy = M[1]*x + M[5]*y + M[9]*z  + M[13];
        float wz = M[2]*x + M[6]*y + M[10]*z + M[14];
        if (wx < out_min[0]) out_min[0] = wx; if (wx > out_max[0]) out_max[0] = wx;
        if (wy < out_min[1]) out_min[1] = wy; if (wy > out_max[1]) out_max[1] = wy;
        if (wz < out_min[2]) out_min[2] = wz; if (wz > out_max[2]) out_max[2] = wz;
    }
}

void GeometryStreamer::run(const std::string& path, int num_threads) {
    try {
        // read_only is a no-op for SPF; for RocksDB it allows concurrent
        // readers and avoids acquiring the exclusive DB lock.
        ifc_file_ = std::make_unique<ifcopenshell::file>(
            path, ifcopenshell::FT_AUTODETECT, /*read_only=*/true);
    } catch (const std::exception& e) {
        emit errorOccurred(QString("Failed to parse IFC file: %1").arg(e.what()));
        return;
    }

    ifcopenshell::geometry::Settings settings;
    // Instancing path: geometry stays in local coords; the transform is
    // applied on the GPU per instance.
    settings.set("use-world-coords", false);
    settings.set("weld-vertices", false);
    settings.set("apply-default-materials", true);
    // Off by default in IfcOpenShell — makes face winding consistent within
    // each shell, which we need for GL_CULL_FACE and for per-vertex normals
    // to shade a solid without dark inside-out patches.  Costs some iterator
    // time, but results are cached in the sidecar so it's a one-shot hit.
    settings.set("reorient-shells", true);

    // @todo parallel mapping on RocksDB-backed files still races somewhere
    // outside the instance cache, producing inconsistent shape counts. Force
    // serial iteration for RocksDB until the read path is fully thread-safe.
    const bool is_rocksdb = std::holds_alternative<ifcopenshell::impl::rocks_db_file_storage>(ifc_file_->storage_);
    const int effective_threads = is_rocksdb ? 1 : num_threads;

    std::unique_ptr<IfcGeom::Iterator> iterator;
    try {
        const std::string geometry_library =
            AppSettings::instance().geometryLibrary().toStdString();
        auto kernel = ifcopenshell::geometry::kernels::construct(
            ifc_file_.get(), geometry_library, settings);
        iterator = std::make_unique<IfcGeom::Iterator>(
            std::move(kernel), settings, ifc_file_.get(),
            std::vector<ifcopenshell::geometry::filter_t>(), effective_threads);
    } catch (const std::exception& e) {
        emit errorOccurred(QString("Failed to create geometry iterator: %1").arg(e.what()));
        return;
    }

    if (!iterator->initialize()) {
        emit errorOccurred("No geometry found in IFC file");
        return;
    }

    int last_progress = 0;

    // geom.id() → local_mesh_id within this model.
    std::unordered_map<std::string, uint32_t> geom_to_local_mesh_id;
    // local_mesh_id → (local AABB) so we can derive world AABBs for later instances.
    struct MeshAabb { float lmin[3], lmax[3]; };
    std::vector<MeshAabb> mesh_aabbs;

    uint32_t total_shapes = 0;
    uint32_t total_meshes = 0;
    QElapsedTimer stream_timer;
    stream_timer.start();

    do {
        if (cancel_requested_.load()) break;

        const IfcGeom::Element* elem = iterator->get();
        if (!elem) continue;

        const auto* tri_elem = dynamic_cast<const IfcGeom::TriangulationElement*>(elem);
        if (!tri_elem) continue;

        const auto& geom = tri_elem->geometry();
        if (geom.verts().empty() || geom.faces().empty()) continue;

        uint32_t object_id = next_object_id_++;

        // Element metadata.
        ElementInfo info;
        info.object_id = object_id;
        info.model_id = model_id_;
        info.ifc_id = tri_elem->id();
        info.guid = tri_elem->guid();
        info.name = tri_elem->name();
        info.type = tri_elem->type();
        info.parent_id = tri_elem->parent_id();
        {
            std::lock_guard<std::mutex> lock(elements_mutex_);
            pending_elements_.push_back(std::move(info));
        }

        // Representation dedup.
        const std::string& geom_id = geom.id();
        uint32_t local_mesh_id;
        bool first_sight = false;
        if (geom_id.empty()) {
            // No representation key — treat as unique.
            local_mesh_id = total_meshes++;
            first_sight = true;
        } else {
            auto it = geom_to_local_mesh_id.find(geom_id);
            if (it == geom_to_local_mesh_id.end()) {
                local_mesh_id = total_meshes++;
                geom_to_local_mesh_id.emplace(geom_id, local_mesh_id);
                first_sight = true;
            } else {
                local_mesh_id = it->second;
            }
        }

        if (first_sight) {
            MeshChunk mesh_chunk = buildMeshChunk(model_id_, local_mesh_id, tri_elem);
            MeshAabb ma;
            for (int a = 0; a < 3; ++a) {
                ma.lmin[a] = mesh_chunk.local_aabb_min[a];
                ma.lmax[a] = mesh_chunk.local_aabb_max[a];
            }
            if (mesh_aabbs.size() <= local_mesh_id) mesh_aabbs.resize(local_mesh_id + 1);
            mesh_aabbs[local_mesh_id] = ma;
            if (!mesh_chunk.indices.empty()) {
                emit meshReady(std::move(mesh_chunk));
            }
        }

        // Transform (column-major 4x4, cast to float).
        const Eigen::Matrix4d& mat_d = tri_elem->transformation().data()->ccomponents();
        InstanceChunk inst;
        inst.model_id = model_id_;
        inst.local_mesh_id = local_mesh_id;
        inst.object_id = object_id;
        inst.color_override_rgba8 = 0; // 0 = use baked vertex color
        for (int i = 0; i < 16; ++i) {
            inst.transform[i] = static_cast<float>(mat_d.data()[i]);
        }

        const MeshAabb& ma = mesh_aabbs[local_mesh_id];
        worldAabbFromLocal(ma.lmin, ma.lmax, inst.transform,
                           inst.world_aabb_min, inst.world_aabb_max);

        emit instanceReady(std::move(inst));
        total_shapes++;

        int p = iterator->progress();
        if (p != last_progress) {
            last_progress = p;
            progress_ = p;
            emit progressChanged(p);
        }
    } while (iterator->next());

    progress_ = 100;
    emit progressChanged(100);

    double dedup_ratio = total_meshes > 0
        ? static_cast<double>(total_shapes) / static_cast<double>(total_meshes) : 1.0;
    qDebug("Streamer done: %s  %.2fs  shapes=%u  unique_meshes=%u  dedup=%.2fx",
           path.c_str(), stream_timer.elapsed() / 1000.0,
           total_shapes, total_meshes, dedup_ratio);
    succeeded_ = !cancel_requested_.load();
}