// mde_extract — extract media files from MDE1/MDE2 export files
//
// Usage:
//   mde_extract <file.mdb> <output_dir>              — extract all media
//   mde_extract <file.mdb> <output_dir> --id <id>    — extract single entry
//   mde_extract <file.mdb> <output_dir> --store <id> — extract one store
//   mde_extract <file.mdb> <output_dir> --album <id> — extract one album
//   mde_extract <file.mdb> <output_dir> --list        — list entries only

#include <cstdint>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <string>
#include <vector>

namespace fs = std::filesystem;

// ---- format constants ----

static constexpr char EXPORT_MAGIC_V1[4] = {'M','D','E','1'};
static constexpr char EXPORT_MAGIC_V2[4] = {'M','D','E','2'};

// ---- binary read helpers ----

static bool read_bytes(std::istream& in, void* dest, std::size_t n) {
    in.read(reinterpret_cast<char*>(dest), static_cast<std::streamsize>(n));
    return static_cast<std::size_t>(in.gcount()) == n;
}

static bool read_u32(std::istream& in, std::uint32_t& out) {
    return read_bytes(in, &out, 4);
}

static bool read_u64(std::istream& in, std::uint64_t& out) {
    return read_bytes(in, &out, 8);
}

static bool read_str(std::istream& in, std::string& out) {
    std::uint32_t len;
    if (!read_u32(in, len)) return false;
    if (len == 0) { out.clear(); return true; }
    out.resize(len);
    return read_bytes(in, out.data(), len);
}

// ---- sanitise filename for filesystem ----

static std::string safe_filename(const std::string& name) {
    std::string out;
    out.reserve(name.size());
    for (char c : name) {
        if (c == '/' || c == '\\' || c == '\0' || c == ':' || c == '*'
            || c == '?' || c == '"' || c == '<' || c == '>' || c == '|')
            out += '_';
        else
            out += c;
    }
    if (out.empty()) out = "unnamed";
    return out;
}

static void usage() {
    std::cerr << "Usage: mde_extract <file.mdb> <output_dir> [options]\n"
              << "  --list           List entries without extracting\n"
              << "  --id <media_id>  Extract only this media entry\n"
              << "  --store <id>     Extract only media from this store\n"
              << "  --album <id>     Extract only media from this album\n";
}

int main(int argc, char* argv[]) {
    if (argc < 3) { usage(); return 1; }

    const std::string input_path = argv[1];
    const std::string output_dir = argv[2];

    bool list_only = false;
    std::string filter_media_id;
    std::string filter_store_id;
    std::string filter_album_id;

    for (int i = 3; i < argc; ++i) {
        std::string arg = argv[i];
        if (arg == "--list") { list_only = true; }
        else if (arg == "--id" && i + 1 < argc) { filter_media_id = argv[++i]; }
        else if (arg == "--store" && i + 1 < argc) { filter_store_id = argv[++i]; }
        else if (arg == "--album" && i + 1 < argc) { filter_album_id = argv[++i]; }
        else { std::cerr << "Unknown option: " << arg << "\n"; usage(); return 1; }
    }

    if (!fs::exists(input_path)) {
        std::cerr << "File not found: " << input_path << "\n";
        return 1;
    }

    std::ifstream in(input_path, std::ios::binary);
    if (!in.is_open()) {
        std::cerr << "Cannot open file: " << input_path << "\n";
        return 1;
    }

    // ---- magic ----
    char magic[4]{};
    if (!read_bytes(in, magic, 4)) {
        std::cerr << "ERROR: file too small\n";
        return 1;
    }

    bool is_v2 = (std::memcmp(magic, EXPORT_MAGIC_V2, 4) == 0);
    bool is_v1 = (std::memcmp(magic, EXPORT_MAGIC_V1, 4) == 0);
    if (!is_v1 && !is_v2) {
        std::cerr << "ERROR: not an MDE file (bad magic)\n";
        return 1;
    }

    if (!list_only) {
        std::error_code ec;
        fs::create_directories(output_dir, ec);
        if (ec) {
            std::cerr << "Cannot create output dir: " << ec.message() << "\n";
            return 1;
        }
    }

    std::uint32_t num_stores;
    if (!read_u32(in, num_stores)) {
        std::cerr << "ERROR: truncated store count\n";
        return 1;
    }

    std::size_t extracted = 0;
    std::size_t skipped = 0;

    for (std::uint32_t si = 0; si < num_stores && in.good(); ++si) {
        std::string store_id, store_name, store_created;
        if (!read_str(in, store_id) || !read_str(in, store_name) ||
            !read_str(in, store_created)) {
            std::cerr << "ERROR: truncated store header\n";
            return 1;
        }

        bool store_match = filter_store_id.empty() || filter_store_id == store_id;

        std::uint32_t num_albums;
        if (!read_u32(in, num_albums)) {
            std::cerr << "ERROR: truncated album count\n";
            return 1;
        }

        for (std::uint32_t ai = 0; ai < num_albums && in.good(); ++ai) {
            std::string album_id, album_name, album_created;
            if (!read_str(in, album_id) || !read_str(in, album_name) ||
                !read_str(in, album_created)) {
                std::cerr << "ERROR: truncated album header\n";
                return 1;
            }

            if (is_v2) {
                std::uint8_t pub;
                if (!read_bytes(in, &pub, 1)) {
                    std::cerr << "ERROR: truncated is_public\n";
                    return 1;
                }
            }

            bool album_match = filter_album_id.empty() || filter_album_id == album_id;

            std::uint32_t num_media;
            if (!read_u32(in, num_media)) {
                std::cerr << "ERROR: truncated media count\n";
                return 1;
            }

            for (std::uint32_t mi = 0; mi < num_media && in.good(); ++mi) {
                std::string id, filename, kind, content_type, created;
                if (!read_str(in, id) || !read_str(in, filename) ||
                    !read_str(in, kind) || !read_str(in, content_type) ||
                    !read_str(in, created)) {
                    std::cerr << "ERROR: truncated media entry\n";
                    return 1;
                }

                std::uint64_t size_bytes;
                if (!read_u64(in, size_bytes)) {
                    std::cerr << "ERROR: truncated size_bytes\n";
                    return 1;
                }

                bool media_match = filter_media_id.empty() || filter_media_id == id;
                bool want = store_match && album_match && media_match;

                if (list_only) {
                    if (want) {
                        std::cout << "store=" << store_id
                                  << " album=" << album_id
                                  << " id=" << id
                                  << " file=\"" << filename << "\""
                                  << " kind=" << kind
                                  << " type=" << content_type
                                  << " size=" << size_bytes << "\n";
                    }
                    // Skip data
                    if (size_bytes > 0)
                        in.seekg(static_cast<std::streamoff>(size_bytes), std::ios::cur);
                    continue;
                }

                if (!want || size_bytes == 0) {
                    // Skip data
                    if (size_bytes > 0)
                        in.seekg(static_cast<std::streamoff>(size_bytes), std::ios::cur);
                    if (want && size_bytes == 0)
                        std::cerr << "  SKIP (zero size): " << id << " " << filename << "\n";
                    ++skipped;
                    continue;
                }

                // Build output path: output_dir/store_name/album_name/filename
                fs::path out_dir = fs::path(output_dir)
                    / safe_filename(store_name)
                    / safe_filename(album_name);
                std::error_code ec;
                fs::create_directories(out_dir, ec);

                // Use original filename, append id if collision
                std::string out_name = safe_filename(filename);
                fs::path out_path = out_dir / out_name;
                if (fs::exists(out_path)) {
                    // Deduplicate: prepend media id
                    out_path = out_dir / (id + "_" + out_name);
                }

                std::ofstream out(out_path, std::ios::binary | std::ios::trunc);
                if (!out.is_open()) {
                    std::cerr << "ERROR: cannot create " << out_path << "\n";
                    in.seekg(static_cast<std::streamoff>(size_bytes), std::ios::cur);
                    ++skipped;
                    continue;
                }

                // Stream copy in 1MB chunks
                constexpr std::size_t BUF_SIZE = 1024 * 1024;
                std::vector<char> buf(BUF_SIZE);
                std::uint64_t remaining = size_bytes;
                bool ok = true;
                while (remaining > 0) {
                    auto chunk = static_cast<std::streamsize>(
                        std::min<std::uint64_t>(remaining, BUF_SIZE));
                    in.read(buf.data(), chunk);
                    auto got = in.gcount();
                    if (got <= 0) {
                        std::cerr << "ERROR: truncated data for " << id << "\n";
                        ok = false;
                        break;
                    }
                    out.write(buf.data(), got);
                    remaining -= static_cast<std::uint64_t>(got);
                }
                out.close();

                if (ok) {
                    std::cout << out_path.string() << "  (" << size_bytes << " bytes)\n";
                    ++extracted;
                } else {
                    ++skipped;
                }
            }
        }
    }

    if (!list_only) {
        std::cout << "\nExtracted: " << extracted << " files\n";
        if (skipped > 0)
            std::cout << "Skipped:   " << skipped << "\n";
    }

    return 0;
}
