"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.securityEvalService = exports.SecurityEvalService = void 0;
const decimal_js_1 = __importDefault(require("decimal.js"));
const drizzle_orm_1 = require("drizzle-orm");
const lodash_1 = __importDefault(require("lodash"));
const uuid_1 = require("uuid");
const cliState_1 = __importDefault(require("../../cliState"));
const constants_1 = require("../../constants");
const database_1 = require("../../database");
const tables_1 = require("../../database/tables");
const accounts_1 = require("../../globalConfig/accounts");
const logger_1 = __importDefault(require("../../logger"));
const shared_1 = require("../../redteam/shared");
const callControlPlaneApi_1 = require("../../util/callContolPlane/callControlPlaneApi");
const time_1 = require("../../util/time");
const restriction_message_1 = require("../constants/restriction.message");
const evalJobsCache_1 = require("../../util/evalJobsCache");
const get_upgrade_meta_1 = require("../../util/get-upgrade-meta");
const billing_logger_1 = __importDefault(require("./billing/billing-logger"));
const activityLogService_1 = require("./activityLogService");
const evalBilling_1 = require("./evaluations/evalBilling");
class SecurityEvalService {
    /**
     * Create a security evaluation with all validation, processing, and execution logic
     */
    async createSecurityEval(params) {
        try {
            const { config, force, verbose, delay, user, projectId, orgId, workspaceId, applicationIds, maxConcurrency, securityConfigId, isCustomProvider = false, token, apiKey, userId, userName, } = params;
            const db = await (0, database_1.getDb)();
            const { restrictionsData, isEnterprise, wallet, isStarter, usages, notifications, thresholds, bucketInfo } = await (0, get_upgrade_meta_1.getUpgradeMetaByOrgId)(orgId, {
                fetchMetricUsage: true,
            });
            // Check bucket charges and calculate extra buckets needed
            if (!isEnterprise
                && !isStarter
                && bucketInfo
                && bucketInfo.metric) {
                const metricBucketSize = bucketInfo.metric.bucketSize ?? 0;
                const metricBucketsCharged = bucketInfo.metric.bucketsCharged ?? 0;
                const metricFreeTierLimit = bucketInfo.metric.freeTierLimit ?? 0;
                const metricCurrentUsage = new decimal_js_1.default(usages?.metric ?? 0).toNumber();
                // Calculate if extra buckets are needed based on current usage
                if (metricBucketSize > 0 && metricCurrentUsage > metricFreeTierLimit) {
                    const usageBeyondFree = new decimal_js_1.default(metricCurrentUsage).minus(metricFreeTierLimit);
                    const alreadyCoveredUsage = new decimal_js_1.default(metricBucketsCharged).mul(metricBucketSize);
                    const uncoveredUsage = usageBeyondFree.minus(alreadyCoveredUsage);
                    const extraBucketsNeeded = uncoveredUsage.gt(0)
                        ? uncoveredUsage.div(metricBucketSize).ceil().toNumber()
                        : 0;
                    if (extraBucketsNeeded > 0) {
                        // Call API to charge bucket - calculates cost, checks wallet, deducts, and updates tracking
                        try {
                            const response = await (0, callControlPlaneApi_1.callControlPlane)(`/trusys/billing-plans/charge-bucket`, {
                                method: 'POST',
                                params: {
                                    orgId,
                                    bucketType: 'metrics',
                                    extraBucketsNeeded,
                                    currentUsage: metricCurrentUsage,
                                },
                            }, '');
                            if (response?.error) {
                                return {
                                    success: false,
                                    message: response.error,
                                    showToast: true,
                                    toastMetadata: {
                                        text: response.error,
                                        type: 'error',
                                        duration: 5000,
                                    },
                                };
                            }
                        }
                        catch (error) {
                            // API throws error if insufficient funds or other issues - halt execution
                            const errorMessage = error instanceof Error ? error.message : String(error);
                            return {
                                success: false,
                                message: errorMessage.includes('Insufficient')
                                    ? errorMessage
                                    : 'Failed to charge metric bucket. Please try again.',
                                showToast: true,
                                toastMetadata: {
                                    text: errorMessage.includes('Insufficient')
                                        ? errorMessage
                                        : 'Failed to charge metric bucket. Please try again.',
                                    type: 'error',
                                    duration: 5000,
                                },
                            };
                        }
                    }
                }
            }
            if (!isEnterprise && restrictionsData?.security_evals > 0) {
                const whereConditions = (0, time_1.createTimespanWhereConditions)({
                    table: tables_1.testRunsTable,
                    timeColumn: 'createdAt',
                    orgId,
                    timespan: restrictionsData?.security_evals_timespan || 'lifetime',
                    extraConditions: (0, drizzle_orm_1.and)((0, drizzle_orm_1.sql) `${tables_1.testRunsTable.type} = ${constants_1.TEST_RUN_TYPE.SECURITY}`, (0, drizzle_orm_1.sql) `${tables_1.testRunsTable.status} = ${constants_1.TEST_RUN_STATUS.COMPLETED}`),
                });
                const count = await db
                    .select({ count: (0, drizzle_orm_1.sql) `count(*)` })
                    .from(tables_1.testRunsTable)
                    .where(whereConditions)
                    .then((result) => Number(result[0]?.count ?? 0));
                if (count != 0
                    && new decimal_js_1.default(count).gt(0)
                    && new decimal_js_1.default(count).gte(restrictionsData?.security_evals ?? 0)) {
                    return {
                        success: false,
                        message: restriction_message_1.ERROR_RESTRICTION_MESSAGE.upgrade.title,
                        showToast: true,
                        toastMetadata: {
                            text: restriction_message_1.ERROR_RESTRICTION_MESSAGE.upgrade.message,
                            type: 'error',
                        },
                    };
                }
            }
            // block run for reaching the wallet limit
            if (!isEnterprise
                && new decimal_js_1.default(wallet?.credits ?? 0).lte(new decimal_js_1.default(0))) {
                if (notifications?.canSendWalletUsageExceeded) {
                    await (0, callControlPlaneApi_1.callControlPlane)(`/trusys-emails/limit-reached`, {
                        method: 'POST',
                        params: {
                            orgId,
                            type: 'Wallet',
                            currentUsage: wallet?.credits ?? 0,
                        },
                    }, '');
                }
                return {
                    success: false,
                    message: restriction_message_1.ERROR_RESTRICTION_MESSAGE.lowCredits.title,
                    showToast: true,
                    toastMetadata: {
                        text: restriction_message_1.ERROR_RESTRICTION_MESSAGE.lowCredits.message,
                        type: 'error',
                        duration: 5000,
                    },
                };
            }
            // send notification for low wallet balance
            if (!isEnterprise
                && new decimal_js_1.default(wallet?.credits ?? 0).lte(thresholds?.walletNotificationThreshold ?? 0)
                && notifications?.canSendLowWalletBalance) {
                await (0, callControlPlaneApi_1.callControlPlane)(`/trusys-emails/limit-reaching`, {
                    method: 'POST',
                    params: {
                        orgId,
                        type: 'Wallet',
                        currentUsage: wallet?.credits ?? 0,
                    },
                }, '');
            }
            // check for metric usage limit
            if (!isEnterprise
                && restrictionsData?.metrics > 0
                && new decimal_js_1.default(usages?.metric ?? 0).gte(restrictionsData?.metrics ?? 0)) {
                if (isStarter) {
                    if (notifications?.canSendMetricUsageExceeded) {
                        await (0, callControlPlaneApi_1.callControlPlane)(`/trusys-emails/limit-reached`, {
                            method: 'POST',
                            params: {
                                orgId,
                                type: 'Metric',
                                currentUsage: usages?.metric ?? 0,
                            },
                        }, '');
                    }
                    return {
                        success: false,
                        message: restriction_message_1.ERROR_RESTRICTION_MESSAGE.metrics.title,
                        showToast: true,
                        toastMetadata: {
                            text: restriction_message_1.ERROR_RESTRICTION_MESSAGE.metrics.message,
                            type: 'error',
                            duration: 5000,
                        },
                    };
                }
                else if (notifications?.canSendMetricRenewal) {
                    await (0, callControlPlaneApi_1.callControlPlane)(`/trusys-emails/limit-renewal`, {
                        method: 'POST',
                        params: {
                            orgId,
                            type: 'Metric',
                            currentUsage: usages?.metric ?? 0,
                        },
                    }, '');
                }
            }
            (0, accounts_1.setUserEmail)(user?.email ?? 'system@trusys.ai');
            const id = (0, uuid_1.v4)();
            const abortController = new AbortController();
            // Initialize job status with empty logs array
            await (0, evalJobsCache_1.setEvalJob)(id, {
                evalId: null,
                status: 'in-progress',
                progress: 0,
                total: 0,
                result: null,
                logs: [],
            });
            const testRunId = id;
            await db.transaction(async (tx) => {
                // Create test run record
                await tx.insert(tables_1.testRunsTable).values({
                    id: testRunId,
                    type: constants_1.TEST_RUN_TYPE.SECURITY,
                    status: constants_1.TEST_RUN_STATUS.PENDING,
                    projectId: projectId || '',
                    workspaceId: workspaceId || '',
                    applicationIds: applicationIds || [],
                });
                if (securityConfigId) {
                    const [securityConfig] = await db.select({ config: tables_1.securityConfigTable.config })
                        .from(tables_1.securityConfigTable)
                        .where((0, drizzle_orm_1.eq)(tables_1.securityConfigTable.id, securityConfigId))
                        .limit(1);
                    await tx.insert(tables_1.testRunsToSecurityConfigsTable).values({
                        testRunId,
                        securityConfigId,
                        config: securityConfig?.config == null
                            ? null
                            : securityConfig.config,
                    });
                }
            });
            // Set web UI mode
            cliState_1.default.webUI = true;
            // Run redteam in background
            (0, shared_1.doRedteamRun)({
                liveRedteamConfig: config,
                token,
                apiKey,
                applicationIds,
                force,
                verbose,
                delay: Number(delay || '0'),
                logCallback: async (message, stats) => {
                    const job = await (0, evalJobsCache_1.getEvalJob)(id);
                    if (job) {
                        const updates = {};
                        if (lodash_1.default.isEmpty(message) === false) {
                            updates.logs = [...(job.logs || []), message];
                        }
                        if (lodash_1.default.isEmpty(stats) === false) {
                            updates.progress = stats?.net ?? 0;
                            updates.total = stats?.total ?? 0;
                        }
                        if (Object.keys(updates).length > 0) {
                            await (0, evalJobsCache_1.updateEvalJob)(id, updates);
                        }
                    }
                },
                maxConcurrency: Number(maxConcurrency || '4'),
                abortSignal: abortController.signal,
                testRunId: id,
                isCustomProvider,
                projectId,
                workspaceId,
            })
                .then(async (evalResult) => {
                const summary = evalResult ? await evalResult.toEvaluateSummary() : null;
                const job = await (0, evalJobsCache_1.getEvalJob)(id);
                if (job) {
                    await (0, evalJobsCache_1.updateEvalJob)(id, {
                        status: 'complete',
                        result: summary,
                        evalId: evalResult?.id ?? null,
                    });
                    // Update test run status and link to eval
                    if (evalResult?.id) {
                        await db.transaction(async (tx) => {
                            // Check if the test run was aborted before updating status
                            const [testRun] = await tx
                                .select({ status: tables_1.testRunsTable.status })
                                .from(tables_1.testRunsTable)
                                .where((0, drizzle_orm_1.eq)(tables_1.testRunsTable.id, testRunId));
                            if (testRun?.status !== constants_1.TEST_RUN_STATUS.ABORTED) {
                                await tx
                                    .update(tables_1.testRunsTable)
                                    .set({ status: constants_1.TEST_RUN_STATUS.COMPLETED })
                                    .where((0, drizzle_orm_1.eq)(tables_1.testRunsTable.id, testRunId));
                            }
                            await tx
                                .insert(tables_1.testRunsToEvalsTable)
                                .values({ testRunId, evalId: evalResult.id })
                                .onConflictDoNothing()
                                .execute();
                            await tx.update(tables_1.evalResultsTable).set({ isRedteam: 1 }).where((0, drizzle_orm_1.eq)(tables_1.evalResultsTable.evalId, evalResult.id));
                        });
                    }
                    let totalMetricUsage = new decimal_js_1.default(0);
                    if (Array.isArray(evalResult?.results) && (evalResult?.results?.length ?? 0) > 0) {
                        // Aggregate and log security eval usage per model (including Deepgram duration)
                        await (0, evalBilling_1.aggregateAndLogEvalUsage)({
                            evalType: 'security',
                            workspaceId,
                            evalId: testRunId,
                            evalResults: evalResult?.results ?? [],
                            baseMetadata: {
                                projectId,
                                applicationIds,
                                testRunId,
                                securityConfigId,
                                evalId: evalResult?.id ?? null,
                            },
                        });
                        lodash_1.default.forEach(evalResult?.results, (result) => {
                            if (result?.testCase &&
                                Array.isArray(result?.testCase?.assert) &&
                                result?.testCase?.assert?.length > 0) {
                                totalMetricUsage = totalMetricUsage.add(new decimal_js_1.default(result?.testCase?.assert?.length ?? 0));
                                billing_logger_1.default.logPluginUsage({
                                    workspaceId,
                                    testRunId,
                                    cost: 0,
                                    tokenUsage: {},
                                    amount: result?.testCase?.assert?.length ?? 0,
                                    metadata: {
                                        subType: 'plugin',
                                        projectId,
                                        applicationIds,
                                        testRunId,
                                        securityConfigId,
                                        evalResultId: result.id,
                                        evalId: evalResult?.id ?? null,
                                    },
                                });
                            }
                        });
                    }
                    // Check bucket charges and calculate extra buckets needed
                    if (!isEnterprise
                        && !isStarter
                        && bucketInfo
                        && bucketInfo.metric) {
                        const metricBucketSize = bucketInfo.metric.bucketSize ?? 0;
                        const metricBucketsCharged = bucketInfo.metric.bucketsCharged ?? 0;
                        const metricFreeTierLimit = bucketInfo.metric.freeTierLimit ?? 0;
                        const metricCurrentUsage = new decimal_js_1.default(usages?.metric ?? 0).plus(totalMetricUsage).toNumber();
                        // Calculate if extra buckets are needed based on current usage
                        if (metricBucketSize > 0 && metricCurrentUsage > metricFreeTierLimit) {
                            const usageBeyondFree = new decimal_js_1.default(metricCurrentUsage).minus(metricFreeTierLimit);
                            const alreadyCoveredUsage = new decimal_js_1.default(metricBucketsCharged).mul(metricBucketSize);
                            const uncoveredUsage = usageBeyondFree.minus(alreadyCoveredUsage);
                            const extraBucketsNeeded = uncoveredUsage.gt(0)
                                ? uncoveredUsage.div(metricBucketSize).ceil().toNumber()
                                : 0;
                            if (extraBucketsNeeded > 0) {
                                // Call API to charge bucket - calculates cost, checks wallet, deducts, and updates tracking
                                try {
                                    (0, callControlPlaneApi_1.callControlPlane)(`/trusys/billing-plans/charge-bucket`, {
                                        method: 'POST',
                                        params: {
                                            orgId,
                                            bucketType: 'metrics',
                                            extraBucketsNeeded,
                                            currentUsage: metricCurrentUsage,
                                            allowNegative: true,
                                        },
                                        retryCount: 3,
                                    }, '');
                                }
                                catch (error) {
                                    console.error('Error charging metric bucket:', error);
                                }
                            }
                        }
                    }
                }
                cliState_1.default.webUI = false;
            })
                .catch(async (error) => {
                console.error('Error running redteam:', error);
                const job = await (0, evalJobsCache_1.getEvalJob)(id);
                if (job) {
                    await (0, evalJobsCache_1.updateEvalJob)(id, {
                        status: 'error',
                        logs: [...(job.logs || []), `Error: ${error.message}`],
                    });
                }
                cliState_1.default.webUI = false;
            });
            // Log activity for security evaluation
            try {
                const [orgName, projectName] = await Promise.all([
                    (0, activityLogService_1.fetchOrgName)(token, workspaceId),
                    (0, activityLogService_1.fetchProjectName)(token, workspaceId, projectId),
                ]);
                // Format test run ID for display (e.g., "Security Eval#testRunId")
                const testRunDisplayId = testRunId;
                if (token) {
                    await (0, activityLogService_1.logActivity)(token, {
                        type: constants_1.ActivityType.RUN_SECURITY_EVALUATION,
                        params: (0, activityLogService_1.buildActivityParams)({
                            test_run_id: { value: testRunDisplayId, _id: testRunId },
                            user_name: { value: userName || '', _id: userId || '' },
                            org_name: { value: orgName || workspaceId, _id: workspaceId },
                            project_name: { value: projectName || projectId, _id: projectId },
                            application_name: { value: config.targets?.[0]?.label ?? '' },
                        }),
                        project_id: projectId,
                        org_id: workspaceId,
                        created_by: userId || '',
                        updated_by: userId || '',
                    });
                }
            }
            catch (activityError) {
                logger_1.default.error(`Error logging activity for run security evaluation: ${activityError}`);
            }
            return { id, testRunId };
        }
        catch (error) {
            logger_1.default.error(`Error while security eval: ${error.message}`);
            throw error;
        }
    }
}
exports.SecurityEvalService = SecurityEvalService;
// Export singleton instance
exports.securityEvalService = new SecurityEvalService();
//# sourceMappingURL=securityEvalService.js.map