"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.runClickhouseTextToSql = void 0;
const drizzle_orm_1 = require("drizzle-orm");
const database_1 = require("../../../../database");
const tables_1 = require("../../../../database/tables");
const logger_1 = __importDefault(require("../../../../logger"));
const clickhouseService_1 = require("../../clickhouseService");
const pulseService_1 = require("../../pulseService");
const argusEvalContext_1 = require("../utils/argusEvalContext");
const clickhouseTextToSqlSqlFixes_1 = require("../utils/clickhouseTextToSqlSqlFixes");
const langChainFromDefaultProvider_1 = require("../utils/langChainFromDefaultProvider");
const systemPrompt_1 = require("../utils/systemPrompt");
const guardrailReferenceTablesClickhouse_1 = require("./guardrailReferenceTablesClickhouse");
const redteamReferenceTablesClickhouse_1 = require("./redteamReferenceTablesClickhouse");
const GUARDRAIL_RESULTS_TABLE_NAME = 'guardrail_results';
const CLICKHOUSE_APPLICATIONS_TABLE_NAME = 'applications';
const describeTable = async (tableName) => {
    try {
        const rows = await clickhouseService_1.clickHouseService.query(`DESCRIBE TABLE ${tableName}`);
        if (!rows || rows.length === 0) {
            return `TABLE ${tableName} (no columns found)`;
        }
        const cols = rows.map((r) => `${r.name} ${r.type}`).join(', ');
        return `TABLE ${tableName} (${cols})`;
    }
    catch (error) {
        logger_1.default.warn(`Failed to describe ClickHouse table ${tableName}: ${error}`);
        return `TABLE ${tableName} (schema unavailable due to error)`;
    }
};
const getActiveApplicationIdsForProject = async (projectId) => {
    if (!projectId) {
        return [];
    }
    try {
        const db = await (0, database_1.getDb)();
        const rows = await db
            .select({ id: tables_1.applicationsTable.id, isArchived: tables_1.applicationsTable.is_archived })
            .from(tables_1.applicationsTable)
            .where((0, drizzle_orm_1.and)((0, drizzle_orm_1.eq)(tables_1.applicationsTable.projectId, projectId), (0, drizzle_orm_1.eq)(tables_1.applicationsTable.is_archived, false)));
        return Array.isArray(rows) ? rows.map((row) => row.id).filter(Boolean) : [];
    }
    catch (error) {
        logger_1.default.warn(`Failed to load active applicationIds for projectId=${projectId}: ${error}`);
        return [];
    }
};
function getTablesToDescribeForDomain(flags) {
    const { isTruguardDomain, isTruevalDomain, isTruscoutDomain } = flags;
    if (isTruguardDomain) {
        return new Set([
            GUARDRAIL_RESULTS_TABLE_NAME,
            CLICKHOUSE_APPLICATIONS_TABLE_NAME,
            guardrailReferenceTablesClickhouse_1.CLICKHOUSE_GUARDRAIL_VALIDATORS_TABLE,
            guardrailReferenceTablesClickhouse_1.CLICKHOUSE_GUARDRAIL_VALIDATOR_CATEGORIES_TABLE,
            guardrailReferenceTablesClickhouse_1.CLICKHOUSE_GUARDRAIL_ACTIONS_TABLE,
        ]);
    }
    if (isTruevalDomain) {
        return new Set([pulseService_1.EVAL_RESULTS_TABLE_NAME, CLICKHOUSE_APPLICATIONS_TABLE_NAME]);
    }
    if (isTruscoutDomain) {
        return new Set([
            pulseService_1.EVAL_RESULTS_TABLE_NAME,
            CLICKHOUSE_APPLICATIONS_TABLE_NAME,
            redteamReferenceTablesClickhouse_1.CLICKHOUSE_REDTEAM_PLUGINS_TABLE,
            redteamReferenceTablesClickhouse_1.CLICKHOUSE_REDTEAM_STRATEGIES_TABLE,
        ]);
    }
    return new Set([
        pulseService_1.OTEL_TRACES_TABLE_NAME,
        pulseService_1.EVAL_RESULTS_TABLE_NAME,
        GUARDRAIL_RESULTS_TABLE_NAME,
        CLICKHOUSE_APPLICATIONS_TABLE_NAME,
        redteamReferenceTablesClickhouse_1.CLICKHOUSE_REDTEAM_PLUGINS_TABLE,
        redteamReferenceTablesClickhouse_1.CLICKHOUSE_REDTEAM_STRATEGIES_TABLE,
        guardrailReferenceTablesClickhouse_1.CLICKHOUSE_GUARDRAIL_VALIDATORS_TABLE,
        guardrailReferenceTablesClickhouse_1.CLICKHOUSE_GUARDRAIL_VALIDATOR_CATEGORIES_TABLE,
        guardrailReferenceTablesClickhouse_1.CLICKHOUSE_GUARDRAIL_ACTIONS_TABLE,
    ]);
}
function getSqlSystemPromptForDomain(flags) {
    const { isTruguardDomain, isTruevalDomain, isTruscoutDomain } = flags;
    if (isTruguardDomain) {
        return systemPrompt_1.truguardClickhouseTextToSqlSystemPrompt;
    }
    if (isTruevalDomain) {
        return systemPrompt_1.truevalClickhouseTextToSqlSystemPrompt;
    }
    if (isTruscoutDomain) {
        return systemPrompt_1.truscoutClickhouseTextToSqlSystemPrompt;
    }
    return systemPrompt_1.clickhouseTextToSqlSystemPrompt;
}
function buildContextBlock(contextIds, flags) {
    const { isTruguardDomain, isEvalOnlyDomain, isTestRunEvalUi } = flags;
    const contextLines = [];
    if (contextIds?.traceId && !isTruguardDomain && !isEvalOnlyDomain) {
        contextLines.push(`traceId: ${contextIds.traceId}`);
    }
    if (contextIds?.spanId && !isTruguardDomain && !isEvalOnlyDomain) {
        contextLines.push(`spanId: ${contextIds.spanId}`);
    }
    if (contextIds?.sessionId && !isTruguardDomain && !isEvalOnlyDomain) {
        contextLines.push(`sessionId: ${contextIds.sessionId}`);
    }
    if (contextIds?.applicationId) {
        if (isTruguardDomain) {
            contextLines.push(`application_id (guardrail_results / applications): ${contextIds.applicationId}`);
        }
        else if (isEvalOnlyDomain) {
            contextLines.push(`application_id (eval_results / applications): ${contextIds.applicationId}`);
        }
        else {
            contextLines.push(`ServiceName: ${contextIds.applicationId}`);
        }
    }
    if (contextIds?.testRunId) {
        contextLines.push(`test_run_id: ${contextIds.testRunId}`);
        if (isTestRunEvalUi || isEvalOnlyDomain) {
            contextLines.push('Scope eval_results with this test_run_id. Do NOT use trace_id, TraceId, or otel_traces for this id.');
        }
    }
    else if (isTestRunEvalUi) {
        contextLines.push('Test-run UI: any UUID in the user question is test_run_id (not trace_id). Filter eval_results.test_run_id.');
    }
    if (contextIds?.evaluationId && !isTruguardDomain && !isEvalOnlyDomain) {
        contextLines.push(`evaluationId: ${contextIds.evaluationId}`);
        contextLines.push(`For eval_results: if SCHEMA includes test_run_id and the user/context implies one test run / many result rows, treat this id as test_run_id (same concept whether the UI/user says functional eval/run id, evaluation id, security eval/run id, or redteam id); otherwise use eval_id or other SCHEMA columns as appropriate.`);
    }
    if (contextIds?.evaluationId && isEvalOnlyDomain) {
        contextLines.push(`evaluationId (eval_id or related id from context): ${contextIds.evaluationId}`);
    }
    if (contextLines.length === 0) {
        return '';
    }
    const contextBlockParts = [`CONTEXT_IDS:\n- ${contextLines.join('\n- ')}`];
    return `\n${contextBlockParts.join('\n\n')}\n`;
}
const runClickhouseTextToSql = async (input) => {
    const { question, maxRows = 500, contextIds, projectId, domain, contextSource } = input;
    const domainLower = String(domain || '').toLowerCase();
    const isTruguardDomain = domainLower === 'truguard';
    const isTruevalDomain = domainLower === 'trueval';
    const isTruscoutDomain = domainLower === 'truscout';
    const isEvalOnlyDomain = isTruevalDomain || isTruscoutDomain;
    if (!question || typeof question !== 'string') {
        throw new Error('question is required and must be a non-empty string');
    }
    const activeApplicationIds = await getActiveApplicationIdsForProject(projectId);
    if (!isTruguardDomain && !isTruevalDomain) {
        try {
            await (0, redteamReferenceTablesClickhouse_1.ensureRedteamPluginsAndStrategiesTables)();
        }
        catch (error) {
            logger_1.default.warn(`Could not ensure ClickHouse red team reference tables (${redteamReferenceTablesClickhouse_1.CLICKHOUSE_REDTEAM_PLUGINS_TABLE}, ${redteamReferenceTablesClickhouse_1.CLICKHOUSE_REDTEAM_STRATEGIES_TABLE}): ${error instanceof Error ? error.message : String(error)}`);
        }
    }
    if (!isTruguardDomain && !isEvalOnlyDomain) {
        try {
            await (0, guardrailReferenceTablesClickhouse_1.ensureGuardrailReferenceTables)();
        }
        catch (error) {
            logger_1.default.warn(`Could not ensure ClickHouse guardrail reference tables (${guardrailReferenceTablesClickhouse_1.CLICKHOUSE_GUARDRAIL_VALIDATORS_TABLE}, ${guardrailReferenceTablesClickhouse_1.CLICKHOUSE_GUARDRAIL_VALIDATOR_CATEGORIES_TABLE}, ${guardrailReferenceTablesClickhouse_1.CLICKHOUSE_GUARDRAIL_ACTIONS_TABLE}): ${error instanceof Error ? error.message : String(error)}`);
        }
    }
    const domainFlags = { isTruguardDomain, isTruevalDomain, isTruscoutDomain };
    const tablesToDescribe = getTablesToDescribeForDomain(domainFlags);
    const schemaParts = [];
    for (const table of tablesToDescribe) {
        const desc = await describeTable(table);
        schemaParts.push(desc);
    }
    const schemaText = schemaParts.join('\n\n');
    const llm = (0, langChainFromDefaultProvider_1.getLangChainFromDefaultProvider)();
    const isTestRunEvalUi = (0, argusEvalContext_1.isTestRunEvalContextSource)(contextSource);
    const contextBlock = buildContextBlock(contextIds, {
        isTruguardDomain,
        isEvalOnlyDomain,
        isTestRunEvalUi,
    });
    const projectScopeBlock = projectId && String(projectId).trim()
        ? `\nPROJECT_SCOPE:\n- project_id: ${String(projectId).trim()}\n- When querying the applications table, always filter by this project_id.\n`
        : '';
    const userPrompt = (0, systemPrompt_1.buildClickhouseTextToSqlUserPrompt)({
        schemaText,
        question,
        contextBlock: `${contextBlock}${projectScopeBlock}`,
        isTruguardDomain,
        isTruevalDomain,
        isTruscoutDomain,
    });
    const sqlSystemPrompt = getSqlSystemPromptForDomain(domainFlags);
    let llmResponse;
    try {
        llmResponse = await llm.invoke([
            { role: 'system', content: sqlSystemPrompt },
            { role: 'user', content: userPrompt },
        ]);
    }
    catch (error) {
        const rawMessage = error?.response?.data?.error?.message ??
            error?.message ??
            String(error);
        logger_1.default.error(`ClickHouse text-to-SQL LLM error: ${rawMessage}`);
        throw new Error('The AI service encountered an error while generating the ClickHouse query. Please try again later or contact your administrator.');
    }
    let sql = String(llmResponse.content || '').trim();
    if (!sql) {
        throw new Error('LLM returned an empty SQL string');
    }
    sql = sql.replace(/;+$/g, '').trim();
    if (!/^select/i.test(sql)) {
        throw new Error(`Generated SQL is not a SELECT statement: ${sql}`);
    }
    if (/(insert|update|delete|alter|drop|truncate)\s/i.test(sql)) {
        throw new Error(`Generated SQL contains forbidden operation: ${sql}`);
    }
    let finalSql = sql;
    if (!/\blimit\s+\d+/i.test(sql)) {
        finalSql = `${sql} LIMIT ${maxRows}`;
    }
    finalSql = (0, clickhouseTextToSqlSqlFixes_1.postProcessGeneratedClickhouseSql)(finalSql, {
        activeApplicationIds,
        isEvalOnlyDomain,
        projectId,
        scopedTestRunId: typeof contextIds?.testRunId === 'string' ? contextIds.testRunId.trim() : undefined,
    });
    const promptChars = (sqlSystemPrompt?.length || 0) + (userPrompt?.length || 0);
    const completionChars = finalSql.length;
    const promptTokens = Math.round(promptChars / 4);
    const completionTokens = Math.round(completionChars / 4);
    const totalTokens = promptTokens + completionTokens;
    let rows = [];
    try {
        rows = await clickhouseService_1.clickHouseService.query(finalSql);
    }
    catch (error) {
        logger_1.default.error(`ClickHouse text-to-sql query failed. SQL="${finalSql}". Error: ${error instanceof Error ? error.message : String(error)}`);
        throw new Error('Unable to fetch details from ClickHouse at the moment. Please try again later.');
    }
    return {
        sql: finalSql,
        rowCount: rows.length,
        rows,
        tokenUsage: {
            total: totalTokens,
            prompt: promptTokens,
            completion: completionTokens,
        },
    };
};
exports.runClickhouseTextToSql = runClickhouseTextToSql;
//# sourceMappingURL=clickhouseTextToSqlTool.js.map