"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.updateTrace = exports.insertTrace = exports.updateStatusByIds = exports.isBlocklet = void 0;
const node_fs_1 = require("node:fs");
const node_path_1 = __importDefault(require("node:path"));
const api_1 = require("@opentelemetry/api");
const decimal_js_1 = __importDefault(require("decimal.js"));
const drizzle_orm_1 = require("drizzle-orm");
const trace_js_1 = require("../server/models/trace.js");
exports.isBlocklet = !!process.env.BLOCKLET_APP_DIR && !!process.env.BLOCKLET_PORT;
const PROVIDERS = [
    "openai",
    "google",
    "anthropic",
    "xai",
    "bedrock",
    "deepseek",
    "openrouter",
    "ollama",
    "doubao",
    "poe",
    "ideogram",
];
const getPriceValue = (prices, model) => {
    let price = prices?.[model];
    for (const provider of PROVIDERS) {
        if (price) {
            break;
        }
        price = prices?.[`${provider}/${model}`];
    }
    return price;
};
const calculateTokenAndCost = async (db, data) => {
    const { id, output } = data;
    const model = output?.model;
    const traces = await db
        .select({
        id: trace_js_1.Trace.id,
        token: trace_js_1.Trace.token,
        cost: trace_js_1.Trace.cost,
    })
        .from(trace_js_1.Trace)
        .where((0, drizzle_orm_1.eq)(trace_js_1.Trace.parentId, id))
        .execute();
    let token = 0;
    let cost = 0;
    if (traces.length > 0) {
        const tokenDecimal = traces.reduce((acc, curr) => acc.plus(curr.token || 0), new decimal_js_1.default(0));
        token = tokenDecimal.toNumber();
        const costDecimal = traces.reduce((acc, curr) => acc.plus(curr.cost || 0), new decimal_js_1.default(0));
        cost = costDecimal.gt(0) ? Number(costDecimal.toFixed(6)) : 0;
    }
    else {
        const inputTokens = output?.usage?.inputTokens || 0;
        const outputTokens = output?.usage?.outputTokens || 0;
        const seconds = output?.seconds || 4;
        token = new decimal_js_1.default(inputTokens).plus(outputTokens).toNumber();
        let prices = {};
        try {
            if (process?.env?.BLOCKLET_APP_DIR) {
                const fullPath = node_path_1.default.resolve(process.env.BLOCKLET_APP_DIR, "dist", "model-prices.json");
                const content = (0, node_fs_1.readFileSync)(fullPath, "utf-8");
                prices = JSON.parse(content);
            }
            else {
                const fullPath = "../../../dist/model-prices.json";
                // @ts-ignore
                prices = await Promise.resolve(`${fullPath}`).then(s => __importStar(require(s))).then((res) => res.default);
            }
        }
        catch (err) {
            console.warn(`[Observability] Failed to load model prices: ${err.message}. Cost calculation will be disabled.`);
            prices = {};
        }
        const getCost = (value) => {
            if (!value)
                return 0;
            if (value.mode === "image_generation") {
                return value.output_cost_per_image || 0;
            }
            if (value.mode === "video_generation") {
                return Number(new decimal_js_1.default(seconds).mul(value.output_cost_per_video_per_second || 0).toFixed(6));
            }
            return Number(new decimal_js_1.default(inputTokens)
                .mul(value.input_cost_per_token || 0)
                .plus(new decimal_js_1.default(outputTokens).mul(value.output_cost_per_token || 0))
                .toFixed(6));
        };
        if (prices && Object.keys(prices).length > 0 && model) {
            const value = getPriceValue(prices, model);
            cost = getCost(value);
        }
        else {
            cost = 0;
        }
    }
    return { token, cost };
};
/**
 * Propagate the error status of the root trace to all child traces that have not set the status.
 */
const propagateErrorStatusToChildren = async (db, trace) => {
    if (trace.rootId && !trace.parentId && trace.status?.code === api_1.SpanStatusCode.ERROR) {
        await db
            .update(trace_js_1.Trace)
            .set({ status: trace.status })
            .where((0, drizzle_orm_1.and)((0, drizzle_orm_1.eq)(trace_js_1.Trace.rootId, trace.rootId), (0, drizzle_orm_1.sql) `json_extract(${trace_js_1.Trace.status}, '$.code') = ${api_1.SpanStatusCode.UNSET}`))
            .execute();
    }
};
const updateStatusByIds = async (db, ids, status) => {
    if (!ids.length)
        return;
    if (!db)
        return;
    await db
        .update(trace_js_1.Trace)
        .set({ status, endTime: Date.now() })
        .where((0, drizzle_orm_1.and)((0, drizzle_orm_1.inArray)(trace_js_1.Trace.id, ids), (0, drizzle_orm_1.sql) `json_extract(${trace_js_1.Trace.status}, '$.code') = ${api_1.SpanStatusCode.UNSET}`))
        .execute();
};
exports.updateStatusByIds = updateStatusByIds;
const insertTrace = async (db, trace) => {
    const insertSql = (0, drizzle_orm_1.sql) `
    INSERT INTO Trace (
      id,
      rootId,
      parentId,
      name,
      startTime,
      endTime,
      attributes,
      status,
      userId,
      sessionId,
      componentId,
      action,
      token,
      cost,
      isImport
    ) VALUES (
      ${trace.id},
      ${trace.rootId},
      ${trace.parentId || null},
      ${trace.name},
      ${trace.startTime},
      ${trace.endTime},
      ${JSON.stringify(trace.attributes)},
      ${JSON.stringify(trace.status)},
      ${trace.userId || null},
      ${trace.sessionId || null},
      ${trace.componentId || null},
      ${trace.action || null},
      ${trace.token || 0},
      ${trace.cost || 0},
      ${trace.isImport || 0}
    )
    ON CONFLICT(id)
    DO UPDATE SET
      name = excluded.name,
      startTime = excluded.startTime,
      endTime = excluded.endTime,
      status = excluded.status,
      componentId = excluded.componentId,
      action = excluded.action,
      isImport = excluded.isImport;
  `;
    await db?.run?.(insertSql);
    // Batch update status to ERROR if the root trace is ERROR
    await propagateErrorStatusToChildren(db, trace);
};
exports.insertTrace = insertTrace;
const updateTrace = async (db, id, data) => {
    // get existing attributes
    const existing = await db.select().from(trace_js_1.Trace).where((0, drizzle_orm_1.eq)(trace_js_1.Trace.id, id)).execute();
    if (!existing.length)
        return;
    const trace = existing[0];
    if (!trace)
        return;
    const currentAttributes = trace?.attributes;
    const hasInput = data.input && Object.keys(data.input).length > 0;
    const hasOutput = data.output && Object.keys(data.output).length > 0;
    const hasUserContext = data.userContext && Object.keys(data.userContext).length > 0;
    const hasMemories = data.memories && Object.keys(data.memories).length > 0;
    let attributes = {};
    if (currentAttributes && typeof currentAttributes === "string") {
        try {
            attributes = JSON.parse(currentAttributes);
        }
        catch (error) {
            console.error("parse attributes error", error.message);
            attributes = {};
        }
    }
    else if (currentAttributes && typeof currentAttributes === "object") {
        attributes = currentAttributes;
    }
    // merge attributes
    const updatedAttributes = {
        ...attributes,
        ...(hasInput && { input: data.input }),
        ...(hasOutput && { output: data.output }),
        ...(hasUserContext && { userContext: data.userContext }),
        ...(hasMemories && { memories: data.memories }),
    };
    // calculate token and cost
    const { token, cost } = hasOutput && data.output
        ? await calculateTokenAndCost(db, { id, output: data.output })
        : { token: 0, cost: 0 };
    const params = {
        attributes: updatedAttributes,
        token: token || 0,
        cost: cost || 0,
    };
    if (data.status) {
        params.status = data.status;
        params.endTime = Date.now();
    }
    if (hasUserContext) {
        params.userId = data.userContext?.userId;
        params.sessionId = data.userContext?.sessionId;
    }
    await db.update(trace_js_1.Trace).set(params).where((0, drizzle_orm_1.eq)(trace_js_1.Trace.id, id)).execute();
    if (data.status) {
        await propagateErrorStatusToChildren(db, { ...trace, status: data.status });
    }
};
exports.updateTrace = updateTrace;
