import fs from "fs";
import path from "path";

console.log("Loading source files...");

// 1. Load the sources
const imagesDb = JSON.parse(
    fs.readFileSync(path.resolve("./data-hf-images.json"), "utf-8")
);

const uploadedDb = JSON.parse(
    fs.readFileSync(path.resolve("./data-hf-uploaded.json"), "utf-8")
);

// 2. Setup mappings
const imgTypeMap = {
    flux: "Flux",
    wan: "WAN",
    zimage: "ZImage",
    sdxl: "SDXL",
    qwen: "Qwen",
};

const complexKeys = ["flux", "wan", "zimage", "sdxl", "qwen"];
const simpleKeys = ["locon", "lora", "embedding"];

const finalDb = {};

// Source of truth
const sourceData = uploadedDb.data || {};
const allNames = Object.keys(sourceData).sort();

console.log(`Processing ${allNames.length} models...`);

for (const name of allNames) {
    const rawModelData = sourceData[name];
    const rawFiles = rawModelData.models || {};

    const entry = {
        lastUpdated: rawModelData.lastUpdated,
        models: {},
    };

    // A. Simple types
    for (const key of simpleKeys) {
        entry.models[key] = rawFiles[key] || [];
    }

    // B. Complex types
    for (const key of complexKeys) {
        const filesList = rawFiles[key] || [];
        let imgList = [];

        const capitalizedKey = imgTypeMap[key];

        if (imagesDb[name] && imagesDb[name][capitalizedKey]) {
            imgList = imagesDb[name][capitalizedKey].map(item => item.url);
        }

        if (filesList.length === 0 && imgList.length === 0) {
            entry.models[key] = [];
        } else {
            entry.models[key] = {
                loras: filesList,
                images: imgList,
            };
        }
    }

    finalDb[name] = entry;
}

// 3. Save
const outputPath = path.resolve("./unified-data.json");
fs.writeFileSync(outputPath, JSON.stringify(finalDb, null, 2), "utf-8");

console.log(`Success! Database built at ${outputPath}`);