#include <destrum/FS/AssetFS.h>

#include <algorithm>
#include <filesystem>
#include <fstream>
#include <stdexcept>
#include <string>

#include "spdlog/spdlog.h"

namespace {
    struct ParsedVirtualPath {
        std::string scheme;
        std::filesystem::path relativePath;
    };

    [[nodiscard]] std::filesystem::path ValidateRelativePath(std::string_view value) {
        if (value.empty()) {
            throw std::runtime_error("asset path is empty");
        }

        const std::filesystem::path path{std::string(value)};
        if (path.empty() || path.is_absolute() || path.has_root_path()) {
            throw std::runtime_error("asset path must be relative");
        }

        for (const auto& part : path) {
            if (part == "..") {
                throw std::runtime_error("asset path traversal is not allowed");
            }
        }

        const std::filesystem::path normalized = path.lexically_normal();
        if (normalized.empty() || normalized == ".") {
            throw std::runtime_error("asset path is invalid");
        }

        return normalized;
    }

    [[nodiscard]] ParsedVirtualPath ParseVirtualPath(std::string_view value) {
        const std::size_t separator = value.find("://");
        if (separator == std::string_view::npos || separator == 0) {
            throw std::runtime_error("asset path must use scheme://relative/path syntax");
        }

        const std::string scheme(value.substr(0, separator));
        if (scheme.find_first_of("/\\:") != std::string::npos) {
            throw std::runtime_error("asset scheme is invalid");
        }

        return {
            scheme,
            ValidateRelativePath(value.substr(separator + 3))
        };
    }

    [[nodiscard]] std::filesystem::path ResolvePath(const FSMount& mount,
                                                     const std::filesystem::path& relativePath) {
        const std::filesystem::path root = mount.root.lexically_normal();
        const std::filesystem::path full = (root / relativePath).lexically_normal();
        const std::filesystem::path relative = full.lexically_relative(root);

        if (relative.empty() || relative.is_absolute() || relative == ".." ||
            relative.generic_string().starts_with("../")) {
            throw std::runtime_error("asset path escapes its mount");
        }

        std::error_code error;
        const std::filesystem::path canonicalRoot = std::filesystem::weakly_canonical(root, error);
        if (error) {
            throw std::runtime_error("failed to resolve asset mount: " + root.string());
        }

        error.clear();
        const std::filesystem::path canonicalFull = std::filesystem::weakly_canonical(full, error);
        if (error) {
            throw std::runtime_error("failed to resolve asset path: " + full.string());
        }

        const std::filesystem::path canonicalRelative =
            canonicalFull.lexically_relative(canonicalRoot);
        if (canonicalRelative.empty() || canonicalRelative.is_absolute() ||
            canonicalRelative == ".." || canonicalRelative.generic_string().starts_with("../")) {
            throw std::runtime_error("asset path escapes its mount");
        }

        return full;
    }

    [[nodiscard]] const FSMount* FindMount(const std::vector<FSMount>& mounts,
                                           const std::string& scheme) {
        for (const FSMount& mount : mounts) {
            if (mount.scheme == scheme) {
                return &mount;
            }
        }

        return nullptr;
    }

    [[nodiscard]] std::optional<std::filesystem::path> RelativeToMount(const FSMount& mount,
                                                                       const std::filesystem::path& fullPath) {
        const std::filesystem::path relative = fullPath.lexically_normal().lexically_relative(mount.root);
        if (relative.empty() || relative.is_absolute() || relative == "." ||
            relative == ".." || relative.generic_string().starts_with("../")) {
            return std::nullopt;
        }
        return relative;
    }

    struct PackLocation {
        const FSMount* mount;
        std::string key;
    };

    // Everything under a packed mount comes from the pack, even files that also exist on disk
    [[nodiscard]] std::optional<PackLocation> FindInPack(const std::vector<FSMount>& mounts,
                                                         const std::filesystem::path& fullPath) {
        for (const FSMount& mount : mounts) {
            if (!mount.pack) {
                continue;
            }
            if (const auto relative = RelativeToMount(mount, fullPath)) {
                return PackLocation{&mount, mount.scheme + "/" + relative->generic_string()};
            }
        }
        return std::nullopt;
    }

    [[nodiscard]] std::vector<uint8_t> ReadDiskFile(const std::filesystem::path& fullPath) {
        std::ifstream file(fullPath, std::ios::binary);
        if (!file) {
            throw std::runtime_error("failed to open file: " + fullPath.string());
        }

        file.seekg(0, std::ios::end);
        const std::streamsize size = file.tellg();
        if (size < 0) {
            throw std::runtime_error("failed to determine file size: " + fullPath.string());
        }
        file.seekg(0, std::ios::beg);

        std::vector<uint8_t> buffer(static_cast<std::size_t>(size));
        if (size > 0 && !file.read(reinterpret_cast<char*>(buffer.data()), size)) {
            throw std::runtime_error("failed to read file: " + fullPath.string());
        }
        return buffer;
    }

    [[nodiscard]] std::optional<AssetManifest> LoadManifest(const FSMount& mount) {
        const std::filesystem::path path = mount.root / "manifest.json";
        std::vector<uint8_t> bytes;
        if (mount.pack) {
            const PackFile::Entry* entry = mount.pack->Find(mount.scheme + "/manifest.json");
            if (entry == nullptr) {
                return std::nullopt;
            }
            bytes = mount.pack->Read(*entry);
        } else {
            std::error_code error;
            if (!std::filesystem::is_regular_file(path, error)) {
                return std::nullopt;
            }
            bytes = ReadDiskFile(path);
        }

        try {
            return FS::ParseAssetManifest(bytes);
        } catch (const std::exception& exception) {
            throw std::runtime_error("invalid manifest " + path.string() + ": " + exception.what());
        }
    }
}

void AssetFS::Init(std::filesystem::path exeDir) {
    if (initialized) {
        return;
    }

    const std::size_t originalMountCount = mounts.size();
    try {
        const std::filesystem::path packPath = exeDir / PackName;
        std::error_code error;
        if (std::filesystem::is_regular_file(packPath, error)) {
            MountPack(packPath, exeDir / "assets");
        } else {
            Mount("engine", exeDir / "assets/engine");
            Mount("game", exeDir / "assets/game");
        }
        initialized = true;
    } catch (...) {
        mounts.resize(originalMountCount);
        throw;
    }
}

void AssetFS::Reset() {
    mounts.clear();
    initialized = false;
}

void AssetFS::Mount(std::string scheme, std::filesystem::path root) {
    AddMount(std::move(scheme), std::move(root), nullptr);
}

// Each scheme in the pack mounts at root/<scheme>, so full paths look the same as with the unpacked folders
void AssetFS::MountPack(const std::filesystem::path& packPath, const std::filesystem::path& root) {
    auto pack = std::make_shared<const PackFile::Reader>(packPath);

    std::vector<std::string> schemes;
    for (const PackFile::Entry& entry : pack->Entries()) {
        const std::size_t separator = entry.path.find('/');
        if (separator == std::string::npos || separator == 0) {
            throw std::runtime_error("asset pack entry has no scheme: " + entry.path);
        }
        std::string scheme = entry.path.substr(0, separator);
        if (std::ranges::find(schemes, scheme) == schemes.end()) {
            schemes.push_back(std::move(scheme));
        }
    }

    const std::size_t originalMountCount = mounts.size();
    try {
        for (std::string& scheme : schemes) {
            std::filesystem::path schemeRoot = root / scheme;
            AddMount(std::move(scheme), std::move(schemeRoot), pack);
        }
    } catch (...) {
        mounts.resize(originalMountCount);
        throw;
    }
    spdlog::info("Mounted asset pack {} ({} files)", packPath.string(), pack->Entries().size());
}

void AssetFS::AddMount(std::string scheme, std::filesystem::path root,
                       std::shared_ptr<const PackFile::Reader> pack) {
    if (scheme.empty() || scheme.find_first_of("/\\:") != std::string::npos) {
        throw std::runtime_error("asset mount scheme is invalid");
    }
    if (FindMount(mounts, scheme) != nullptr) {
        throw std::runtime_error("asset mount already exists: " + scheme);
    }

    FSMount mount;
    mount.scheme = std::move(scheme);
    mount.root = std::filesystem::absolute(std::move(root)).lexically_normal();
    mount.pack = std::move(pack);
    mount.manifest = LoadManifest(mount);

    spdlog::debug("Mounting assetfs scheme '{}' to root '{}'", mount.scheme, mount.root.string());
    mounts.push_back(std::move(mount));
}

void AssetFS::ReloadManifest(std::string_view scheme) {
    for (FSMount& mount : mounts) {
        if (mount.scheme != scheme) {
            continue;
        }
        mount.manifest.reset();
        mount.manifest = LoadManifest(mount);
        return;
    }
    throw std::runtime_error("asset mount not found: " + std::string{scheme});
}

std::vector<uint8_t> AssetFS::ReadBytes(std::string_view vpath) {
    return ReadFile(GetFullPath(vpath));
}

std::vector<uint8_t> AssetFS::ReadFile(const std::filesystem::path& fullPath) const {
    if (const auto location = FindInPack(mounts, fullPath)) {
        const PackFile::Entry* entry = location->mount->pack->Find(location->key);
        if (entry == nullptr) {
            throw std::runtime_error("failed to open file: " + fullPath.string() + " is not in " +
                                     location->mount->pack->Path().string());
        }
        return location->mount->pack->Read(*entry);
    }
    return ReadDiskFile(fullPath);
}

bool AssetFS::Exists(const std::filesystem::path& fullPath) const {
    if (const auto location = FindInPack(mounts, fullPath)) {
        return location->mount->pack->Find(location->key) != nullptr;
    }
    std::error_code error;
    return std::filesystem::is_regular_file(fullPath, error);
}

std::vector<std::filesystem::path> AssetFS::ListFiles(const FSMount& mount) const {
    std::vector<std::filesystem::path> files;
    if (mount.pack) {
        const std::string prefix = mount.scheme + "/";
        for (const PackFile::Entry& entry : mount.pack->Entries()) {
            if (entry.path.starts_with(prefix)) {
                files.push_back((mount.root / entry.path.substr(prefix.size())).lexically_normal());
            }
        }
        return files;
    }

    std::error_code error;
    for (const auto& entry : std::filesystem::recursive_directory_iterator(mount.root, error)) {
        if (entry.is_regular_file(error)) {
            files.push_back(entry.path());
        }
    }
    return files;
}

bool AssetFS::IsPacked() const {
    return std::ranges::any_of(mounts, [](const FSMount& mount) { return mount.pack != nullptr; });
}

std::filesystem::path AssetFS::GetFullPath(std::string_view vpath) const {
    if (!initialized) {
        throw std::runtime_error("AssetFS is not initialized");
    }

    const ParsedVirtualPath parsed = ParseVirtualPath(vpath);
    const FSMount* mount = FindMount(mounts, parsed.scheme);
    if (mount == nullptr) {
        throw std::runtime_error("asset mount not found: " + parsed.scheme);
    }

    return ResolvePath(*mount, parsed.relativePath);
}

std::filesystem::path AssetFS::GetCookedPathForFile(std::string_view vpath) const {
    if (!initialized) {
        throw std::runtime_error("AssetFS is not initialized");
    }

    const ParsedVirtualPath parsed = ParseVirtualPath(vpath);
    const FSMount* mount = FindMount(mounts, parsed.scheme);
    if (mount == nullptr) {
        throw std::runtime_error("asset mount not found: " + parsed.scheme);
    }

    const std::filesystem::path rawPath = ResolvePath(*mount, parsed.relativePath);
    if (mount->manifest) {
        if (const ManifestAsset* asset = mount->manifest->FindBySrc(parsed.relativePath.generic_string())) {
            if (asset->out) {
                const std::filesystem::path cookedRelative = ValidateRelativePath(*asset->out);
                const std::filesystem::path cookedPath = ResolvePath(*mount, cookedRelative);

                if (mount->pack && Exists(cookedPath)) {
                    return cookedPath;
                }

                std::error_code error;
                if (std::filesystem::is_regular_file(cookedPath, error)) {
                    error.clear();
                    const auto sourceTime = std::filesystem::last_write_time(rawPath, error);
                    error.clear();
                    const auto cookedTime = std::filesystem::last_write_time(cookedPath, error);
                    if (!error && sourceTime > cookedTime) {
                        throw std::runtime_error(
                            "cooked asset output is stale: " + cookedPath.string());
                    }
                    return cookedPath;
                }

                throw std::runtime_error(
                    "cooked asset output is missing: " + cookedPath.string());
            }
        }
    }

    // Assets without a cooked output, such as shader includes, use the source path.
    return rawPath;
}

std::optional<std::string> AssetFS::ToVirtualPath(const std::filesystem::path& fullPath) const {
    for (const FSMount& mount : mounts) {
        if (const auto relative = RelativeToMount(mount, fullPath)) {
            return mount.scheme + "://" + relative->generic_string();
        }
    }
    return std::nullopt;
}

void AssetFS::SetSourceRoot(std::string_view scheme, std::filesystem::path sourceRoot) {
    for (FSMount& mount : mounts) {
        if (mount.scheme == scheme) {
            mount.sourceRoot = std::filesystem::absolute(std::move(sourceRoot)).lexically_normal();
            return;
        }
    }
    throw std::runtime_error("asset mount not found: " + std::string{scheme});
}

std::optional<std::filesystem::path> AssetFS::GetSourcePath(std::string_view vpath) const {
    const ParsedVirtualPath parsed = ParseVirtualPath(vpath);
    const FSMount* mount = FindMount(mounts, parsed.scheme);
    if (mount == nullptr || !mount->sourceRoot) {
        return std::nullopt;
    }
    return (*mount->sourceRoot / parsed.relativePath).lexically_normal();
}

bool AssetFS::WriteFile(std::string_view vpath, std::string_view contents) {
    std::vector<std::filesystem::path> targets{GetFullPath(vpath)};
    if (FindInPack(mounts, targets.front())) {
        spdlog::error("Can't write {}, it's in a read only asset pack", vpath);
        return false;
    }
    if (auto source = GetSourcePath(vpath)) {
        targets.push_back(std::move(*source));
    }

    for (const auto& target : targets) {
        std::error_code error;
        std::filesystem::create_directories(target.parent_path(), error);
        std::filesystem::path temporary = target;
        temporary += ".tmp";
        {
            std::ofstream file(temporary, std::ios::binary | std::ios::trunc);
            if (!file.write(contents.data(), static_cast<std::streamsize>(contents.size()))) {
                spdlog::error("Failed to write {}", temporary.string());
                return false;
            }
        }
        std::filesystem::rename(temporary, target, error);
        if (error) {
            spdlog::error("Failed to replace {}: {}", target.string(), error.message());
            std::filesystem::remove(temporary, error);
            return false;
        }
    }
    return true;
}