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

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

if (!values.dataDir || !values.workingDir || !values.script) {
    console.error("Usage: bun run script/train_factor.ts --workingDir <dir> --dataDir <dir> --script <python_script> --trainSize <num> --validSize <num>");
    process.exit(1);
}

// 1. factors.json 생성 (없을 경우)
generateFactors();

// 2. 파이썬 학습 스크립트 실행
console.log(`Starting training with factor-based script: ${values.script}`);
const child = spawn("python3", [
    values.script,
    "--workingDir", values.workingDir,
    "--dataDir", values.dataDir,
    "--trainSize", (Number(values.trainSize) || 1000).toString(),
    "--validSize", (Number(values.validSize) || 200).toString(),
]);

child.stdout.pipe(process.stdout);
child.stderr.pipe(process.stderr);

child.on("close", (code) => {
    console.log(`Training process exited with code ${code}`);
    process.exit(code || 0);
});

// functions
function generateFactors() {
    const workingDir = values.workingDir ?? '';
    if (!fs.existsSync(workingDir)) fs.mkdirSync(workingDir, { recursive: true });
    
    const factorsPath = path.join(workingDir, 'factors.json');
    if (fs.existsSync(factorsPath)) {
        console.log('factors.json already exists, skipping generation.');
        return;
    }

    const dataDir = values.dataDir ?? '';
    const tjaDir = path.join(dataDir, 'tja');
    if (!fs.existsSync(tjaDir)) {
        console.error(`TJA directory not found: ${tjaDir}`);
        return;
    }

    const files = fs.readdirSync(tjaDir).filter(f => f.endsWith('.tja'));
    console.log(`Generating factors from ${files.length} TJA files...`);

    const results: any[] = [];
    for (const file of files) {
        try {
            const tja = fs.readFileSync(path.join(tjaDir, file), 'utf-8');
            const songno = path.basename(file, '.tja');
            const parsed = parseTja(tja);
            
            const courses = [
                { diff: 'oni', data: parsed?.oni },
                { diff: 'ura', data: parsed?.edit }
            ];

            for (const course of courses) {
                if (course.data) {
                    results.push({
                        songno,
                        difficulty: course.diff,
                        factors: factorize(course.data)
                    });
                }
            }
        } catch (err) {
            console.error(`Error processing ${file}:`, err);
        }
    }

    fs.writeFileSync(factorsPath, JSON.stringify(results, null, 2), 'utf-8');
    console.log(`factors.json generated at ${factorsPath}`);
}