/*******************************************************************************
 * Danbooru Diffusion Prompt Builder
 * Copyright (C) 2022  Jabasukuriputo Wang
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 *
 ******************************************************************************/

import yaml from 'js-yaml'
import fs from 'fs'
import glob from 'glob'
import path from 'path'
import { fileURLToPath } from 'url'
import { type TagFile } from '../types/file'

const dirname = path.dirname(fileURLToPath(import.meta.url))
const root = path.resolve(dirname, '../../')

const tagFiles = glob.sync('**/*.yaml', {
    cwd: path.resolve(dirname, '../../data/tags'),
})
const tagSet: Set<string> = new Set()

const uc =
    'lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, nsfw'
const ucA = uc.split(', ')

for (const file of tagFiles) {
    if (file.startsWith('restricted')) continue
    const tagData: TagFile = yaml.load(
        fs.readFileSync(path.resolve(dirname, '../../data/tags', file), 'utf-8')
    ) as TagFile
    for (const [tag, meta] of Object.entries(tagData.content)) {
        if (ucA.includes(tag)) continue
        if (!meta.image) {
            tagSet.add(tag)
        }
    }
}

const generatedPath = path.resolve(root, 'workspace/generated_images')
fs.mkdirSync(generatedPath, { recursive: true })

for (const tag of tagSet) {
    if (
        fs.existsSync(
            path.resolve(generatedPath, `${encodeURIComponent(tag)}.png`)
        )
    ) {
        continue
    }
    console.log(tag)

    const params = JSON.stringify({
        prompt: `masterpiece, best quality, ${tag}`,
        negative_prompt: uc,
        steps: 28,
        cfg_scale: 12,
        do_not_save_samples: true,
        width: 512,
        height: 512,
        clip_skip: 2,
        sampler_index: 'Euler a',
        sd_model: 'modelsfw.pruned.ckpt [1d4a34af]',
    })

    const res = await fetch('http://127.0.0.1:7860/sdapi/v1/txt2img', {
        method: 'POST',
        body: params,
        headers: {
            'Content-Type': 'application/json',
        },
    })
    const json = await res.json()
    await fs.promises.writeFile(
        path.resolve(generatedPath, `${encodeURIComponent(tag)}.png`),
        Buffer.from(json.images[0], 'base64')
    )
}
