const fs = require('fs');
const path = require('path');
const Training = require('../models/training.js');
const TrainingProject = require('../models/TrainingProject.js');

// Remove Python comments and legacy code
const exp_names = [
    'YOLOX-s',
    'YOLOX-m',
    'YOLOX-l',
    'YOLOX-x', 
    'YOLOX-Darknet53', //todo
    'YOLOX-Nano',
    'YOLOX-Tiny'
]

//TODO: Clean up generation of exp_names.py and remove second exp creation!!!


// Refactored: Accept trainingId, fetch info from DB
async function generateYoloxExp(trainingId) {
    // Fetch training row from DB by project_details_id if not found by PK
    let training = await Training.findByPk(trainingId);
    if (!training) {
        training = await Training.findOne({ where: { project_details_id: trainingId } });
    }
    if (!training) throw new Error('Training not found for trainingId or project_details_id: ' + trainingId);

    // If transfer_learning is 'coco', just return the path to the default exp.py
    if (training.transfer_learning === 'coco') {
        const selectedModel = training.selected_model.toLowerCase().replace('-', '_');
        const expSourcePath = `/home/kitraining/Yolox/YOLOX-main/exps/default/${selectedModel}.py`;
        if (!fs.existsSync(expSourcePath)) {
            throw new Error(`Default exp.py not found for model: ${selectedModel} at ${expSourcePath}`);
        }
        // Copy to project folder (e.g., /home/kitraining/coco_tool/backend/project_XX/YY/exp.py)
        const projectDetailsId = training.project_details_id;
        const projectFolder = path.resolve(__dirname, `../project_23/${projectDetailsId}`);
        if (!fs.existsSync(projectFolder)) {
            fs.mkdirSync(projectFolder, { recursive: true });
        }
        const expDestPath = path.join(projectFolder, 'exp.py');
        fs.copyFileSync(expSourcePath, expDestPath);
        return { type: 'default', expPath: expDestPath };
    }

    // If transfer_learning is 'sketch', generate a custom exp.py as before
    if (training.transfer_learning === 'sketch') {
        // ...existing custom exp.py generation logic here (copy from previous implementation)...
        // For brevity, you can call generateYoloxInferenceExp or similar here, or inline the logic.
        // Example:
        const expContent = await generateYoloxInferenceExp(trainingId);
        return { type: 'custom', expContent };
    }

    throw new Error('Unknown transfer_learning type: ' + training.transfer_learning);
}

async function saveYoloxExp(trainingId, outPath) {
    const expResult = await generateYoloxExp(trainingId);
    if (expResult.type === 'custom' && expResult.expContent) {
        fs.writeFileSync(outPath, expResult.expContent);
        return outPath;
    } else if (expResult.type === 'default' && expResult.expPath) {
        // Optionally copy the file if outPath is different
        if (expResult.expPath !== outPath) {
            fs.copyFileSync(expResult.expPath, outPath);
        }
        return outPath;
    } else {
        throw new Error('Unknown expResult type or missing content');
    }
}

// Generate a second exp.py for inference, using the provided template and DB values
async function generateYoloxInferenceExp(trainingId, options = {}) {
    let training = await Training.findByPk(trainingId);
    if (!training) {
        training = await Training.findOne({ where: { project_details_id: trainingId } });
    }
    if (!training) throw new Error('Training not found for trainingId or project_details_id: ' + trainingId);
    // Always use the trainingId (project_details_id) for annotation file names
    const projectDetailsId = training.project_details_id;
    const dataDir = options.data_dir || '/home/kitraining/To_Annotate/';
    const trainAnn = options.train_ann || `coco_project_${trainingId}_train.json`;
    const valAnn = options.val_ann || `coco_project_${trainingId}_valid.json`;
    const testAnn = options.test_ann || `coco_project_${trainingId}_test.json`;
    // Get num_classes from TrainingProject.classes JSON
    let numClasses = 80;
    try {
        const trainingProject = await TrainingProject.findByPk(projectDetailsId);
        if (trainingProject && trainingProject.classes) {
            let classesArr = trainingProject.classes;
            if (typeof classesArr === 'string') {
                classesArr = JSON.parse(classesArr);
            }
            if (Array.isArray(classesArr)) {
                numClasses = classesArr.filter(c => c !== null && c !== undefined && c !== '').length;
            } else if (typeof classesArr === 'object' && classesArr !== null) {
                numClasses = Object.keys(classesArr).filter(k => classesArr[k] !== null && classesArr[k] !== undefined && classesArr[k] !== '').length;
            }
        }
    } catch (e) {
        console.warn('Could not determine num_classes from TrainingProject.classes:', e);
    }
    const depth = options.depth || training.depth || 1.00;
    const width = options.width || training.width || 1.00;
    const inputSize = options.input_size || training.input_size || [640, 640];
    const mosaicScale = options.mosaic_scale || training.mosaic_scale || [0.1, 2];
    const randomSize = options.random_size || training.random_size || [10, 20];
    const testSize = options.test_size || training.test_size || [640, 640];
    const expName = options.exp_name || 'inference_exp';
    const enableMixup = options.enable_mixup !== undefined ? options.enable_mixup : false;
    let expContent = '';
    expContent += `#!/usr/bin/env python3\n# -*- coding:utf-8 -*-\n# Copyright (c) Megvii, Inc. and its affiliates.\n\nimport os\n\nfrom yolox.exp import Exp as MyExp\n\n\nclass Exp(MyExp):\n    def __init__(self):\n        super(Exp, self).__init__()\n        self.data_dir = "${dataDir}"\n        self.train_ann = "${trainAnn}"\n        self.val_ann = "${valAnn}"\n        self.test_ann = "coco_project_${trainingId}_test.json"\n        self.num_classes = ${numClasses}\n`;
    // Set pretrained_ckpt if transfer_learning is 'coco'
    if (training.transfer_learning && typeof training.transfer_learning === 'string' && training.transfer_learning.toLowerCase() === 'coco') {
        const yoloxBaseDir = '/home/kitraining/Yolox/YOLOX-main';
        const selectedModel = training.selected_model ? training.selected_model.replace(/\.pth$/i, '') : '';
        if (selectedModel) {
            expContent += `        self.pretrained_ckpt = r'${yoloxBaseDir}/pretrained/${selectedModel}.pth'\n`;
        }
    }
    expContent += `        self.depth = ${depth}\n        self.width = ${width}\n        self.input_size = (${Array.isArray(inputSize) ? inputSize.join(', ') : inputSize})\n        self.mosaic_scale = (${Array.isArray(mosaicScale) ? mosaicScale.join(', ') : mosaicScale})\n        self.random_size = (${Array.isArray(randomSize) ? randomSize.join(', ') : randomSize})\n        self.test_size = (${Array.isArray(testSize) ? testSize.join(', ') : testSize})\n        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]\n        self.enable_mixup = ${enableMixup ? 'True' : 'False'}\n`;
    return expContent;
}

// Save inference exp.py to a custom path
async function saveYoloxInferenceExp(trainingId, outPath, options = {}) {
    const expContent = await generateYoloxInferenceExp(trainingId, options);
    fs.writeFileSync(outPath, expContent);
    return outPath;
}

module.exports = { generateYoloxExp, saveYoloxExp, generateYoloxInferenceExp, saveYoloxInferenceExp };