#include "mesh.hpp"
#include <fstream>
#include <sstream>
#include <stdexcept>
#include <vector>
#include <map>

std::map<std::string, std::weak_ptr<const gfx::Texture>> assets::Mesh::s_texture_cache;
std::map<std::string, std::weak_ptr<const assets::Skeleton>> assets::Mesh::s_skeleton_cache;

static int GetVertexAttrFlags(int mesh_flags)
{
    int attrs = gfx::VA_POSITION | gfx::VA_NORMAL | gfx::VA_UV;

    if (mesh_flags & assets::MF_HAS_LIGHTMAP_UV)
    {
        attrs |= gfx::VA_LIGHTMAP_UV;
    }
    
    if (mesh_flags & assets::MF_IS_SKELETAL)
    {
        attrs |= gfx::VA_BONE_INDICES | gfx::VA_BONE_WEIGHTS;
	}
    
    return attrs;
}

template <class T>
static void BufferPut(std::vector<char>& buffer, const T& val)
{
    const char* data = reinterpret_cast<const char*>(&val);
	buffer.insert(buffer.end(), data, data + sizeof(T));
}

assets::Mesh::Mesh(std::span<MeshVertex> verts, std::span<MeshTriangle> tris, std::span<MeshMaterial> materials, int flags, std::shared_ptr<const Skeleton> skeleton) :
    va_(GetVertexAttrFlags(flags), gfx::VF_CREATE_EBO),
	skeleton_(std::move(skeleton))
{
	verts_.assign(verts.begin(), verts.end());
	tris_.assign(tris.begin(), tris.end());
	materials_.assign(materials.begin(), materials.end());

    std::vector<char> buffer;

    for (const MeshVertex& vert : verts)
    {
        BufferPut(buffer, vert.pos);
		BufferPut(buffer, vert.normal);
        BufferPut(buffer, vert.uv);

        if (flags & MF_HAS_LIGHTMAP_UV)
        {
            BufferPut(buffer, vert.lightmap_uv);
        }
        
        if (flags & MF_IS_SKELETAL)
        {
            for (int i = 0; i < 4; ++i)
            {
				BufferPut(buffer, static_cast<int32_t>(vert.bones[i].bone_index));
            }

            for (int i = 0; i < 4; ++i)
            {
                BufferPut(buffer, vert.bones[i].weight);
            }
        }
    }

	va_.SetVBOData(buffer.data(), buffer.size());
	va_.SetIndices(reinterpret_cast<const GLuint*>(tris.data()), tris.size() * 3);
}

std::shared_ptr<const assets::Mesh> assets::Mesh::LoadFromFile(const std::string& filename, bool load_collision)
{
    std::shared_ptr<collision::TriangleMesh> collision_mesh;

    if (load_collision)
		collision_mesh = std::make_shared<collision::TriangleMesh>();

    std::ifstream file(filename, std::ios::binary);

    if (!file.is_open())
    {
        throw std::runtime_error("Failed to open mesh file: " + filename);
    }

    int flags = 0;
	std::shared_ptr<const Skeleton> skeleton;

    std::vector<MeshVertex> verts;
    std::vector<MeshTriangle> tris;
	std::vector<MeshMaterial> materials;
    std::vector<int> bone_skel_ids;

    std::string line;

    while (std::getline(file, line))
    {
        if (line.empty() || line[0] == '#') // Skip empty lines and comments
            continue;

        std::istringstream iss(line);

        std::string command;
        iss >> command;

        if (command == "v") // Vertex
        {
            MeshVertex& vertex = verts.emplace_back();
            iss >> vertex.pos.x >> vertex.pos.y >> vertex.pos.z;
            iss >> vertex.normal.x >> vertex.normal.y >> vertex.normal.z;
            iss >> vertex.uv.x >> vertex.uv.y;

            if (flags & MF_HAS_LIGHTMAP_UV)
            {
                iss >> vertex.lightmap_uv.x >> vertex.lightmap_uv.y;
            }

            if (flags & MF_IS_SKELETAL)
            {
                int num_bones;
				iss >> num_bones;

                if (num_bones > 4)
                    num_bones = 4;

                for (int i = 0; i < num_bones; ++i)
                {
					int bone_idx;
					float weight;
                    iss >> bone_idx >> weight;

                    if (bone_idx < 0 || bone_idx >= static_cast<int>(bone_skel_ids.size()))
                    {
                        throw std::runtime_error("Bone index out of bounds in mesh file: " + filename);
					}

                    vertex.bones[i].bone_index = bone_skel_ids[bone_idx];
					vertex.bones[i].weight = weight;
                }

				// Fill remaining influences with zeroes
                for (int i = num_bones; i < 4; ++i)
                {
                    vertex.bones[i].bone_index = -1;
                    vertex.bones[i].weight = 0.0f;
                }
            }
        }
        else if (command == "m") // Material switch
        {
            if (materials.size() > 0)
            {
				MeshMaterial& last_material = materials.back();
				last_material.num_tris = tris.size() - last_material.first_tri;
            }

			MeshMaterial& material = materials.emplace_back();
            iss >> material.name;
			material.first_tri = tris.size();

        }
        else if (command == "f") // Face
        {
			MeshTriangle& tri = tris.emplace_back();
            glm::vec3 tri_verts[3];
            for (size_t i = 0; i < 3; ++i)
            {
				uint32_t vert_index;
                iss >> vert_index;
                if (vert_index >= verts.size())
                {
                    throw std::runtime_error("Vertex index out of bounds in mesh file: " + filename);
                }
                tri.vert[i] = vert_index;
				tri_verts[i] = verts[vert_index].pos;
            }

            if (load_collision)
            {
                // Add triangle to collision mesh
                collision_mesh->AddTriangle(tri_verts[0], tri_verts[1], tri_verts[2]);
			}
        }
        else if (command == "d")
        {
            if (!skeleton)
            {
				throw std::runtime_error("Bone definition found in non-skeletal mesh file: " + filename);
            }

            std::string bone_name;
			iss >> bone_name;
			int idx = skeleton->GetBoneIndex(bone_name);
            
            if (idx < 0)
            {
				throw std::runtime_error("Bone referenced in mesh not found in provided skeleton: " + bone_name);
            }

			bone_skel_ids.push_back(idx);
        }
        else if (command == "luv")
        {
            flags |= MF_HAS_LIGHTMAP_UV;

            if (verts.size() > 0)
            {
                throw std::runtime_error("Lightmap UV flag must be specified before vertex definitions in mesh file: " + filename);
			}
        }
        else if (command == "skeleton")
        {
            if (verts.size() > 0)
            {
                throw std::runtime_error("Skeletal flag must be specified before vertex definitions in mesh file: " + filename);
            }

            flags |= MF_IS_SKELETAL;

			std::string skel_filename;
			iss >> skel_filename;

            // load skeleton
			skeleton = LoadSkeleton("data/" + skel_filename + ".skel");

        }
        else
        {
            throw std::runtime_error("Unknown command in mesh file: " + command);
        }
    }

    if (materials.size() > 0)
    {
        MeshMaterial& last_material = materials.back();
        last_material.num_tris = tris.size() - last_material.first_tri;
	}

    file.close();
    
    if (verts.empty() || tris.empty())
    {
        throw std::runtime_error("Mesh file is empty or malformed: " + filename);
    }

	// Load textures for materials
    for (MeshMaterial& material : materials)
    {
        if (!material.name.empty())
        {
            try {
                material.texture = LoadTexture("data/" + material.name + ".png");
            } catch (const std::exception& e) {
                throw std::runtime_error("Failed to load texture for material '" + material.name + "': " + e.what());
            }
        }
    }
	// Create the mesh object

	std::shared_ptr<Mesh> mesh = std::make_shared<Mesh>(verts, tris, materials, flags, skeleton);

    if (load_collision)
    {
        collision_mesh->Build();
	    mesh->SetCollisionMesh(std::move(collision_mesh));
    }
	return mesh;
}

std::shared_ptr<const gfx::Texture> assets::Mesh::LoadTexture(const std::string& filename)
{
	auto it = s_texture_cache.find(filename);
    if (it != s_texture_cache.end())
    {
        if (auto texture = it->second.lock())
        {
            return texture; // Return cached texture
        }
	}

	auto texture = gfx::Texture::LoadFromFile(filename);
	s_texture_cache[filename] = texture; // Cache the texture
	
    return texture;
}

std::shared_ptr<const assets::Skeleton> assets::Mesh::LoadSkeleton(const std::string& filename)
{
	auto it = s_skeleton_cache.find(filename);
    if (it != s_skeleton_cache.end())
    {
        if (auto skel = it->second.lock())
        {
            return skel; // Return cached skeleton
        }
	}

	auto skel = Skeleton::LoadFromFile(filename);
    s_skeleton_cache[filename] = skel; // Cache the skeleton
    
	return skel;
}