#!/usr/bin/env node
"use strict";
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    var desc = Object.getOwnPropertyDescriptor(m, k);
    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
      desc = { enumerable: true, get: function() { return m[k]; } };
    }
    Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
    Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
    o["default"] = v;
});
var __importStar = (this && this.__importStar) || function (mod) {
    if (mod && mod.__esModule) return mod;
    var result = {};
    if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
    __setModuleDefault(result, mod);
    return result;
};
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
const commander_1 = require("commander");
const lib_1 = require("./lib");
const cli_table3_1 = __importDefault(require("cli-table3"));
const fs = __importStar(require("fs"));
const config_1 = require("./config");
const chalk_1 = __importDefault(require("chalk"));
const outdent_1 = require("outdent");
const program = new commander_1.Command();
// if (!config.protege.apiKey) {
//     console.log(`API Key not found. Please run protege config to configure the CLI`);
//     process.exit(1);
// }
const options = {
    endpoint: config_1.config.protege.endpoint,
    auth: new lib_1.DeviceAuth()
};
const protege = new lib_1.ProtegeEngineSDK(options);
program
    .version('0.0.1')
    .description('Protege Engine SDK CLI');
// Operations for configuring the CLI 
// Updates or creates an env file at ~/.protege
program
    .command('config')
    .alias('c')
    .description('Configure the CLI')
    .option('-e, --endpoint <endpoint>', 'endpoint for the Protege Engine')
    .option('-k, --apiKey <apiKey>', 'apiKey for the Protege Engine')
    .action(async (options) => {
    // create a new env file
    const envFile = (0, outdent_1.outdent) `# Protege Engine SDK CLI config file
        # This file was created by the Protege Engine SDK CLI
        # You can edit this file manually or use the CLI to update it
        
        # Protege Engine endpoint
        PROTEGE_ENGINE_ENDPOINT=${options.endpoint ?? config_1.config.protege.endpoint}
        # Protege Engine apiKey
        PROTEGE_ENGINE_API_KEY=${options.apiKey ?? config_1.config.protege.apiKey}
        `;
    fs.writeFileSync(`${process.env.HOME}/.protege`, envFile);
    console.log(`Config file saved to ${process.env.HOME}/.protege`);
});
// CRUD operations for the User entity
// TODO 
// CRUD operations for the InferenceBackend entity
// create a new inference backend
program
    .command('backend:create <name> <urls...>')
    .alias('bc')
    .option('-a, --apiKey <apiKey>', 'apiKey for the inference backend')
    .option('--user <userId>', 'userId for the inference backend')
    .description('Create a new inference backend')
    .action(async (name, urls, apiKey, user) => {
    const result = await protege.inferenceBackend.createInferenceBackend({
        user: {
            connect: {
                id: user ? parseInt(user) : 1 // TODO: get the user id from the config
            }
        },
        name,
        urls: {
            set: urls
        },
        apiKey: apiKey ? apiKey.apiKey : undefined,
    });
    // display the result in a table
    const table = new cli_table3_1.default({
        head: ['id', 'name', 'active', 'urls', 'apiKey']
    });
    table.push([
        result.createOneInferenceBackend.id.toString(),
        result.createOneInferenceBackend.name,
        result.createOneInferenceBackend.active.toString(),
        result.createOneInferenceBackend.urls.join(', '),
        result.createOneInferenceBackend.apiKey
    ]);
    console.log(table.toString());
});
// list all inference backends
program
    .command('backend:list')
    .alias('bl')
    .description('List all inference backends')
    .action(async () => {
    const result = await protege.inferenceBackend.listInferenceBackends({});
    // display the result in a table
    const table = new cli_table3_1.default({
        head: ['id', 'name', 'active', 'urls', 'apiKey']
    });
    result.inferenceBackends.forEach((backend) => {
        table.push([
            backend.id.toString(),
            backend.name,
            backend.active.toString(),
            backend.urls.join(', '),
            backend.apiKey
        ]);
    });
    console.log(table.toString());
});
// get a inference backend by id
program
    .command('backend:get <id>')
    .alias('bg')
    .description('Get a inference backend by id')
    .action(async (id) => {
    const result = await protege.inferenceBackend.getInferenceBackend(parseInt(id));
    if (!result.inferenceBackend) {
        console.log(`Inference backend with id ${id} not found`);
        return;
    }
    // display the result in a table
    const table = new cli_table3_1.default({
        head: ['id', 'name', 'active', 'urls', 'apiKey'],
    });
    table.push([
        result.inferenceBackend.id.toString(),
        result.inferenceBackend.name,
        result.inferenceBackend.active.toString(),
        result.inferenceBackend.urls.join(', '),
        result.inferenceBackend.apiKey
    ]);
    console.log(table.toString());
});
// update a inference backend
program
    .command('backend:update <id>')
    .alias('bu')
    .option('-n, --name <name>', 'name of the inference backend')
    .option('-a, --active <active>', 'active status of the inference backend')
    .option('-u, --urls <urls...>', 'urls of the inference backend')
    .option('-k, --apiKey <apiKey>', 'apiKey for the inference backend')
    .description('Update a inference backend')
    .action(async (id, options) => {
    let result;
    try {
        result = await protege.inferenceBackend.updateInferenceBackend(parseInt(id), {
            name: options.name ? { set: options.name } : undefined,
            active: options.active ? { set: options.active } : undefined,
            urls: options.urls ? { set: options.urls } : undefined,
            apiKey: options.apiKey ? { set: options.apiKey } : undefined,
        });
    }
    catch (e) {
        const clientError = e;
        if (clientError.response.errors[0].message.includes("Record to update not found.")) {
            console.log(`Inference backend with id ${id} not found`);
        }
        else {
            // display the error in a table
            const table = new cli_table3_1.default({
                head: ['message', 'locations', 'path', 'extensions']
            });
            table.push([
                clientError.response.errors[0].message,
                JSON.stringify(clientError.response.errors[0].locations, null, 2),
                clientError.response.errors[0].path.join(', '),
                JSON.stringify(clientError.response.errors[0].extensions, null, 2)
            ]);
        }
        return;
    }
    if (!result.updateOneInferenceBackend) {
        console.log(`Inference backend with id ${id} not found`);
        return;
    }
    // display the result in a table
    const table = new cli_table3_1.default({
        head: ['id', 'name', 'active', 'urls', 'apiKey']
    });
    table.push([
        result.updateOneInferenceBackend.id.toString(),
        result.updateOneInferenceBackend.name,
        result.updateOneInferenceBackend.active.toString(),
        result.updateOneInferenceBackend.urls.join(', '),
        result.updateOneInferenceBackend.apiKey
    ]);
    console.log(table.toString());
});
// delete a inference backend
program
    .command('backend:delete <id>')
    .alias('bd')
    .description('Delete a inference backend')
    .action(async (id) => {
    let result;
    try {
        result = await protege.inferenceBackend.deleteInferenceBackend(parseInt(id));
    }
    catch (e) {
        const clientError = e;
        console.log(`Error: ${clientError.message}`);
        return;
    }
    if (!result.deleteOneInferenceBackend) {
        console.log(`Inference backend with id ${id} not found`);
        return;
    }
    // display the result in a table
    const table = new cli_table3_1.default({
        head: ['id']
    });
    table.push([
        result.deleteOneInferenceBackend.id.toString()
    ]);
    console.log(table.toString());
});
// CRUD operations for the Prediction entity
// list all predictions, with optional datasetId, skip and limit
program
    .command('prediction:list')
    .alias('pl')
    .option('-d, --datasetId <datasetId>', 'datasetId for the prediction')
    .option('-s, --skip <skip>', 'skip for the prediction', "0")
    .option('-l, --limit <limit>', 'limit for the prediction', "10")
    .option('--labels', 'include labels in the prediction')
    .option('--label-contains <labelContains>', 'filter predictions by label value')
    .description('List all predictions')
    .action(async (options) => {
    console.log(`DEBUG -- options: ${JSON.stringify(options, null, 2)}`);
    const where = {
        datasetId: options.datasetId ? parseInt(options.datasetId) : undefined,
        labels: options.labelContains ? { some: { value: { contains: options.labelContains } } } : undefined
    };
    const result = await protege.prediction.listPredictions({
        where,
        skip: parseInt(options.skip),
        take: parseInt(options.limit)
    });
    // display the result in a table
    const table = new cli_table3_1.default({
        head: ['id', 'createdAt', 'model', 'labels', 'prompt', "completion"],
        colWidths: [5, 28, 15, options.labels ? 35 : 10, 50, 20],
        wordWrap: true
    });
    result.predictions.forEach((prediction) => {
        var labelString = "";
        if (options.labels) {
            prediction.labels.forEach((label) => {
                if (label.confidence < 1 && label.confidence > 0.8) {
                    labelString += `${label.value} (${chalk_1.default.green(label.confidence.toString())})\n`;
                }
                else if (label.confidence < 0.8 && label.confidence > 0.5) {
                    labelString += `${label.value} (${chalk_1.default.yellow(label.confidence.toString())})\n`;
                }
                else if (label.confidence < 0.5) {
                    labelString += `${label.value} (${chalk_1.default.red(label.confidence.toString())})\n`;
                }
                else {
                    labelString += `${label.value} (${label.confidence})\n`;
                }
            });
        }
        table.push([
            prediction.id.toString(),
            prediction.createdAt.toString(),
            prediction.modelId.toString(),
            options.labels ? labelString : prediction._count.labels.toString(),
            prediction.prompt,
            prediction.completion,
        ]);
        // // if the labels option is set, display the labels for each prediction in an additional row
        // if (options.labels) {
        //     // display a sub-table with the labels in rows of 4
        //     for (let i = 0; i < prediction.labels.length; i+=4) {
        //         table.push(
        //             prediction.labels.slice(i, i+4).map((label) => {
        //                 return { colspan: 1, content: `${label.name}`,  }
        //             }), 
        //             prediction.labels.slice(i, i+4).map((label) => {
        //                 return { colspan: 1, content: `${label.value} (${label.confidence})` }
        //             })
        //         );
        //     }
        // }  
    });
    console.log(table.toString());
});
// count all predictions, with optional datasetId
program
    .command('prediction:count')
    .alias('pc')
    .option('-d, --datasetId <datasetId>', 'datasetId for the prediction')
    .description('Count all predictions')
    .action(async (options) => {
    const result = await protege.prediction.countPredictions(options.datasetId);
    if (!result.aggregatePrediction) {
        console.log(`No predictions found`);
        return;
    }
    // display the result in a table
    const table = new cli_table3_1.default({
        head: ['_all']
    });
    table.push([
        result.aggregatePrediction._count?._all.toString()
    ]);
    console.log(table.toString());
});
// CRUD operations for the Dataset entity
// list all datasets
program
    .command('dataset:list')
    .alias('dl')
    .description('List all datasets')
    .action(async () => {
    const result = await protege.dataset.listDatasets({});
    if (!result.datasets) {
        console.log(`No datasets found`);
        return;
    }
    // display the result in a table
    const table = new cli_table3_1.default({
        head: ['name', 'predictions', 'trainingJobs', 'trainingPredictions', 'evaluationJobs', 'updatedAt']
    });
    result.datasets.forEach((dataset) => {
        table.push([
            dataset.name,
            dataset._count.predictions.toString(),
            dataset._count.trainingPredictions.toString(),
            dataset.updatedAt.toString()
        ]);
    });
    console.log(table.toString());
});
// list all predictions for a given dataset
program
    .command('dataset:listPredictions <datasetId>')
    .alias('dlp')
    .description('List all predictions for a given dataset')
    .action(async (datasetId) => {
    const result = await protege.prediction.listPredictions({
        where: {
            datasets: {
                some: {
                    id: {
                        equals: parseInt(datasetId)
                    }
                }
            }
        }
    });
    // display the result in a table
    const table = new cli_table3_1.default({
        head: ['id', 'createdAt', 'model', 'labels', 'prompt', "completion"],
        colWidths: [5, 28, 15, 10, 50, 50],
        wordWrap: true
    });
    result.predictions.forEach((prediction) => {
        table.push([
            prediction.id.toString(),
            prediction.createdAt.toString(),
            prediction.modelId.toString(),
            prediction._count.labels.toString(),
            prediction.prompt,
            prediction.completion
        ]);
    });
    console.log(table.toString());
});
// download dataset CSV
program
    .command('dataset:download <id>')
    .alias('dd')
    .option('-o, --output <output>', 'output file', '{name}-dataset.csv')
    .option('-s, --split', 'split the dataset into training and evaluation sets')
    .description('Download dataset CSV')
    .action(async (id, options) => {
    const datasetId = parseInt(id);
    const dataset = await protege.dataset.getDataset(datasetId);
    if (!dataset.dataset) {
        console.log(`Dataset with id ${id} not found`);
        return;
    }
    const csv = await protege.dataset.getDatasetCSV(datasetId);
    // write the result to a file
    console.log(`DEBUG -- outputFile: ${JSON.stringify(options.output, null, 2)}`);
    let outputFile;
    if (options.output.includes("{name}")) {
        outputFile = options.output.replace("{name}", dataset.dataset.name);
    }
    if (options.split) {
        // shuffle the dataset
        const shuffledDataset = csv.datasetCSV.split("\r\n\r\n").slice(1).sort(() => Math.random() - 0.5);
        console.log(`DEBUG -- shuffledDataset: Dataset has ${shuffledDataset.length} predictions`);
        // split the dataset into training and evaluation sets
        const splitIndex = Math.floor(shuffledDataset.length * 0.8);
        console.log(`DEBUG -- splitIndex: ${splitIndex}`);
        const headers = "prompt,completion";
        const trainingCSV = headers + "\n" + shuffledDataset.slice(0, splitIndex).join("\n");
        console.log(`DEBUG -- trainingCSV has ${trainingCSV.split("\n").length} lines`);
        const evaluationCSV = headers + "\n" + shuffledDataset.slice(splitIndex).join("\n");
        console.log(`DEBUG -- evaluationCSV has ${evaluationCSV.split("\n").length} lines`);
        const trainingOutputFile = outputFile.replace(".csv", "-training.csv");
        const evaluationOutputFile = outputFile.replace(".csv", "-evaluation.csv");
        fs.writeFileSync(trainingOutputFile, trainingCSV, { flag: "w" });
        fs.writeFileSync(evaluationOutputFile, evaluationCSV, { flag: "w" });
        // write first 10 lines of each file to the console
        console.log(`DEBUG -- trainingCSV: ${trainingCSV.split("\n").slice(0, 10).join("\n")}`);
        console.log(`DEBUG -- evaluationCSV: ${evaluationCSV.split("\n").slice(0, 10).join("\n")}`);
    }
    else {
        fs.writeFileSync(outputFile, csv.datasetCSV);
    }
    console.log(`Dataset CSV saved to ${outputFile}`);
});
// Create a dataset
program
    .command('dataset:create <name>')
    .alias('dc')
    .description('Create a dataset')
    .action(async (name) => {
    const result = await protege.dataset.createDataset({
        name,
        user: {
            connect: {
                id: 1
            }
        }
    });
    // display the result in a table
    const table = new cli_table3_1.default({
        head: ['id', 'name']
    });
    table.push([
        result.createOneDataset.id.toString(),
        result.createOneDataset.name
    ]);
    console.log(table.toString());
});
// Add a range of prediction ids to a dataset
program
    .command('dataset:addPredictions <datasetId> <predictionIdStart> <predictionIdEnd>')
    .alias('dap')
    .description('Add a range of prediction ids to a dataset')
    .action(async (datasetId, predictionIdStart, predictionIdEnd) => {
    const dataset = await protege.dataset.getDataset(parseInt(datasetId));
    if (!dataset.dataset) {
        console.log(`Dataset with id ${datasetId} not found`);
        return;
    }
    const predictionIds = [];
    for (let i = parseInt(predictionIdStart); i <= parseInt(predictionIdEnd); i++) {
        predictionIds.push(i);
    }
    // batch add the predictions to the dataset in batches of 1000
    const batchSize = 1000;
    for (let i = 0; i < predictionIds.length; i += batchSize) {
        const result = await protege.prediction.addPredictionsToDataset(predictionIds.slice(i, i + batchSize), parseInt(datasetId));
        if (!result.updateOneDataset) {
            console.log(`Dataset with id ${datasetId} not found`);
            return;
        }
    }
    // // display the result in a table
    // const table = new Table({
    //     head: ['id', 'predictions']
    // });
    // table.push([
    //     result.updateOneDataset.id.toString(), 
    //     result.updateOneDataset._count!.predictions.toString()
    // ]);
    // console.log(table.toString());
    console.log(`Added predictions ${predictionIdStart} to ${predictionIdEnd} to dataset ${datasetId}`);
});
/// CRUD operations for the Label entity
// list all labels, with optional predictionId, skip and limit
// program 
//     .command('label:list')
//     .alias('ll')
//     .option('-p, --predictionId <predictionId>', 'predictionId for the label')
//     .option('-s, --skip <skip>', 'skip for the label', "0")
//     .option('-l, --limit <limit>', 'limit for the label', "10")
//     .description('List all labels')
//     .action(async (options) => {
//         const result = await protege.label.listLabels(
//             options.predictionId ? parseInt(options.predictionId) : undefined, 
//             parseInt(options.skip), 
//             parseInt(options.limit)
//         );
//         // display the result in a table
//         const table = new Table({
//             head: ['id', 'value', 'confidence', 'predictionId']
//         });
//         result.labels.forEach((label) => {
//             table.push([
//                 label.id.toString(), 
//                 label.value, 
//                 label.confidence.toString(), 
//                 label.predictionId ? label.predictionId.toString() : ""
//             ]);
//         });
//         console.log(table.toString());
//     });
// list problematic labels
// fetch all labels and aggregate by value, grouping by average confidence medium (0.8 > confidence > 0.5) and low (0.5 > confidence)
// display the worst performing 10 labels in each category where the number of occurrences is greater than 1
program
    .command('label:listProblematic')
    .alias('lpr')
    .description('List problematic labels')
    .action(async () => {
    const result = await protege.label.listLabels({
        take: 10000,
        skip: 0,
    });
    const labelMap = new Map();
    result.labels.forEach((label) => {
        if (labelMap.has(label.value)) {
            labelMap.get(label.value).push(label.confidence);
        }
        else {
            labelMap.set(label.value, [label.confidence]);
        }
    });
    // We filter out labels with a single occurrence before calculating the averages
    const filteredLabelMap = new Map([...labelMap].filter(([_, confidences]) => confidences.length > 1));
    const labelConfidenceMap = new Map();
    filteredLabelMap.forEach((confidenceArray, label) => {
        const averageConfidence = confidenceArray.reduce((a, b) => a + b, 0) / confidenceArray.length;
        labelConfidenceMap.set(label, averageConfidence);
    });
    // Separating labels into two categories: medium and low confidence
    const mediumConfidenceLabels = [];
    const lowConfidenceLabels = [];
    labelConfidenceMap.forEach((averageConfidence, label) => {
        if (averageConfidence > 0.5 && averageConfidence <= 0.8) {
            mediumConfidenceLabels.push(label);
        }
        else if (averageConfidence <= 0.5) {
            lowConfidenceLabels.push(label);
        }
        // High confidence labels are not required as per the comment
    });
    // sort the labels by average confidence and number of occurrences
    const sortLabels = (a, b) => {
        if (labelConfidenceMap.get(a) < labelConfidenceMap.get(b)) {
            return -1;
        }
        else if (labelConfidenceMap.get(a) > labelConfidenceMap.get(b)) {
            return 1;
        }
        else {
            return filteredLabelMap.get(b).length - filteredLabelMap.get(a).length;
        }
    };
    // We need to sort in ascending order
    mediumConfidenceLabels.sort(sortLabels);
    lowConfidenceLabels.sort(sortLabels);
    // display the result in a table color-coded by confidence
    const table = new cli_table3_1.default({
        head: ['label', 'count', 'average confidence'],
        colAligns: ['left', 'right', 'right'],
    });
    // Since we're looking for the 10 worst, we take the first 10 after sorting
    mediumConfidenceLabels.slice(0, 10).forEach((label) => {
        table.push([
            label,
            filteredLabelMap.get(label).length.toString(),
            chalk_1.default.yellow(labelConfidenceMap.get(label).toFixed(2))
        ]);
    });
    lowConfidenceLabels.slice(0, 10).forEach((label) => {
        table.push([
            label,
            filteredLabelMap.get(label).length.toString(),
            chalk_1.default.red(labelConfidenceMap.get(label).toFixed(2))
        ]);
    });
    console.log(table.toString());
});
exports.default = program;
(async () => {
    program.parse(process.argv);
    if (process.argv.length === 2) {
        program.help();
    }
})();
