"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 () {
    var ownKeys = function(o) {
        ownKeys = Object.getOwnPropertyNames || function (o) {
            var ar = [];
            for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
            return ar;
        };
        return ownKeys(o);
    };
    return function (mod) {
        if (mod && mod.__esModule) return mod;
        var result = {};
        if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
        __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 });
exports.DEFAULT_TURN_COUNT = exports.ArgsSchema = void 0;
exports.doTargetPurposeDiscovery = doTargetPurposeDiscovery;
exports.mergePurposes = mergePurposes;
exports.discoverCommand = discoverCommand;
const chalk_1 = __importDefault(require("chalk"));
const cli_progress_1 = __importDefault(require("cli-progress"));
const commander_1 = require("commander");
const dedent_1 = __importDefault(require("dedent"));
const fs = __importStar(require("fs"));
const js_yaml_1 = __importDefault(require("js-yaml"));
const path_1 = __importDefault(require("path"));
const zod_1 = require("zod");
const constants_1 = require("../../constants");
const fetch_1 = require("../../fetch");
const accounts_1 = require("../../globalConfig/accounts");
const cloud_1 = require("../../globalConfig/cloud");
const logger_1 = __importDefault(require("../../logger"));
const providers_1 = require("../../providers");
const telemetry_1 = __importDefault(require("../../telemetry"));
const cloud_2 = require("../../util/cloud");
const load_1 = require("../../util/config/load");
const manage_1 = require("../../util/config/manage");
const invariant_1 = __importDefault(require("../../util/invariant"));
const constants_2 = require("../constants");
const remoteGeneration_1 = require("../remoteGeneration");
const remoteGeneration_2 = require("../remoteGeneration");
exports.ArgsSchema = zod_1.z
    .object({
    config: zod_1.z.string().optional(),
    output: zod_1.z.string().optional(),
    target: zod_1.z.string().optional(),
    preview: zod_1.z.boolean(),
    turns: zod_1.z.number().optional(),
    overwrite: zod_1.z.boolean(),
})
    // Config and target are mutually exclusive:
    .refine((data) => !(data.config && data.target), {
    message: 'Cannot specify both config and target!',
    path: ['config', 'target'],
})
    // Either config or target must be provided:
    .refine((data) => data.config || data.target, {
    message: 'Either config or target must be provided!',
    path: ['config', 'target'],
})
    // `output` and `preview` are mutually exclusive:
    .refine((data) => !(data.output && data.preview), {
    message: 'Cannot specify both output and preview!',
    path: ['output', 'preview'],
})
    // `overwrite` can only be used if `output` is provided:
    .refine((data) => !(data.overwrite && !data.output), {
    message: 'Cannot specify overwrite without output!',
    path: ['overwrite', 'output'],
})
    // if `preview` is false, `output` must be provided:
    .refine((data) => !(data.preview === false && !data.output), {
    message: 'If preview is false, output must be provided!',
    path: ['preview', 'output'],
});
// A larger turn count is more accurate (b/c more probes) but slower.
// TODO: Optimize this default to balance quality/runtime using the Discover eval.
// NOTE: Set to 5 because UI lacks ability to set the count.
exports.DEFAULT_TURN_COUNT = 5;
/**
 * Queries Cloud for the purpose-discovery logic, sends each logic to the target,
 * and summarizes the results.
 *
 * @param target - The target API provider.
 * @param maxTurns - The maximum number of turns to run the discovery process.
 * @returns The purpose of the target.
 */
async function doTargetPurposeDiscovery(target, maxTurns = exports.DEFAULT_TURN_COUNT) {
    const conversationHistory = [];
    let turnCounter = 0;
    const pbar = new cli_progress_1.default.SingleBar({
        format: `Discovering purpose {bar} {percentage}% | {value}${maxTurns ? '/{total}' : ''} turns`,
        barCompleteChar: '\u2588',
        barIncompleteChar: '\u2591',
        hideCursor: true,
    });
    pbar.start(maxTurns, turnCounter);
    try {
        for (turnCounter = 0; turnCounter < maxTurns; turnCounter++) {
            const res = await (0, fetch_1.fetchWithProxy)((0, remoteGeneration_1.getRemoteGenerationUrl)(), {
                body: JSON.stringify({
                    task: 'target-purpose-discovery',
                    conversationHistory,
                    maxTurns,
                    version: constants_1.VERSION,
                    email: (0, accounts_1.getUserEmail)(),
                }),
                headers: { 'Content-Type': 'application/json' },
                method: 'POST',
            });
            const { done, question, purpose } = (await res.json());
            if (done) {
                pbar.increment();
                pbar.stop();
                logger_1.default.info(`\nPurpose:\n\n${chalk_1.default.green(purpose)}\n`);
                return purpose;
            }
            else {
                if (!question) {
                    throw new Error(`Failed to discover purpose: ${res.statusText}`);
                }
                conversationHistory.push({ type: 'promptfoo', content: question });
            }
            // Call the target with the question:
            const response = await target.callApi(question);
            logger_1.default.debug(JSON.stringify({ question, output: response.output }, null, 2));
            conversationHistory.push({ type: 'target', content: response.output });
            pbar.increment();
        }
        // If we've exhausted all turns, assume the last purpose is available
        // This will only be reached if done=false for all turns
        const res = await (0, fetch_1.fetchWithProxy)((0, remoteGeneration_1.getRemoteGenerationUrl)(), {
            body: JSON.stringify({
                task: 'target-purpose-discovery',
                conversationHistory,
                maxTurns,
                version: constants_1.VERSION,
                email: (0, accounts_1.getUserEmail)(),
                forceReturn: true,
            }),
            headers: { 'Content-Type': 'application/json' },
            method: 'POST',
        });
        const { purpose } = (await res.json());
        if (purpose) {
            return purpose;
        }
        throw new Error('Failed to discover purpose after exhausting maximum turns');
    }
    finally {
        // Ensure progress bar is always stopped, even if an exception occurs
        if (pbar) {
            pbar.stop();
        }
    }
}
/**
 * For targets hosted on Cloud, save the purpose to the database.
 *
 * @param targetId - The target ID.
 * @param purpose - The purpose.
 * @returns The response from the database.
 */
async function saveCloudTargetPurpose(targetId, purpose) {
    (0, invariant_1.default)(cloud_1.cloudConfig.isEnabled(), 'Cloud config should have been enabled for a target to be provided');
    const url = `${cloud_1.cloudConfig.getApiHost()}/api/v1/providers/${targetId}`;
    logger_1.default.debug(`Saving purpose to ${url}`);
    const res = await (0, fetch_1.fetchWithProxy)(url, {
        method: 'PATCH',
        headers: {
            'Content-Type': 'application/json',
            Authorization: `Bearer ${cloud_1.cloudConfig.getApiKey()}`,
        },
        body: JSON.stringify({ applicationDescription: { purpose } }),
    });
    if (res.ok) {
        logger_1.default.info('Purpose updated');
    }
    else {
        logger_1.default.error(`Failed to save purpose to database: ${res.statusText}`);
    }
}
function mergePurposes(humanDefinedPurpose, discoveredPurpose) {
    return `${humanDefinedPurpose}\n\nDiscovered Purpose:\n\n${discoveredPurpose}`;
}
/**
 * Registers the `discover` command with the CLI.
 */
function discoverCommand(program) {
    program
        .command('discover')
        .description((0, dedent_1.default) `
        Automatically discover a target application's purpose, enhancing attack probe efficacy.

        If neither a config file nor a target ID is provided, the current working directory will be checked for a promptfooconfig.yaml file,
        and the target will be discovered from the first provider in that config.
      `)
        .option('-c, --config <path>', 'Path to `promptfooconfig.yaml` configuration file.')
        .option('-o, --output <path>', 'Path to output file. Discovered purpose will be appended to the file if it already exists.', constants_2.DEFAULT_OUTPUT_PATH)
        .option('--overwrite', 'Overwrite the existing purpose if it already exists.', false)
        .option('-t, --target <id>', 'UUID of a Cloud-defined target to run the discovery on')
        .option('--preview', 'Preview discovery results without writing to an output file', false)
        .addOption(new commander_1.Option('--turns <turns>', 'A maximum number of turns to run the discovery process. Lower is faster but less accurate.')
        .argParser(Number.parseInt)
        .default(exports.DEFAULT_TURN_COUNT))
        .action(async (rawArgs) => {
        // If preview is true and output is DEFAULT_OUTPUT_PATH, set output to undefined to satisfy
        // the schema. Defaults are defined within the `option` definitions to include them within the
        // help message; however the schema enforces combinations of options that are mutually exclusive,
        // such as `output` and `preview`.
        if (rawArgs.preview && rawArgs.output === constants_2.DEFAULT_OUTPUT_PATH) {
            rawArgs.output = undefined;
        }
        // Check that remote generation is enabled:
        if ((0, remoteGeneration_2.neverGenerateRemote)()) {
            logger_1.default.error((0, dedent_1.default) `
          Discovery relies on remote generation which is disabled.

          To enable remote generation, unset the PROMPTFOO_DISABLE_REDTEAM_REMOTE_GENERATION environment variable.
        `);
            process.exit(1);
        }
        // Validate the arguments:
        const { success, data: args, error } = exports.ArgsSchema.safeParse(rawArgs);
        if (!success) {
            logger_1.default.error('Invalid options:');
            error.issues.forEach((issue) => {
                logger_1.default.error(`  ${issue.path.join('.')}: ${issue.message}`);
            });
            process.exitCode = 1;
            return;
        }
        // Record telemetry:
        telemetry_1.default.record('command_used', {
            name: 'redteam discover',
        });
        let config = null;
        // Although the providers/targets property supports multiple values, Redteaming only supports
        // a single target at a time.
        let target = undefined;
        // Fallback to the default config path:
        const fallbackConfigPath = path_1.default.join(process.cwd(), 'promptfooconfig.yaml');
        // If user provides a config, read the target from it:
        if (args.config) {
            // Validate that the config is a valid path:
            if (!fs.existsSync(args.config)) {
                throw new Error(`Config not found at ${args.config}`);
            }
            config = await (0, load_1.readConfig)(args.config);
            if (!config) {
                throw new Error(`Config is invalid at ${args.config}`);
            }
            if (!config.providers) {
                throw new Error('Config must contain at least one target or provider');
            }
            const providers = await (0, providers_1.loadApiProviders)(config.providers);
            target = providers[0];
        }
        // If the target flag is provided, load it from Cloud:
        else if (args.target) {
            // Let the internal error handling bubble up:
            const providerOptions = await (0, cloud_2.getProviderFromCloud)(args.target);
            target = await (0, providers_1.loadApiProvider)(providerOptions.id, { options: providerOptions });
        }
        // Check the current working directory for a promptfooconfig.yaml file:
        else if (fs.existsSync(fallbackConfigPath)) {
            config = await (0, load_1.readConfig)(fallbackConfigPath);
            if (!config) {
                throw new Error(`Config is invalid at ${fallbackConfigPath}`);
            }
            if (!config.providers) {
                throw new Error('Config must contain at least one target or provider');
            }
            const providers = await (0, providers_1.loadApiProviders)(config.providers);
            target = providers[0];
            // Alert the user that we're using a config from the current working directory:
            logger_1.default.info(`Using config from ${chalk_1.default.italic(fallbackConfigPath)}`);
        }
        // At this point, we should have at least one target:
        (0, invariant_1.default)(target != undefined, 'An error occurred loading the target config');
        // Discover the purpose for the target:
        let purpose = undefined;
        try {
            purpose = await doTargetPurposeDiscovery(target, args.turns);
        }
        catch (error) {
            logger_1.default.error(`An unexpected error occurred during target discovery: ${error instanceof Error ? error.message : String(error)}\n${error instanceof Error ? error.stack : ''}`);
            process.exit(1);
        }
        // If not previewing, persist the purposes:
        if (!args.preview) {
            // Persist the purposes:
            if (args.target) {
                await saveCloudTargetPurpose(args.target, purpose);
            }
            else {
                (0, invariant_1.default)(config, 'Config is required');
                if (args.output === constants_2.DEFAULT_OUTPUT_PATH) {
                    args.output = path_1.default.relative(process.cwd(), constants_2.DEFAULT_OUTPUT_PATH);
                }
                logger_1.default.debug(`Writing purpose to ${args.output}`);
                if (fs.existsSync(args.output)) {
                    const existingYaml = js_yaml_1.default.load(fs.readFileSync(args.output, 'utf8'));
                    // Either append or overwrite the existing purpose:
                    const existingPurpose = existingYaml['redteam']?.purpose;
                    if (existingPurpose) {
                        if (args.overwrite) {
                            logger_1.default.warn((0, dedent_1.default) `
                  Output file already contains a value at \`redteam.purpose\`; overwriting it.
                `);
                        }
                        else {
                            logger_1.default.warn((0, dedent_1.default) `
                  Output file already contains a value at \`redteam.purpose\`; appending discovered purpose to it.

                  To overwrite the existing purpose, use the \`--overwrite\` flag.
                `);
                        }
                    }
                    existingYaml['redteam'] = {
                        ...(existingYaml['redteam'] || {}),
                        purpose: existingPurpose && !args.overwrite
                            ? mergePurposes(existingPurpose, purpose)
                            : purpose,
                    };
                    (0, manage_1.writePromptfooConfig)(existingYaml, args.output);
                }
                else {
                    // Create a new config file with the purpose.
                    (0, manage_1.writePromptfooConfig)({ redteam: { purpose } }, args.output);
                }
                logger_1.default.info(`\nPurpose written to ${chalk_1.default.italic(args.output)}`);
            }
        }
    });
}
//# sourceMappingURL=discover.js.map