#include "ModelLoader.h"
#include "MorphManager.h"

#include <iostream>
#include <filesystem>

#include <assimp/Importer.hpp>
#include <assimp/scene.h>
#include <assimp/postprocess.h>

#include <osg/Geode>
#include <osg/Geometry>
#include <osg/Array>
#include <osg/PrimitiveSet>
#include <osg/Material>
#include <osg/Texture2D>
#include <osg/BlendFunc>
#include <osg/StateSet>
#include <osg/MatrixTransform>
#include <osgDB/ReadFile>

namespace fs = std::filesystem;

// ── Helpers ───────────────────────────────────────────────────────────────────

namespace {

// Section header morphs like "------EYE------" are not real morphs
bool isSectionHeader(const std::string& name) {
    return name.size() >= 4 &&
           name.front() == '-' && name.back() == '-';
}

osg::ref_ptr<osg::Texture2D> loadTexture(const std::string& path) {
    osg::ref_ptr<osg::Image> img = osgDB::readImageFile(path);
    if (!img) {
        std::cerr << "[loader] texture not found: " << path << "\n";
        return {};
    }
    auto tex = new osg::Texture2D(img);
    tex->setWrap(osg::Texture::WRAP_S, osg::Texture::REPEAT);
    tex->setWrap(osg::Texture::WRAP_T, osg::Texture::REPEAT);
    tex->setFilter(osg::Texture::MIN_FILTER, osg::Texture::LINEAR_MIPMAP_LINEAR);
    tex->setFilter(osg::Texture::MAG_FILTER, osg::Texture::LINEAR);
    return tex;
}

osg::ref_ptr<osg::Geode> convertMesh(const aiMesh*      mesh,
                                      const aiScene*     scene,
                                      const std::string& baseDir,
                                      MorphManager*      morphMgr) {
    auto geode = new osg::Geode;
    auto geom  = new osg::Geometry;

    // ── Base vertices ─────────────────────────────────────────────────────────
    auto baseVerts = new osg::Vec3Array;
    baseVerts->reserve(mesh->mNumVertices);
    for (unsigned i = 0; i < mesh->mNumVertices; ++i)
        baseVerts->push_back({mesh->mVertices[i].x,
                              mesh->mVertices[i].y,
                              mesh->mVertices[i].z});

    // We set a COPY as the live array; the base is kept separately for morphing
    auto liveVerts = new osg::Vec3Array(*baseVerts);
    geom->setVertexArray(liveVerts);

    // ── Base normals ──────────────────────────────────────────────────────────
    osg::ref_ptr<osg::Vec3Array> baseNormals;
    if (mesh->HasNormals()) {
        baseNormals = new osg::Vec3Array;
        baseNormals->reserve(mesh->mNumVertices);
        for (unsigned i = 0; i < mesh->mNumVertices; ++i)
            baseNormals->push_back({mesh->mNormals[i].x,
                                    mesh->mNormals[i].y,
                                    mesh->mNormals[i].z});
        auto liveNormals = new osg::Vec3Array(*baseNormals);
        geom->setNormalArray(liveNormals, osg::Array::BIND_PER_VERTEX);
    }

    // ── UVs ───────────────────────────────────────────────────────────────────
    if (mesh->HasTextureCoords(0)) {
        auto uvs = new osg::Vec2Array;
        uvs->reserve(mesh->mNumVertices);
        for (unsigned i = 0; i < mesh->mNumVertices; ++i)
            uvs->push_back({mesh->mTextureCoords[0][i].x,
                            mesh->mTextureCoords[0][i].y});
        geom->setTexCoordArray(0, uvs, osg::Array::BIND_PER_VERTEX);
    }

    // ── Vertex colours ────────────────────────────────────────────────────────
    if (mesh->HasVertexColors(0)) {
        auto cols = new osg::Vec4Array;
        cols->reserve(mesh->mNumVertices);
        for (unsigned i = 0; i < mesh->mNumVertices; ++i)
            cols->push_back({mesh->mColors[0][i].r, mesh->mColors[0][i].g,
                             mesh->mColors[0][i].b, mesh->mColors[0][i].a});
        geom->setColorArray(cols, osg::Array::BIND_PER_VERTEX);
    }

    // ── Indices ───────────────────────────────────────────────────────────────
    auto indices = new osg::DrawElementsUInt(osg::PrimitiveSet::TRIANGLES);
    indices->reserve(mesh->mNumFaces * 3);
    for (unsigned f = 0; f < mesh->mNumFaces; ++f) {
        const aiFace& face = mesh->mFaces[f];
        if (face.mNumIndices != 3) continue;
        indices->push_back(face.mIndices[0]);
        indices->push_back(face.mIndices[1]);
        indices->push_back(face.mIndices[2]);
    }
    // Guard: skip meshes with no valid triangles (avoids front() crash on empty vector)
    if (indices->empty()) {
        std::cerr << "[loader] Skipping mesh with no valid triangles: "
                  << mesh->mName.C_Str() << "\n";
        return geode;
    }
    geom->addPrimitiveSet(indices);

    // ── Material ──────────────────────────────────────────────────────────────
    osg::StateSet* ss = geom->getOrCreateStateSet();
    if (mesh->mMaterialIndex < scene->mNumMaterials) {
        const aiMaterial* mat = scene->mMaterials[mesh->mMaterialIndex];
        auto osgMat = new osg::Material;
        aiColor4D colour;

        if (AI_SUCCESS == mat->Get(AI_MATKEY_COLOR_DIFFUSE, colour))
            osgMat->setDiffuse(osg::Material::FRONT_AND_BACK,
                               {colour.r, colour.g, colour.b, colour.a});
        if (AI_SUCCESS == mat->Get(AI_MATKEY_COLOR_AMBIENT, colour))
            osgMat->setAmbient(osg::Material::FRONT_AND_BACK,
                               {colour.r, colour.g, colour.b, colour.a});
        if (AI_SUCCESS == mat->Get(AI_MATKEY_COLOR_SPECULAR, colour))
            osgMat->setSpecular(osg::Material::FRONT_AND_BACK,
                                {colour.r, colour.g, colour.b, colour.a});
        float shininess = 0.f;
        if (AI_SUCCESS == mat->Get(AI_MATKEY_SHININESS, shininess))
            osgMat->setShininess(osg::Material::FRONT_AND_BACK,
                                 std::min(shininess, 128.f));
        ss->setAttribute(osgMat, osg::StateAttribute::ON);

        if (mat->GetTextureCount(aiTextureType_DIFFUSE) > 0) {
            aiString texPath;
            mat->GetTexture(aiTextureType_DIFFUSE, 0, &texPath);
            std::string fullPath = baseDir + "/" + texPath.C_Str();
            for (char& c : fullPath) if (c == '\\') c = '/';
            // Try the path as-is first, then common extension variants
            // (FBX often references .png when textures are actually .jpg or .tga)
            osg::ref_ptr<osg::Texture2D> tex = loadTexture(fullPath);
            if (!tex) {
                // Try swapping extension
                auto swapExt = [](const std::string& p, const std::string& newExt) {
                    auto dot = p.rfind('.');
                    return dot != std::string::npos ? p.substr(0, dot) + newExt : p;
                };
                for (auto& ext : {".jpg", ".jpeg", ".png", ".tga", ".bmp"}) {
                    tex = loadTexture(swapExt(fullPath, ext));
                    if (tex) break;
                }
            }
            if (tex) ss->setTextureAttributeAndModes(0, tex, osg::StateAttribute::ON);
        }

        float opacity = 1.f;
        mat->Get(AI_MATKEY_OPACITY, opacity);
        if (opacity < 1.f) {
            ss->setMode(GL_BLEND, osg::StateAttribute::ON);
            ss->setRenderingHint(osg::StateSet::TRANSPARENT_BIN);
            ss->setAttribute(new osg::BlendFunc(
                osg::BlendFunc::SRC_ALPHA, osg::BlendFunc::ONE_MINUS_SRC_ALPHA));
        }
    }

    geode->addDrawable(geom);

    // ── Morph targets ─────────────────────────────────────────────────────────
    if (morphMgr && mesh->mNumAnimMeshes > 0) {
        // Mark geometry and arrays as DYNAMIC — tells OSG this data changes
        // every frame and must not be double-buffered or cached in display lists.
        geom->setDataVariance(osg::Object::DYNAMIC);
        geom->setUseDisplayList(false);
        geom->setUseVertexBufferObjects(true);

        auto* vArr = dynamic_cast<osg::Vec3Array*>(geom->getVertexArray());
        auto* nArr = dynamic_cast<osg::Vec3Array*>(geom->getNormalArray());
        if (vArr) {
            vArr->setDataVariance(osg::Object::DYNAMIC);
            vArr->setBinding(osg::Array::BIND_PER_VERTEX);
        }
        if (nArr) {
            nArr->setDataVariance(osg::Object::DYNAMIC);
            nArr->setBinding(osg::Array::BIND_PER_VERTEX);
        }

        morphMgr->registerMesh(geom, baseVerts, baseNormals);

        int registered = 0;
        for (unsigned a = 0; a < mesh->mNumAnimMeshes; ++a) {
            const aiAnimMesh* am = mesh->mAnimMeshes[a];
            std::string name = am->mName.C_Str();

            if (name.empty() || isSectionHeader(name)) continue;
            if (am->mNumVertices != mesh->mNumVertices)  continue;

            // Compute vertex deltas (animMesh stores ABSOLUTE positions)
            auto deltaVerts = new osg::Vec3Array;
            deltaVerts->resize(mesh->mNumVertices);
            for (unsigned i = 0; i < mesh->mNumVertices; ++i) {
                (*deltaVerts)[i] = osg::Vec3(
                    am->mVertices[i].x - mesh->mVertices[i].x,
                    am->mVertices[i].y - mesh->mVertices[i].y,
                    am->mVertices[i].z - mesh->mVertices[i].z);
            }

            // Normal deltas (optional)
            osg::ref_ptr<osg::Vec3Array> deltaNormals;
            if (am->mNormals && mesh->HasNormals()) {
                deltaNormals = new osg::Vec3Array;
                deltaNormals->resize(mesh->mNumVertices);
                for (unsigned i = 0; i < mesh->mNumVertices; ++i) {
                    (*deltaNormals)[i] = osg::Vec3(
                        am->mNormals[i].x - mesh->mNormals[i].x,
                        am->mNormals[i].y - mesh->mNormals[i].y,
                        am->mNormals[i].z - mesh->mNormals[i].z);
                }
            }

            morphMgr->addTarget(geom, name, deltaVerts, deltaNormals);
            ++registered;
        }

        if (registered > 0)
            std::cout << "[loader] Mesh \"" << mesh->mName.C_Str()
                      << "\": registered " << registered << " morph targets\n";
    }

    return geode;
}

osg::ref_ptr<osg::Group> convertNode(const aiNode*      node,
                                      const aiScene*     scene,
                                      const std::string& baseDir,
                                      MorphManager*      morphMgr) {
    const aiMatrix4x4& m = node->mTransformation;
    osg::Matrixf mat(m.a1, m.b1, m.c1, m.d1,
                     m.a2, m.b2, m.c2, m.d2,
                     m.a3, m.b3, m.c3, m.d3,
                     m.a4, m.b4, m.c4, m.d4);

    auto xform = new osg::MatrixTransform(mat);
    xform->setName(node->mName.C_Str());

    for (unsigned i = 0; i < node->mNumMeshes; ++i)
        xform->addChild(convertMesh(scene->mMeshes[node->mMeshes[i]],
                                    scene, baseDir, morphMgr));

    for (unsigned i = 0; i < node->mNumChildren; ++i)
        xform->addChild(convertNode(node->mChildren[i], scene, baseDir, morphMgr));

    return xform;
}

} // namespace

// ── ModelLoader ───────────────────────────────────────────────────────────────

osg::ref_ptr<osg::Node> ModelLoader::load(const std::string& filepath,
                                           MorphManager*      morphMgr) {
    Assimp::Importer importer;

    // NOTE: JoinIdenticalVertices is intentionally omitted —
    // it destroys the vertex correspondence that morph targets rely on.
    //
    // aiProcess_FlipUVs is omitted — FBX exports from Blender already have
    // UVs in the correct OpenGL orientation (Y=0 at bottom). Flipping them
    // was causing textures to appear mirrored/upside-down.
    // If loading a raw PMX via Assimp ever becomes needed, this flag would
    // need to be added back conditionally based on file extension.
    constexpr unsigned flags =
        aiProcess_Triangulate        |
        aiProcess_GenSmoothNormals   |
        aiProcess_SortByPType        |
        aiProcess_ImproveCacheLocality;

    const aiScene* scene = importer.ReadFile(filepath, flags);
    if (!scene || (scene->mFlags & AI_SCENE_FLAGS_INCOMPLETE) || !scene->mRootNode) {
        std::cerr << "[loader] Assimp error: " << importer.GetErrorString() << "\n";
        return {};
    }

    std::cout << "[loader] Meshes: "    << scene->mNumMeshes
              << "  Materials: "        << scene->mNumMaterials
              << "  Animations: "       << scene->mNumAnimations << "\n";

    const std::string baseDir = fs::path(filepath).parent_path().string();
    return buildOsgScene(scene, baseDir, morphMgr);
}

osg::ref_ptr<osg::Node> ModelLoader::buildOsgScene(const aiScene*     scene,
                                                    const std::string& baseDir,
                                                    MorphManager*      morphMgr) {
    return convertNode(scene->mRootNode, scene, baseDir, morphMgr);
}