const TrainingProject = require('../models/TrainingProject.js');
const TrainingProjectDetails = require('../models/TrainingProjectDetails.js')
const LabelStudioProject = require('../models/LabelStudioProject.js')
const Annotation = require('../models/Annotation.js')
const Images = require('../models/Images.js')
const fs = require('fs');


async function generateTrainingJson(trainingId){
    // trainingId is now project_details_id
    const trainingProjectDetails = await TrainingProjectDetails.findByPk(trainingId);
    if (!trainingProjectDetails) throw new Error('No TrainingProjectDetails found for project_details_id ' + trainingId);
    const detailsObj = trainingProjectDetails.get({ plain: true });
    // Get parent project for name
    const trainingProject = await TrainingProject.findByPk(detailsObj.project_id);
    // Get split percentages (assume they are stored as train_percent, valid_percent, test_percent)
    const trainPercent = detailsObj.train_percent || 85;
    const validPercent = detailsObj.valid_percent || 10;
    const testPercent = detailsObj.test_percent || 5;

    let cocoImages = [];
    let cocoAnnotations = [];
    let cocoCategories = [];
    let categoryMap = {};
    let categoryId = 0;
    let imageid = 0;
    let annotationid = 0;

    for (const cls of detailsObj.class_map) {
        const asgMap = [];
        const listAsg = cls[1];
        for(const asg of listAsg){
            asgMap.push ({ original: asg[0], mapped: asg[1] });
            // Build category list and mapping
            if (asg[1] && !(asg[1] in categoryMap)) {
                categoryMap[asg[1]] = categoryId;
                cocoCategories.push({ id: categoryId, name: asg[1], supercategory: '' });
                categoryId++;
            }
        }
        const images = await Images.findAll({ where: { project_id: cls[0] } });
        for(const image of images){
            imageid += 1;
            let fileName = image.image_path;
            if (fileName.includes('%20')) {
                fileName = fileName.replace(/%20/g, ' ');
            }
            if (fileName && fileName.startsWith('/data/local-files/?d=')) {
                fileName = fileName.replace('/data/local-files/?d=', '');
                fileName = fileName.replace('/home/kitraining/home/kitraining/', '');
            }
            if (fileName && fileName.startsWith('home/kitraining/To_Annotate/')) {
                fileName = fileName.replace('home/kitraining/To_Annotate/','');
            }
            // Get annotations for this image
            const annotations = await Annotation.findAll({ where: { image_id: image.image_id } });
            // Use image.width and image.height from DB (populated from original_width/original_height)
            cocoImages.push({
                id: imageid,
                file_name: fileName,
                width: image.width || 0,
                height: image.height || 0
            });
            for (const annotation of annotations) {
                // Translate class name using asgMap
                let mappedClass = annotation.Label;
                for (const mapEntry of asgMap) {
                    if (annotation.Label === mapEntry.original) {
                        mappedClass = mapEntry.mapped;
                        break;
                    }
                }
                // Only add annotation if mappedClass is valid
                if (mappedClass && mappedClass in categoryMap) {
                    annotationid += 1;
                    let area = 0;
                    if (annotation.width && annotation.height) {
                        area = annotation.width * annotation.height;
                    }
                    cocoAnnotations.push({
                        id: annotationid,
                        image_id: imageid,
                        category_id: categoryMap[mappedClass],
                        bbox: [annotation.x, annotation.y, annotation.width, annotation.height],
                        area: area,
                        iscrowd: annotation.iscrowd || 0
                    });
                }
            }
        }
    }

    // Shuffle images for random split using seed
    function seededRandom(seed) {
        let x = Math.sin(seed++) * 10000;
        return x - Math.floor(x);
    }
    function shuffle(array, seed) {
        for (let i = array.length - 1; i > 0; i--) {
            const j = Math.floor(seededRandom(seed + i) * (i + 1));
            [array[i], array[j]] = [array[j], array[i]];
        }
    }
    // Use seed from detailsObj if present, else default to 42
    const splitSeed = detailsObj.seed !== undefined && detailsObj.seed !== null ? Number(detailsObj.seed) : 42;
    shuffle(cocoImages, splitSeed);

    // Split images
    const totalImages = cocoImages.length;
    const trainCount = Math.floor(totalImages * trainPercent / 100);
    const validCount = Math.floor(totalImages * validPercent / 100);
    const testCount = totalImages - trainCount - validCount;

    const trainImages = cocoImages.slice(0, trainCount);
    const validImages = cocoImages.slice(trainCount, trainCount + validCount);
    const testImages = cocoImages.slice(trainCount + validCount);

    // Helper to get image ids for each split
    const trainImageIds = new Set(trainImages.map(img => img.id));
    const validImageIds = new Set(validImages.map(img => img.id));
    const testImageIds = new Set(testImages.map(img => img.id));

    // Split annotations
    const trainAnnotations = cocoAnnotations.filter(ann => trainImageIds.has(ann.image_id));
    const validAnnotations = cocoAnnotations.filter(ann => validImageIds.has(ann.image_id));
    const testAnnotations = cocoAnnotations.filter(ann => testImageIds.has(ann.image_id));

    // Build final COCO JSONs with info section
    const buildCocoJson = (images, annotations, categories) => ({
        images,
        annotations,
        categories
    });

    // Build COCO JSONs with info section
    const trainJson = buildCocoJson(trainImages, trainAnnotations, cocoCategories);
    const validJson = buildCocoJson(validImages, validAnnotations, cocoCategories);
    const testJson = buildCocoJson(testImages, testAnnotations, cocoCategories);

    // Create output directory: projectname/trainingid/annotations
    const projectName = trainingProject && trainingProject.name ? trainingProject.name.replace(/\s+/g, '_') : `project_${detailsObj.project_id}`;
    const outDir = `${projectName}/${trainingId}`;
    const annotationsDir = `/home/kitraining/To_Annotate/annotations`;
    if (!fs.existsSync(annotationsDir)) {
        fs.mkdirSync(annotationsDir, { recursive: true });
    }

    // Write to files in the annotations directory
    const trainPath = `${annotationsDir}/coco_project_${trainingId}_train.json`;
    const validPath = `${annotationsDir}/coco_project_${trainingId}_valid.json`;
    const testPath = `${annotationsDir}/coco_project_${trainingId}_test.json`;
    fs.writeFileSync(trainPath, JSON.stringify(trainJson, null, 2));
    fs.writeFileSync(validPath, JSON.stringify(validJson, null, 2));
    fs.writeFileSync(testPath, JSON.stringify(testJson, null, 2));
    console.log(`COCO JSON splits written to ${annotationsDir} for trainingId ${trainingId}`);



    // Also generate inference exp.py in the same output directory as exp.py (project folder in workspace)
    const { generateYoloxInferenceExp } = require('./generate-yolox-exp');
    const path = require('path');
    const projectFolder = path.join(__dirname, '..', projectName, String(trainingId));
    if (!fs.existsSync(projectFolder)) {
        fs.mkdirSync(projectFolder, { recursive: true });
    }
    const inferenceExpPath = path.join(projectFolder, 'exp_infer.py');
    generateYoloxInferenceExp(trainingId).then(expContent => {
        fs.writeFileSync(inferenceExpPath, expContent);
        console.log(`Inference exp.py written to ${inferenceExpPath}`);
    }).catch(err => {
        console.error('Failed to generate inference exp.py:', err);
    });
}


module.exports = {generateTrainingJson};