#include <filesystem>
#include <fstream>
#include <string>
#include <vector>
#include <cctype>
#include <chrono>
#include <stdexcept>

#include <nlohmann/json.hpp>

#include <spdlog/spdlog.h>
#include <spdlog/sinks/stdout_color_sinks.h>

#include "ShaderCompiler.h"

namespace fs = std::filesystem;
using json = nlohmann::json;

static void ensure_parent_dir(const fs::path& p) {
    fs::create_directories(p.parent_path());
}

static std::string to_lower(std::string s) {
    for (char& c : s) c = (char)std::tolower((unsigned char)c);
    return s;
}

struct Args {
    fs::path input;
    fs::path output;
    bool clean = false;
};

static Args parse_args(int argc, char** argv) {
    Args a;
    for (int i = 1; i < argc; ++i) {
        std::string s = argv[i];
        auto next = [&]() -> std::string {
            if (i + 1 >= argc) throw std::runtime_error("Missing value after " + s);
            return argv[++i];
        };

        if (s == "--input") a.input = next();
        else if (s == "--output") a.output = next();
        else if (s == "--clean") a.clean = true;
        else if (s == "--help" || s == "-h") {
            spdlog::info("AssetCooker --input <dir> --output <dir> [--clean]");
            std::exit(0);
        } else {
            throw std::runtime_error("Unknown arg: " + s);
        }
    }

    if (a.input.empty() || a.output.empty())
        throw std::runtime_error("Usage: AssetCooker --input <dir> --output <dir> [--clean]");

    return a;
}

static std::string classify_type(const fs::path& src) {
    const std::string ext = to_lower(src.extension().string());
    if (ext == ".png" || ext == ".jpg" || ext == ".jpeg" || ext == ".tga") return "texture";
    // Treat these as "shader family" (entry: .vert/.frag/.comp, include-only: .glsl)
    if (ext == ".vert" || ext == ".frag" || ext == ".comp" || ext == ".glsl") return "shader";
    if (ext == ".wav" || ext == ".ogg" || ext == ".mp3") return "audio";
    if (ext == ".gltf" || ext == ".glb" || ext == ".fbx" || ext == ".obj") return "mesh";
    return "raw";
}

static bool needs_rebuild(const fs::path& src, const fs::path& dst) {
    if (!fs::exists(dst)) return true;
    return fs::last_write_time(src) > fs::last_write_time(dst);
}

static bool copy_if_needed(const fs::path& src, const fs::path& dst) {
    if (!needs_rebuild(src, dst)) return false;
    ensure_parent_dir(dst);
    fs::copy_file(src, dst, fs::copy_options::overwrite_existing);
    return true;
}

// Choose your SPV output naming scheme.
// Option A: shader.vert -> shader.vert.spv (keeps original extension)
// Option B: shader.vert -> shader.spv (replace with .spv)
static fs::path make_spv_output_path(const fs::path& dstLikeSrc) {
    fs::path out = dstLikeSrc;
    // Option A:
    out.replace_extension(dstLikeSrc.extension().string() + ".spv");
    // Option B (uncomment to use):
    // out.replace_extension(".spv");
    return out;
}

static std::int64_t filetime_to_epoch_ns(const fs::file_time_type& ft) {
    return (std::int64_t)std::chrono::duration_cast<std::chrono::nanoseconds>(
        ft.time_since_epoch()
    ).count();
}

int main(int argc, char** argv) {
    auto log = spdlog::stdout_color_mt("AssetCooker");
    spdlog::set_default_logger(log);

    spdlog::set_pattern("[%H:%M:%S] [%^%l%$] %v");
    // spdlog::set_level(spdlog::level::debug);

    try {
        const Args args = parse_args(argc, argv);

        if (!fs::exists(args.input) || !fs::is_directory(args.input)) {
            spdlog::error("Input dir does not exist: {}", args.input.string());
            return 1;
        }

        if (args.clean && fs::exists(args.output)) {
            spdlog::info("Cleaning output dir: {}", args.output.string());
            fs::remove_all(args.output);
            return 0;
        }
        fs::create_directories(args.output);

        // ----------------------------
        // Shader compiler setup
        // ----------------------------
        SpirvShaderCompiler::Options scOpt;
        scOpt.debugInfo = false; // set true if you want debug info in SPIR-V
        scOpt.defines = { "DEBUG=1" };


        // Common conventions (adjust to your folder layout):
        // - shaders live under <input>/shaders
        // - includes live under <input>/shaders/include
        //
        // If you don't have these folders, either remove these or point them elsewhere.
        scOpt.includeDirs = {
            args.input / "shaders",
            args.input / "shaders" / "include"
        };

        // Optional defines:
        // scOpt.defines = { "MY_DEFINE=1", "USE_FOG" };

        SpirvShaderCompiler shaderCompiler(scOpt);

        // ----------------------------
        // Manifest init
        // ----------------------------
        json manifest;
        manifest["version"] = 1;
        manifest["assets"] = json::array();

        std::size_t cooked = 0, skipped = 0;

        for (auto it = fs::recursive_directory_iterator(args.input);
             it != fs::recursive_directory_iterator();
             ++it) {

            if (!it->is_regular_file()) continue;

            fs::path src = it->path();
            fs::path rel = fs::relative(src, args.input);
            fs::path dst = args.output / rel;

            const std::string type = classify_type(src);

            bool changed = false;
            fs::path finalOut = dst;      // what actually ends up in output (if any)
            std::string finalType = type; // stored in manifest

            if (type == "shader") {
                // Include-only shader (.glsl): typically NOT cooked to output.
                if (SpirvShaderCompiler::IsIncludeOnlyShader(src)) {
                    changed = false;
                    finalType = "shader_include";
                    // If you *do* want to copy includes, uncomment:
                    // changed = copy_if_needed(src, dst);
                    // finalOut = dst;
                }
                // Entry shader (.vert/.frag): compile to SPIR-V.
                else if (SpirvShaderCompiler::IsEntryShader(src)) {
                    fs::path spvOut = make_spv_output_path(dst);
                    changed = shaderCompiler.CompileFileIfNeeded(src, spvOut);

                    finalOut = spvOut;
                    finalType = "shader_spirv";
                }
                // Something else classified as shader (ex: .comp in your classifier)
                // but not supported by the compiler yet -> copy raw as fallback.
                else {
                    changed = copy_if_needed(src, dst);
                    finalOut = dst;
                    finalType = "shader_raw";
                }
            } else {
                // Normal assets: copy if needed.
                changed = copy_if_needed(src, dst);
                finalOut = dst;
            }

            if (changed) {
                cooked++;
                spdlog::debug("Cooked: {} -> {}",
                              rel.generic_string(),
                              fs::relative(finalOut, args.output).generic_string());
            } else {
                skipped++;
                spdlog::debug("Skipped (up-to-date): {}", rel.generic_string());
            }

            // ----------------------------
            // Manifest entry
            // ----------------------------
            json entry;
            entry["src"] = rel.generic_string(); // relative to input
            entry["type"] = finalType;

            // If this asset produced an output file, record it.
            // If not (e.g. shader includes), set out to null.
            if (fs::exists(finalOut) && fs::is_regular_file(finalOut)) {
                entry["out"] = fs::relative(finalOut, args.output).generic_string();
                entry["size_bytes"] = (std::uintmax_t)fs::file_size(finalOut);
                entry["mtime_epoch_ns"] = filetime_to_epoch_ns(fs::last_write_time(finalOut));
            } else {
                entry["out"] = nullptr;
                entry["size_bytes"] = (std::uintmax_t)fs::file_size(src);
                entry["mtime_epoch_ns"] = filetime_to_epoch_ns(fs::last_write_time(src));
            }

            manifest["assets"].push_back(std::move(entry));
        }

        // Write manifest.json
        {
            fs::path manifestPath = args.output / "manifest.json";
            std::ofstream out(manifestPath);
            if (!out) {
                throw std::runtime_error("Failed to write: " + manifestPath.string());
            }
            out << manifest.dump(2) << "\n";
        }

        spdlog::info("Cook done. cooked={} skipped={} output={}",
                     cooked, skipped, args.output.string());
        return 0;

    } catch (const std::exception& e) {
        spdlog::critical("Cook failed: {}", e.what());
        return 1;
    }
}