import { readdirSync, writeFileSync, existsSync } from 'fs';
import { join, extname, dirname } from 'path';
import { execSync } from 'child_process';

async function predictAll() {
    const args = process.argv.slice(2);
    if (args.length < 3) {
        console.log("Usage: bun script/predict-all.ts <workingDir> <tjaDir> <outputPath>");
        process.exit(1);
    }

    const workingDir = args[0];
    const tjaDir = args[1];
    const outputPath = args[2];
    const targetDifficulties = ['oni', 'edit'];

    if (!existsSync(tjaDir)) {
        console.error(`❌ Directory not found: ${tjaDir}`);
        process.exit(1);
    }

    const files = readdirSync(tjaDir).filter(file => extname(file).toLowerCase() === '.tja');
    console.log(`📂 Found ${files.length} TJA files. Processing Oni & Edit...`);

    const finalResults: any[] = [];

    let i = 0;
    for (const file of files) {
        const fullPath = join(tjaDir, file);

        for (const diff of targetDifficulties) {
            try {
                // 개별 예측 실행
                const cmd = `bun script/predict.ts "${workingDir}" "${fullPath}" "${diff}"`;
                const output = execSync(cmd, { stdio: ['pipe', 'pipe', 'ignore'] }).toString();

                const match = output.match(/(?:Prediction|Predicted Difficulty):\s*(\d+\.\d+)/);

                if (match) {
                    const value = parseFloat(match[1]);
                    console.log(`✅ [${file}] ${diff.toUpperCase()}: ${value}`);
                    finalResults.push({
                        filename: file,
                        difficulty_type: diff,
                        predicted_measure: value
                    });
                }
            } catch (e) {
                // 해당 난이도가 없는 경우 조용히 넘어감
            }
        }
        i++;
        console.log(`${i}/${files.length}`)
    }

    // 결과 저장 (확장자에 따라 분기)
    if (outputPath.endsWith('.json')) {
        writeFileSync(outputPath, JSON.stringify(finalResults, null, 2));
    } else {
        // CSV 저장 (기본값)
        const csvContent = [
            "filename,difficulty,predicted_measure",
            ...finalResults.map(r => `"${r.filename}",${r.difficulty_type},${r.predicted_measure}`)
        ].join('\n');
        writeFileSync(outputPath, csvContent);
    }

    console.log(`\n✨ Prediction complete! Results saved to: ${outputPath}`);
}

predictAll();