"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.PulseSessionService = exports.EVAL_RESULTS_TABLE_NAME = exports.OTEL_TRACES_TABLE_NAME = void 0;
const database_1 = require("../../database");
const drizzle_orm_1 = require("drizzle-orm");
const tables_1 = require("../../database/tables");
const clickhouseService_1 = require("./clickhouseService");
const clickhouse_1 = require("../../util/clickhouse");
const pulseEvaluationService_1 = require("./pulseEvaluationService");
const lodash_1 = __importDefault(require("lodash"));
const pulseService_1 = require("./pulseService");
const logger_1 = require("../../logger");
const get_upgrade_meta_1 = require("../../util/get-upgrade-meta");
const decimal_js_1 = __importDefault(require("decimal.js"));
const callControlPlaneApi_1 = require("../../util/callContolPlane/callControlPlaneApi");
const restriction_message_1 = require("../constants/restriction.message");
const evalResultsClickhouse_1 = require("./evalResultsClickhouse");
exports.OTEL_TRACES_TABLE_NAME = 'otel_traces';
exports.EVAL_RESULTS_TABLE_NAME = 'eval_results';
class PulseSessionService {
    constructor() {
        this.packSpansByTraces = (tracesData) => {
            const pulseService = new pulseService_1.PulseService();
            const traces = {};
            lodash_1.default.forEach(tracesData, (trace) => {
                if (trace?.traceId) {
                    traces[trace?.traceId] = {
                        isEvaluationDone: trace?.functionalComponentPassCount > 0 || trace?.functionalComponentFailCount > 0 || trace?.securityComponentPassCount > 0 || trace?.securityComponentFailCount > 0,
                        date: trace?.spans?.[0]?.Timestamp,
                        functionalComponentPassCount: trace?.functionalComponentPassCount ?? 0,
                        functionalComponentFailCount: trace?.functionalComponentFailCount ?? 0,
                        securityComponentPassCount: trace?.securityComponentPassCount ?? 0,
                        securityComponentFailCount: trace?.securityComponentFailCount ?? 0,
                        spans: trace?.spans ? trace?.spans?.map(pulseService.parseSpanResult) : [],
                        count: trace?.spans?.length ?? 0,
                    };
                }
            });
            const sortedEntries = Object.entries(traces).sort(([, a], [, b]) => {
                const dateA = new Date(a.date || 0).getTime();
                const dateB = new Date(b.date || 0).getTime();
                return dateB - dateA;
            });
            const sortedTraces = {};
            for (const [key, value] of sortedEntries) {
                sortedTraces[key] = value;
            }
            return sortedTraces;
        };
        this.logger = logger_1.logger;
    }
    async getSessions(params) {
        const { limit = 10, page = 1, applicationIds } = params;
        const offset = clickhouse_1.ClickHouseHelper.safeConvertToNumber(limit) *
            (clickhouse_1.ClickHouseHelper.safeConvertToNumber(page) - 1);
        const sortingType = params?.sorting?.type ?? '';
        const sortingDirection = params?.sorting?.direction ?? 'desc';
        const countQuery = `
SELECT uniq(SpanAttributes['session.id']) AS total	FROM ${exports.OTEL_TRACES_TABLE_NAME} 
WHERE ${clickhouse_1.ClickHouseHelper.getFilterWhereCondition(params, true)}
AND ResourceAttributes['service.name'] IN (${applicationIds?.map((id) => `'${id}'`).join(',')})
AND SpanAttributes['session.id'] IS NOT NULL
AND SpanAttributes['session.id'] IS NOT NULL
AND SpanAttributes['session.id'] != ''
AND SpanAttributes['session.id'] != 'null'
`;
        const countResult = await clickhouseService_1.clickHouseService.query(countQuery);
        const sessionQuery = `  
WITH parentSpans AS (
  SELECT
    TraceId,
    argMin(ifNull(Duration, 0), Timestamp) AS parentDuration
  FROM ${exports.OTEL_TRACES_TABLE_NAME}
  WHERE ResourceAttributes['service.name'] IN (${applicationIds.map((id) => `'${id}'`).join(',')})
  AND ResourceAttributes IS NOT NULL
  AND SpanAttributes['session.id'] IS NOT NULL
  AND SpanAttributes['session.id'] != ''
  AND SpanAttributes['session.id'] != 'null'
  AND Timestamp >= '${params.timeLimit.start}' AND Timestamp <= '${params.timeLimit.end}'
  AND (ParentSpanId IS NULL OR ParentSpanId = '' OR ParentSpanId = 'null')
  GROUP BY TraceId
),
traceDurations AS (
  SELECT
    SpanAttributes['session.id'] AS sessionId,
    TraceId,
    if(
      ifNull(p.parentDuration, 0) != 0,
      p.parentDuration,
      sum(coalesce(Duration, 0))
    ) AS traceDuration
  FROM ${exports.OTEL_TRACES_TABLE_NAME}
  LEFT JOIN parentSpans AS p ON p.TraceId = ${exports.OTEL_TRACES_TABLE_NAME}.TraceId
  WHERE ResourceAttributes['service.name'] IN (${applicationIds.map((id) => `'${id}'`).join(',')})
  AND ResourceAttributes IS NOT NULL
  AND SpanAttributes['session.id'] IS NOT NULL
  AND SpanAttributes['session.id'] != ''
  AND SpanAttributes['session.id'] != 'null'
  AND Timestamp >= '${params.timeLimit.start}' AND Timestamp <= '${params.timeLimit.end}'
  GROUP BY sessionId, TraceId, p.parentDuration
),
sessionDurations AS (
  SELECT
    sessionId,
    sum(traceDuration) AS totalDuration
  FROM traceDurations
  GROUP BY sessionId
),
sessions AS (
  SELECT
    SpanAttributes['session.id'] AS sessionId,
    uniq(TraceId) AS totalTraces,
    min(Timestamp) AS createdAt,
    max(Timestamp) AS updatedAt,
    argMinIf(ResourceAttributes['service.name'], Timestamp, ResourceAttributes['service.name'] IS NOT NULL) AS applicationId
  FROM ${exports.OTEL_TRACES_TABLE_NAME}
  WHERE ResourceAttributes['service.name'] IN (${applicationIds.map((id) => `'${id}'`).join(',')})
  AND ResourceAttributes IS NOT NULL
  AND SpanAttributes['session.id'] IS NOT NULL
  AND SpanAttributes['session.id'] != ''
  AND SpanAttributes['session.id'] != 'null'
  AND Timestamp >= '${params.timeLimit.start}' AND Timestamp <= '${params.timeLimit.end}'
  GROUP BY sessionId
  ${params.sorting
            ? sortingType.includes('cost')
                ? `ORDER BY argMax(toFloat64OrZero(${sortingType}), toFloat64OrZero(${sortingType})) ${sortingDirection} `
                : sortingType.includes('tokens')
                    ? `ORDER BY argMax(toInt32OrZero(${sortingType}), toInt32OrZero(${sortingType})) ${sortingDirection} `
                    : `ORDER BY argMax(${sortingType}, ${sortingType}) ${sortingDirection} `
            : `ORDER BY createdAt desc `}
  LIMIT ${clickhouse_1.ClickHouseHelper.safeConvertToNumber(limit)} OFFSET ${clickhouse_1.ClickHouseHelper.safeConvertToNumber(offset)}
),
evals AS (
    ${(0, evalResultsClickhouse_1.buildSessionListEvalCteSqlDenorm)('SELECT sessionId FROM sessions', applicationIds?.map((id) => `'${id.replace(/'/g, "''")}'`).join(',') ?? "''")}
)
SELECT
  sessions.sessionId as sessionId,
  sessions.totalTraces as totalTraces,
  ifNull(sd.totalDuration, 0) AS totalDuration,
  sessions.createdAt as createdAt,
  sessions.updatedAt as updatedAt,
  sessions.applicationId as applicationId,
  ${(0, evalResultsClickhouse_1.totalEvaluationsFromComponentCountsSql)('e')} AS totalEvaluations,
  e.functionalComponentPassCount,
  e.functionalComponentFailCount,
  e.securityComponentPassCount,
  e.securityComponentFailCount
FROM sessions
LEFT JOIN sessionDurations AS sd ON sd.sessionId = sessions.sessionId
LEFT JOIN evals AS e ON e.sessionId = sessions.sessionId;
`;
        const result = await clickhouseService_1.clickHouseService.query(sessionQuery);
        const traceIdsToApplicationIds = result.reduce((acc, item) => {
            acc[item.sessionId] = item.applicationId;
            return acc;
        }, {});
        const db = await (0, database_1.getDb)();
        const applicationNames = await db
            .select()
            .from(tables_1.applicationsTable)
            .where((0, drizzle_orm_1.inArray)(tables_1.applicationsTable.id, Object.values(traceIdsToApplicationIds)));
        const applicationNamesMap = new Map(applicationNames.map((item) => [item.id, item.name]));
        result.forEach((item) => {
            item.applicationName = applicationNamesMap.get(traceIdsToApplicationIds[item.sessionId]);
        });
        // return this.packSpansBySessions(result);
        return {
            sessions: result,
            total: countResult?.[0]?.total ?? 0,
        };
    }
    async getSessionById(id, shouldPack = true) {
        const query = ` 
WITH sessions AS (
    SELECT DISTINCT TraceId as id,
           min(Timestamp) as createdAt
    FROM ${exports.OTEL_TRACES_TABLE_NAME}
    WHERE ResourceAttributes IS NOT NULL
      AND SpanAttributes['session.id'] IS NOT NULL
      AND SpanAttributes['session.id'] != ''
      AND SpanAttributes['session.id'] != 'null'
      AND SpanAttributes['session.id'] = '${id}'
    GROUP BY TraceId
    ORDER BY createdAt ASC
),
evals AS (
    SELECT
        traceId,
        countIf(is_redteam = 0 AND overall_pass)       AS functionalPassCount,
        countIf(is_redteam = 0 AND NOT overall_pass)   AS functionalFailCount,
        countIf(is_redteam = 1 AND overall_pass)       AS securityPassCount,
        countIf(is_redteam = 1 AND NOT overall_pass)   AS securityFailCount,
        sumIf(
            arrayCount(x -> JSONExtractBool(x, 'pass'),
                       component_objects),
            is_redteam = 0
        )                                              AS functionalComponentPassCount,
        sumIf(
            arrayCount(x -> NOT JSONExtractBool(x, 'pass'),
                       component_objects),
            is_redteam = 0
        )                                              AS functionalComponentFailCount,
        sumIf(
            arrayCount(x -> JSONExtractBool(x, 'pass'),
                       component_objects),
            is_redteam = 1
        )                                              AS securityComponentPassCount,
        sumIf(
            arrayCount(x -> NOT JSONExtractBool(x, 'pass'),
                       component_objects),
            is_redteam = 1
        )                                              AS securityComponentFailCount
    FROM (
        SELECT
            traceId,
            is_redteam,
            JSONExtractBool(CAST(grading_result AS String), 'pass') AS overall_pass,
            coalesce(
                JSONExtract(CAST(grading_result AS String),
                            'componentResults',
                            'Array(String)'),
                []
            ) AS component_objects
        FROM ${exports.EVAL_RESULTS_TABLE_NAME}
        WHERE grading_result IS NOT NULL
        AND trace_id IN (SELECT id FROM sessions)
    ) AS e
    GROUP BY traceId
),
allSpans AS (
    SELECT 
        TraceId,
        arrayMap(
           span -> map(
               'Timestamp', toString(span.1),
               'TraceId', span.2,
               'SpanId', span.3,
               'ParentSpanId', span.4,
               'TraceState', span.5,
               'SpanName', span.6,
               'SpanKind', span.7,
               'ServiceName', span.8,
               'ResourceAttributes', toJSONString(span.9),
               'ScopeName', span.10,
               'ScopeVersion', span.11,
               'SpanAttributes', toJSONString(span.12),
               'Duration', toString(span.13),
               'StatusCode', span.14,
               'StatusMessage', span.15,
               'Events.Timestamp', toJSONString(span.16),
               'Events.Name', toJSONString(span.17),
               'Events.Attributes', toJSONString(span.18)
           ),
           groupArray((
               t.Timestamp,
               t.TraceId,
               t.SpanId,
               t.ParentSpanId,
               t.TraceState,
               t.SpanName,
               t.SpanKind,
               t.ServiceName,
               t.ResourceAttributes,
               t.ScopeName,
               t.ScopeVersion,
               t.SpanAttributes,
               t.Duration,
               t.StatusCode,
               t.StatusMessage,
               t.Events.Timestamp,
               t.Events.Name,
               t.Events.Attributes
           ))
       ) AS spans
    FROM ${exports.OTEL_TRACES_TABLE_NAME} t
    WHERE TraceId IN (SELECT id FROM sessions)
    GROUP BY TraceId
)
SELECT 
    t.TraceId as traceId,
    t.spans,
    e.functionalPassCount,
    e.functionalFailCount,
    e.securityPassCount,
    e.securityFailCount,
    e.functionalComponentPassCount,
    e.functionalComponentFailCount,
    e.securityComponentPassCount,
    e.securityComponentFailCount
FROM allSpans t
LEFT JOIN evals e ON t.TraceId = e.traceId
ORDER BY traceId;
    `;
        const result = await clickhouseService_1.clickHouseService.query(query);
        return this.packSpansByTraces(result);
    }
    async getSessionSpans(id) {
        const query = ` 
WITH traces AS (
    SELECT DISTINCT TraceId as id,
           min(Timestamp) as createdAt
    FROM ${exports.OTEL_TRACES_TABLE_NAME}
    WHERE ResourceAttributes IS NOT NULL
      AND SpanAttributes['session.id'] IS NOT NULL
      AND SpanAttributes['session.id'] != ''
      AND SpanAttributes['session.id'] != 'null'
      AND SpanAttributes['session.id'] = '${id}'
    GROUP BY TraceId
    ORDER BY createdAt ASC
)
SELECT * FROM ${exports.OTEL_TRACES_TABLE_NAME} WHERE TraceId IN (SELECT id FROM traces);
`;
        const result = await clickhouseService_1.clickHouseService.query(query);
        const pulseService = new pulseService_1.PulseService();
        return (result ?? [])?.map(pulseService.parseSpanResult);
    }
    async evaluationSessionById({ id, projectId, 
    // applicationId,
    orgId, }) {
        const spans = await this.getSessionSpans(id);
        if (!lodash_1.default.isArray(spans) || spans.length === 0) {
            throw new Error('Session not found');
        }
        const { restrictionsData, wallet, isStarter, thresholds, notifications, usages, isEnterprise, } = await (0, get_upgrade_meta_1.getUpgradeMetaByOrgId)(orgId, {
            fetchMetricUsage: true,
        });
        const skipMap = new Map();
        const renewalNotificationMap = new Map();
        const usageNotificationMap = new Map();
        const exhaustedNotificationMap = new Map();
        if (!isEnterprise
            && restrictionsData?.metrics > 0
            && new decimal_js_1.default(usages?.metric ?? 0).gte(restrictionsData?.metrics ?? 0)
            && !exhaustedNotificationMap?.get(orgId)?.metric) {
            if (isStarter) {
                if (notifications?.canSendMetricUsageExceeded) {
                    await (0, callControlPlaneApi_1.callControlPlane)(`/trusys-emails/limit-reached`, {
                        method: 'POST',
                        params: {
                            orgId: orgId,
                            type: 'Metric',
                            currentUsage: usages?.metric ?? 0,
                        },
                    }, '');
                }
                throw new Error(restriction_message_1.ERROR_RESTRICTION_MESSAGE.metrics.title);
            }
            if (!isStarter
                && notifications?.canSendMetricRenewal) {
                if (renewalNotificationMap) {
                    renewalNotificationMap.set(orgId, {
                        spans: renewalNotificationMap?.get(orgId)?.spans ?? false,
                        metric: true,
                    });
                }
                await (0, callControlPlaneApi_1.callControlPlane)(`/trusys-emails/limit-renewal`, {
                    method: 'POST',
                    params: {
                        orgId: orgId,
                        type: 'Metric',
                        currentUsage: usages?.metric ?? 0,
                    },
                }, '');
            }
        }
        if (!isEnterprise
            && new decimal_js_1.default(wallet?.credits ?? 0).lte(0)) {
            if (notifications?.canSendWalletUsageExceeded) {
                await (0, callControlPlaneApi_1.callControlPlane)(`/trusys-emails/limit-reached`, {
                    method: 'POST',
                    params: {
                        orgId: orgId,
                        type: 'Wallet',
                        currentUsage: wallet?.credits ?? 0,
                    },
                }, '');
            }
            throw new Error(restriction_message_1.ERROR_RESTRICTION_MESSAGE.lowCredits.title);
        }
        else if (!isEnterprise
            && new decimal_js_1.default(wallet?.credits ?? 0).gt(0)
            && new decimal_js_1.default(wallet?.credits ?? 0).lte(thresholds?.walletNotificationThreshold ?? 0)
            && notifications?.canSendLowWalletBalance) {
            if (usageNotificationMap) {
                usageNotificationMap.set(orgId, {
                    spans: usageNotificationMap?.get(orgId)?.spans ?? false,
                    metric: usageNotificationMap?.get(orgId)?.metric ?? false,
                    wallet: true,
                });
            }
            await (0, callControlPlaneApi_1.callControlPlane)(`/trusys-emails/limit-reaching`, {
                method: 'POST',
                params: {
                    orgId: orgId,
                    type: 'Wallet',
                    currentUsage: wallet?.credits ?? 0,
                },
            }, '');
        }
        let applicationId = '';
        for (const span of spans) {
            if (lodash_1.default.isEmpty(span.resourceAttributes?.['service.name']) === false) {
                applicationId = span.resourceAttributes['service.name'];
                break;
            }
        }
        const db = await (0, database_1.getDb)();
        const [application] = await db
            .select()
            .from(tables_1.applicationsTable)
            .where((0, drizzle_orm_1.eq)(tables_1.applicationsTable.id, applicationId)).limit(1);
        if (lodash_1.default.isEmpty(application)) {
            throw new Error(`Application not found for applicationId ${applicationId}`);
        }
        const [monitoringConfig] = await db
            .select()
            .from(tables_1.pulseMonitoringTable)
            .where((0, drizzle_orm_1.eq)(tables_1.pulseMonitoringTable.applicationId, applicationId)).limit(1);
        if (lodash_1.default.isEmpty(monitoringConfig)) {
            throw new Error(`Monitoring config not found for applicationId ${applicationId}`);
        }
        const evaluationService = new pulseEvaluationService_1.PulseEvaluationService();
        const evaluations = await evaluationService.runEvaluation({
            sessionId: id,
            spans,
            applicationId,
            projectId,
            monitoringConfig,
            application,
            workspaceId: orgId,
            skipMap,
            orgRestrictions: restrictionsData,
            renewalNotificationMap,
            usageNotificationMap,
            exhaustedNotificationMap,
        });
        return evaluations;
    }
    async getEvaluation(params) {
        const { id, applicationId, orgId } = params;
        this.logger.info(`Getting evaluation for session ${id}`);
        const { upgradeMetadata, restrictionsData, isEnterprise, } = await (0, get_upgrade_meta_1.getUpgradeMetaByOrgId)(orgId);
        if (!isEnterprise && restrictionsData?.spans > 0) {
            const count = await pulseService_1.pulseService.getSpansUsage({
                timeLimit: pulseService_1.pulseService.getTimeLimit(restrictionsData?.spans_timespan ?? 'lifetime'),
                orgId: orgId,
            });
            if (count !== 0 && new decimal_js_1.default(count).gte(restrictionsData?.spans ?? 0)) {
                upgradeMetadata.upgradeRequired = true;
                upgradeMetadata.topUpRequired = false;
                upgradeMetadata.upgradeTitle = restriction_message_1.RESTRICTION_MESSAGE.spans.title;
                upgradeMetadata.upgradeMessage = restriction_message_1.RESTRICTION_MESSAGE.spans.message;
                upgradeMetadata.currentPlan = restrictionsData?.currentPlan ?? '';
                upgradeMetadata.currentPlanId = restrictionsData?.currentPlanId ?? '';
            }
        }
        // Query ClickHouse for this traceId
        const query = `SELECT * FROM eval_results WHERE sessionId = '${id}' AND application_id = '${applicationId}'`;
        const results = await clickhouseService_1.clickHouseService.query(query);
        // Parse results
        const securityTestRuns = [];
        const testRuns = [];
        let prompt = null;
        let response = null;
        let hasChat = false;
        let hasAudio = false;
        let hasError = false;
        let errorText = '';
        let hasRunEvaluation = false;
        for (const result of results) {
            if (result.is_redteam === 1) {
                securityTestRuns.push(result);
            }
            else {
                testRuns.push(result);
            }
            if (result.SpanAttributes?.['gen_ai.output.type'] === 'chat') {
                hasChat = true;
            }
            if (result.SpanAttributes?.['gen_ai.output.type'] === 'audio') {
                hasAudio = true;
            }
            if (prompt === null
                && result?.prompt) {
                prompt = result?.prompt;
            }
            if (response === null
                && result?.response) {
                response = result?.response;
            }
            if (result?.failure_reason === 2) {
                hasError = true;
                errorText = 'Internal Server Error';
                result.error = 'Internal Server Error';
            }
            if (!hasRunEvaluation && result?.failure_reason !== 2) {
                hasRunEvaluation = true;
            }
        }
        for (const result of results) {
            if (hasAudio && result?.SpanAttributes?.['gen_ai.output.type'] === 'audio') {
                result.prompt = {};
                result.response = {};
            }
        }
        return {
            securityTestRuns,
            testRuns,
            hasError,
            errorText,
            hasChat,
            hasAudio,
            prompt,
            response,
            hasRunEvaluation,
            upgradeMetadata,
        };
    }
}
exports.PulseSessionService = PulseSessionService;
//# sourceMappingURL=pulseSessionService.js.map