import Bun from 'bun';
import { execSync, spawn, spawnSync } from 'node:child_process';
import { parseArgs } from 'node:util';
import fs from 'fs';
import path from 'path';
import { featurize } from '../preprocess/featurize';
import { parseTja } from '../preprocess/parse'

const { values } = parseArgs({
    args: Bun.argv,
    options: {
        workingDir: {
            type: "string"
        },
        dataDir: {
            type: "string"
        },
        trainScript: {
            type: "string"
        },
        trainSize: {
            type: 'string'
        },
        validSize: {
            type: 'string'
        }
    },
    allowPositionals: true,
});

if (!values.dataDir || !values.workingDir || !values.trainScript) {
    console.error("--workingDir --dataDir --trainDir");
    process.exit(1);
}

generateFeatures();
const child = spawn("python3", [values.trainScript,
    "--workingDir", values.workingDir,
    "--dataDir", values.dataDir,
    "--trainSize", (Number(values.trainSize) || 500).toString(),
    "--validSize", (Number(values.validSize) || 100).toString(),
]);
child.stdout.pipe(process.stdout);
child.stderr.pipe(process.stderr);
child.on("close", () => {
    process.exit()
})

// funcs
function generateFeatures() {
    const workingDir = values.workingDir ?? '';
    if (!fs.existsSync(workingDir)) fs.mkdirSync(workingDir);
    const featurePath = path.join(workingDir, 'features.json');
    if (fs.existsSync(featurePath)) return;
    const dataDir = values.dataDir ?? '';

    const tjaDir = path.join(dataDir, 'tja');
    const files = fs.readdirSync(tjaDir);

    const features: ({ songno: string, difficulty: 'oni' | 'ura' } & {})[] = [];
    for (const file of files) {
        const tja = fs.readFileSync(path.join(tjaDir, file), 'utf-8');
        const songno = path.basename(file, '.tja');
        try {
            const parsed = parseTja(tja);
            const oni = parsed?.oni;
            const edit = parsed?.edit;
            if (oni) {
                features.push({
                    songno,
                    difficulty: 'oni',
                    ...featurize(oni)
                })
            }
            if (edit) {
                features.push({
                    songno,
                    difficulty: 'ura',
                    ...featurize(edit)
                })
            }
        }
        catch (err) {
            console.error(err);
            console.error(file);
        }
    }

    fs.writeFileSync(featurePath, JSON.stringify(features, null, 2), 'utf-8');
    console.log('features.json generated')
}