#include "assetload_funcs.hpp"

#include "animation/rigging.hpp"
#include "Assets.hpp"
#include "AssetsLoader.hpp"
#include "audio/audio.hpp"
#include "coders/commons.hpp"
#include "coders/GLSLExtension.hpp"
#include "coders/imageio.hpp"
#include "coders/json.hpp"
#include "coders/obj.hpp"
#include "coders/vcm.hpp"
#include "coders/vec3.hpp"
#include "coders/vector_fonts.hpp"
#include "constants.hpp"
#include "debug/Logger.hpp"
#include "engine/EnginePaths.hpp"
#include "frontend/UiDocument.hpp"
#include "graphics/commons/Model.hpp"
#include "graphics/core/Atlas.hpp"
#include "graphics/core/Font.hpp"
#include "graphics/core/ImageData.hpp"
#include "graphics/core/Shader.hpp"
#include "graphics/core/Texture.hpp"
#include "graphics/core/TextureAnimation.hpp"
#include "logic/scripting/scripting.hpp"
#include "io/io.hpp"
#include "util/stringutil.hpp"

#include <array>
#include <filesystem>
#include <stdexcept>

static debug::Logger logger("assetload-funcs");

namespace fs = std::filesystem;

static bool load_animated_texture(
    Assets& assets,
    const ResPaths& paths,
    const std::string& atlasName,
    const std::string& directory,
    const std::string& name,
    Atlas* dstAtlas
);

assetload::postfunc assetload::animation(
    AssetsLoader& loader,
    const ResPaths& paths,
    const std::string& filename,
    const std::string& name,
    const std::shared_ptr<AssetCfg>&
) {
    auto path = paths.find(filename + ".vca");
    loader.attachToFile(path, {name, AssetType::ANIMATION});

    if (io::is_regular_file(path)) {
        scripting::load_vca_animation(path, nullptr, name);
    }
    return [](auto&) {};
}

assetload::postfunc assetload::texture(
    AssetsLoader& loader,
    const ResPaths& paths,
    const std::string& filename,
    const std::string& name,
    const std::shared_ptr<AssetCfg>&
) {
    auto file = paths.find(filename + ".png");
    loader.attachToFile(file, {name, AssetType::TEXTURE});
    try {
        std::shared_ptr<ImageData> image(imageio::read(file));
        return [name, image, file](auto& assets) {
            assets.store(Texture::from(image.get()), name);
        };
    } catch (const std::runtime_error& err) {
        logger.error() << file.string() << ": " << err.what();
        return [](auto&) {};
    }
}

static auto process_program(const ResPaths& paths, const std::string& filename) {
    io::path vertexFile = paths.find(filename + ".glslv");
    io::path fragmentFile = paths.find(filename + ".glslf");

    std::string vertexSource = io::read_string(vertexFile);
    std::string fragmentSource = io::read_string(fragmentFile);

    auto& preprocessor = *Shader::preprocessor;

    auto vertex = preprocessor.process(vertexFile, vertexSource, false, {});
    auto fragment = preprocessor.process(fragmentFile, fragmentSource, false, {});
    return std::make_pair(vertex, fragment);
}

static auto read_program(const ResPaths& paths, const std::string& filename) {
    io::path vertexFile = paths.find(filename + ".glslv");
    io::path fragmentFile = paths.find(filename + ".glslf");
    return std::make_pair(
        io::read_string(vertexFile), io::read_string(fragmentFile)
    );
}

assetload::postfunc assetload::shader(
    AssetsLoader&,
    const ResPaths& paths,
    const std::string& filename,
    const std::string& name,
    const std::shared_ptr<AssetCfg>&
) {
    auto result = read_program(paths, filename);
    auto vertex = result.first;
    auto fragment = result.second;

    io::path vertexFile = paths.find(filename + ".glslv");
    io::path fragmentFile = paths.find(filename + ".glslf");

    return [=](auto& assets) {
        assets.store(
            Shader::create(
                {vertexFile.string(), vertex},
                {fragmentFile.string(), fragment}
            ),
            name
        );
    };
}

assetload::postfunc assetload::posteffect(
    AssetsLoader& loader,
    const ResPaths& paths,
    const std::string& file,
    const std::string& name,
    const std::shared_ptr<AssetCfg>& settings
) {
    io::path effectFile = paths.find(file + ".glsl");
    loader.attachToFile(effectFile, {name, AssetType::POST_EFFECT});

    std::string effectSource = io::read_string(effectFile);

    auto& preprocessor = *Shader::preprocessor;
    preprocessor.addHeader(
        "__effect__", preprocessor.process(effectFile, effectSource, true, {})
    );

    auto [vertex, fragment] = process_program(paths, SHADERS_FOLDER + "/effect");
    auto params = std::move(fragment.params);

    std::string vertexSource = std::move(vertex.code);
    std::string fragmentSource = std::move(fragment.code);

    return [=](auto& assets) {
        auto program = Shader::create(
            {effectFile.string(), vertexSource},
            {effectFile.string(), fragmentSource}
        );
        bool advanced = false;
        if (settings) {
            advanced = dynamic_cast<const PostEffectCfg*>(settings.get())->advanced;
        }
        assets.store(
            std::make_shared<PostEffect>(advanced, std::move(program), params),
            name
        );
    };
}

static bool append_atlas(AtlasBuilder& atlas, const io::path& file) {
    std::string name = file.stem();
    // skip duplicates
    if (atlas.has(name)) {
        return false;
    }
    auto image = imageio::read(file);
    image->fixAlphaColor();
    atlas.add(name, std::move(image));
    return true;
}

assetload::postfunc assetload::atlas(
    AssetsLoader& loader,
    const ResPaths& paths,
    const std::string& directory,
    const std::string& name,
    const std::shared_ptr<AssetCfg>& config
) {
    auto atlasConfig = std::dynamic_pointer_cast<AtlasCfg>(config);
    if (atlasConfig && atlasConfig->type == AtlasType::SEPARATE) {
        for (const auto& file : paths.listdir(directory)) {
            if (!imageio::is_read_supported(file.extension()))
                continue;
            loader.add(
                AssetType::TEXTURE,
                directory + "/" + file.stem(),
                name + "/" + file.stem()
            );
        }
        return [](auto&){};
    }
    AtlasBuilder builder;
    for (const auto& file : paths.listdir(directory)) {
        if (!imageio::is_read_supported(file.extension())) continue;
        if (!append_atlas(builder, file)) continue;
    }
    std::set<std::string> names = builder.getNames();
    Atlas* atlas = builder.build(ATLAS_EXTRUSION, false).release();
    return [=](Assets& assets) {
        atlas->prepare();
        assets.store(std::unique_ptr<Atlas>(atlas), name);
        for (const auto& file : names) {
            load_animated_texture(assets, paths, name, directory, file, atlas);
        }
    };
}

assetload::postfunc assetload::font(
    AssetsLoader&,
    const ResPaths& paths,
    const std::string& filename,
    const std::string& name,
    const std::shared_ptr<AssetCfg>& config
) {
    auto cfg = std::dynamic_pointer_cast<FontCfg>(config);
    auto ext = fs::path(filename).extension().string();
    if (!ext.empty()) {
        logger.info() << "loading vector font " << util::quote(filename);
        return [=](Assets& assets) {
            using FontFile = vector_fonts::FontFile;
            FontFile* fontFile;
            if ((fontFile = assets.get<FontFile>(filename)) == nullptr) {
                auto fontFilePtr = vector_fonts::load_font(paths.find(filename).string());
                fontFile = fontFilePtr.get();
                assets.store(fontFilePtr, filename);
            }
            assets.store(fontFile->createInstance(cfg ? cfg->size : 16), name);
        };
    }

    auto pages = std::make_shared<std::vector<std::unique_ptr<ImageData>>>();
    for (size_t i = 0; i <= 1024; i++) {
        std::string pagefile = filename + "_" + std::to_string(i) + ".png";
        auto file = paths.find(pagefile);
        if (io::exists(file)) {
            pages->push_back(imageio::read(file));
        } else if (i == 0) {
            throw std::runtime_error("font must have page 0");
        } else {
            pages->push_back(nullptr);
        }
    }
    return [=](auto& assets) {
        assets.store(Font::createBitmapFont(std::move(*pages)), name);
    };
}

assetload::postfunc assetload::layout(
    AssetsLoader&,
    const ResPaths&,
    const std::string& file,
    const std::string& name,
    const std::shared_ptr<AssetCfg>& config
) {
    return [=](auto& assets) {
        try {
            auto cfg = std::dynamic_pointer_cast<LayoutCfg>(config);
            size_t pos = name.find(':');
            auto prefix = name.substr(0, pos);
            assets.store(
                UiDocument::read(
                    *cfg->gui,
                    cfg->env,
                    name,
                    file,
                    prefix + ":layouts/" + name.substr(pos + 1) + ".xml"
                ),
                name
            );
        } catch (const parsing_error& err) {
            throw std::runtime_error(
                "failed to parse layout XML '" + file + "':\n" + err.errorLog()
            );
        }
    };
}
assetload::postfunc assetload::sound(
    AssetsLoader&,
    const ResPaths& paths,
    const std::string& file,
    const std::string& name,
    const std::shared_ptr<AssetCfg>& config
) {
    auto cfg = std::dynamic_pointer_cast<SoundCfg>(config);
    bool keepPCM = cfg ? cfg->keepPCM : false;

    std::unique_ptr<audio::Sound> baseSound = nullptr;
    static std::vector<std::string> extensions {".ogg", ".wav"};
    std::string extension;
    for (size_t i = 0; i < extensions.size(); i++) {
        extension = extensions[i];
        // looking for 'sound_name' as base sound
        auto soundFile = paths.find(file + extension);
        if (io::exists(soundFile)) {
            baseSound = audio::load_sound(soundFile, keepPCM);
            break;
        }
        // looking for 'sound_name_0' as base sound
        auto variantFile = paths.find(file + "_0" + extension);
        if (io::exists(variantFile)) {
            baseSound = audio::load_sound(variantFile, keepPCM);
            break;
        }
    }
    if (baseSound == nullptr) {
        throw std::runtime_error("could not to find sound: " + file);
    }

    // loading sound variants
    for (uint i = 1;; i++) {
        auto variantFile =
            paths.find(file + "_" + std::to_string(i) + extension);
        if (!io::exists(variantFile)) {
            break;
        }
        baseSound->variants.emplace_back(audio::load_sound(variantFile, keepPCM));
    }

    auto sound = baseSound.release();
    return [=](auto& assets) {
        assets.store(std::unique_ptr<audio::Sound>(sound), name);
    };
}

static void request_textures(AssetsLoader& loader, const model::Model& model) {
    for (auto& mesh : model.meshes) {
        if (mesh.texture.find('$') == std::string::npos &&
            mesh.texture.find(':') == std::string::npos) {
            auto filename = TEXTURES_FOLDER + "/" + mesh.texture;
            loader.add(
                AssetType::TEXTURE, filename, mesh.texture, nullptr
            );
        }
    }
}

static assetload::postfunc load_vec3_model(
    const io::path& file,
    AssetsLoader& loader,
    const std::string& name,
    const std::shared_ptr<ModelCfg>& config
) {
    loader.attachToFile(file, {name, AssetType::MODEL});
    auto bytes = io::read_bytes_buffer(file);
    auto modelVEC3 = std::make_shared<vec3::File>(vec3::load(file.string(), bytes));
    return [&loader, name, file, config, modelVEC3 = std::move(modelVEC3)](
                Assets& assets
            ) {
        loader.attachToFile(file, {name, AssetType::MODEL});
        if (config && config->squashed) {
            model::Model fullModel;
            for (auto& entry : modelVEC3->models) {
                auto& vec3model = entry.second;
                auto& model = vec3model.model;
                model.translate(vec3model.origin);
                fullModel.merge(std::move(model));
            }
            request_textures(loader, fullModel);
            assets.store(
                std::make_unique<model::Model>(fullModel),
                name
            );
            logger.info() << "store model " << util::quote(name);
            return;
        }
        for (auto& [modelName, model] : modelVEC3->models) {
            request_textures(loader, model.model);
            std::string fullName = name;
            if (name != modelName) {
                fullName += "." + modelName;
            }
            assets.store(
                std::make_unique<model::Model>(model.model),
                fullName
            );
            logger.info() << "store model " << util::quote(modelName)
                            << " as " << util::quote(fullName);
        }
    };
}

static assetload::postfunc load_obj_model(
    const io::path& file,
    AssetsLoader& loader,
    const std::string& name,
    const std::shared_ptr<ModelCfg>&
) {
    loader.attachToFile(file, {name, AssetType::MODEL});
    auto text = io::read_string(file);
    try {
        auto model = obj::parse(file.string(), text).release();
        return [=, &loader](Assets& assets) {
            request_textures(loader, *model);
            assets.store(std::unique_ptr<model::Model>(model), name);
        };
    } catch (const parsing_error& err) {
        logger.error() << "error on loading model " << name << ": "
                       << err.errorLog();
        throw;
    }
}

static assetload::postfunc load_vcm_model(
    const io::path& file,
    AssetsLoader& loader,
    const std::string& name,
    const std::shared_ptr<ModelCfg>& cfg
) {
    loader.attachToFile(file, {name, AssetType::MODEL});
    auto text = io::read_string(file);
    try {
        auto vcmModel = vcm::parse(file.string(), text, file.extension() == ".xml");

        assert(vcmModel.parts.size() > 0);

        if (vcmModel.parts.size() == 1 || (cfg && cfg->squashed)) {
            auto modelPtr =
                std::make_unique<model::Model>(std::move(vcmModel.squash()))
                    .release();
            return [=, &loader](Assets& assets) {
                auto model = std::unique_ptr<model::Model>(modelPtr);
                request_textures(loader, *model);
                assets.store(std::move(model), name);
                logger.info() << "store model " << util::quote(name);
            };
        }
        auto vcmModelPtr =
            std::make_unique<vcm::VcmModel>(std::move(vcmModel)).release();
        return [=, &loader](Assets& assets) {
            auto vcmModel = std::unique_ptr<vcm::VcmModel>(vcmModelPtr);
            for (auto& [partName, model] : vcmModel->parts) {
                auto fullName = name + "." + partName;
                logger.info()
                    << "store model part " << util::quote(partName)
                    << " as " << util::quote(fullName);
                request_textures(loader, model);
                assets.store(
                    std::make_unique<model::Model>(std::move(model)),
                    fullName
                );
            }
            for (auto& bone : vcmModel->skeleton->getBones()) {
                bone->setModel(name + "." + bone->model.name);
            }
            logger.info() << "store skeleton " << util::quote(name);
            assets.store<rigging::SkeletonConfig>(
                std::make_unique<rigging::SkeletonConfig>(
                    std::move(*vcmModel->skeleton)
                ),
                name
            );
        };
    } catch (const parsing_error& err) {
        logger.error() << "error on loading model " << name << ": "
                       << err.errorLog();
        throw;
    }
}

assetload::postfunc assetload::model(
    AssetsLoader& loader,
    const ResPaths& paths,
    const std::string& file,
    const std::string& name,
    const std::shared_ptr<AssetCfg>& config
) {
    auto modelConfig = std::dynamic_pointer_cast<ModelCfg>(config);

    auto path = paths.find(file + ".vec3");
    if (io::exists(path)) {
        return load_vec3_model(path, loader, name, modelConfig);
    }
    path = paths.find(file + ".obj");
    if (io::exists(path)) {
        return load_obj_model(path, loader, name, modelConfig);
    }

    std::array<std::string, 2> extensions {
        ".xml",
        ".vcm"
    };

    path = "";
    for (const auto& ext : extensions) {
        auto newPath = paths.find(file + ext);
        if (io::exists(newPath)) {
            path = std::move(newPath);
            break;
        }
    }
    if (path.empty()) {
        throw std::runtime_error("could not to find model " + util::quote(file));
    }

    return load_vcm_model(path, loader, name, modelConfig);
}

assetload::postfunc assetload::skeleton(
    AssetsLoader& loader,
    const ResPaths& paths,
    const std::string& file,
    const std::string& name,
    const std::shared_ptr<AssetCfg>& settings
) {
    return [&loader, &paths, file, name](auto& assets) {
        auto path = paths.find(file + ".json");
        std::string text = io::read_string(path);
        auto skeleton = rigging::SkeletonConfig::parse(text, file, name);
        for (auto& bone : skeleton->getBones()) {
            std::string model = bone->model.name;
            size_t pos = model.rfind('.');
            if (pos != std::string::npos) {
                model = model.substr(0, pos);
            }
            if (!model.empty()) {
                loader.add(
                    AssetType::MODEL, MODELS_FOLDER + "/" + model, model
                );
            }
        }
        assets.store(std::move(skeleton), name);
    };
}

static void read_anim_file(
    const std::string& animFile,
    std::vector<std::pair<std::string, int>>& frameList
) {
    auto root = io::read_json(animFile);
    float frameDuration = DEFAULT_FRAME_DURATION;
    std::string frameName;

    if (auto found = root.at("frames")) {
        auto& frameArr = *found;
        for (size_t i = 0; i < frameArr.size(); i++) {
            const auto& currentFrame = frameArr[i];

            frameName = currentFrame[0].asString();
            if (currentFrame.size() > 1) {
                frameDuration = currentFrame[1].asNumber();
            }
            frameList.emplace_back(frameName, frameDuration);
        }
    }
}

static TextureAnimation create_animation(
    Atlas* srcAtlas,
    Atlas* dstAtlas,
    const std::string& name,
    const std::set<std::string>& frameNames,
    const std::vector<std::pair<std::string, int>>& frameList
) {
    Texture* srcTex = srcAtlas->getTexture();
    Texture* dstTex = dstAtlas->getTexture();
    UVRegion region = dstAtlas->get(name);

    TextureAnimation animation(srcTex, dstTex);
    Frame frame;

    uint dstWidth = dstTex->getWidth();
    uint dstHeight = dstTex->getHeight();

    uint srcWidth = srcTex->getWidth();
    uint srcHeight = srcTex->getHeight();

    const int extension = 2;

    frame.dstPos =
        glm::ivec2(region.u1 * dstWidth, region.v1 * dstHeight) - extension;
    frame.size = glm::ivec2(region.u2 * dstWidth, region.v2 * dstHeight) -
                 frame.dstPos + extension;

    for (const auto& elem : frameList) {
        if (!srcAtlas->has(elem.first)) {
            logger.error() << "unknown frame name: " << elem.first;
            continue;
        }
        region = srcAtlas->get(elem.first);
        if (elem.second > 0) {
            frame.duration = static_cast<float>(elem.second) / 1000.0f;
        }
        frame.srcPos = glm::ivec2(
            region.u1 * srcWidth, srcHeight - region.v2 * srcHeight
        ) - extension;
        animation.addFrame(frame);
    }
    return animation;
}

inline bool contains(
    const std::vector<std::pair<std::string, int>>& frameList,
    const std::string& frameName
) {
    for (const auto& elem : frameList) {
        if (frameName == elem.first) {
            return true;
        }
    }
    return false;
}

static bool load_animated_texture(
    Assets& assets,
    const ResPaths& paths,
    const std::string& atlasName,
    const std::string& directory,
    const std::string& name,
    Atlas* dstAtlas
) {
    std::string animsDir = directory + "/animation";

    for (const auto& folder : paths.listdir(animsDir)) {
        if (!io::is_directory(folder)) continue;
        if (folder.name() != name) continue;
        //FIXME: if (fs::is_empty(folder)) continue;

        AtlasBuilder builder;
        append_atlas(builder, paths.find(directory + "/" + name + ".png"));

        std::vector<std::pair<std::string, int>> frameList;
        std::string animFile = folder.string() + "/animation.json";
        if (io::exists(animFile)) {
            read_anim_file(animFile, frameList);
        }
        for (const auto& file : paths.listdir(animsDir + "/" + name)) {
            if (!frameList.empty() && !contains(frameList, file.stem())) {
                continue;
            }
            if (!append_atlas(builder, file)) continue;
        }
        auto srcAtlas = builder.build(ATLAS_EXTRUSION, true);
        if (frameList.empty()) {
            for (const auto& frameName : builder.getNames()) {
                frameList.emplace_back(frameName, 0);
            }
        }
        auto animation = create_animation(
            srcAtlas.get(), dstAtlas, name, builder.getNames(), frameList
        );
        assets.store(
            std::move(srcAtlas), atlasName + "/" + name + "_animation"
        );
        assets.store(animation);
        return true;
    }
    return true;
}