"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
const node_fs_1 = require("node:fs");
const promises_1 = require("node:fs/promises");
const drizzle_orm_1 = require("drizzle-orm");
const express_1 = __importDefault(require("express"));
const yaml_1 = require("yaml");
const zod_1 = require("zod");
const util_js_1 = require("../../core/util.js");
const trace_js_1 = require("../models/trace.js");
const index_js_1 = require("../utils/index.js");
const router = express_1.default.Router();
const traceTreeQuerySchema = zod_1.z.object({
    page: zod_1.z.coerce
        .number()
        .int()
        .min(0)
        .catch(() => 0)
        .default(0),
    pageSize: zod_1.z.coerce
        .number()
        .int()
        .min(1)
        .catch(() => 10)
        .default(10),
    searchText: zod_1.z.string().optional().default(""),
    componentId: zod_1.z.string().optional().default(""),
    startDate: zod_1.z.string().optional().default(""),
    endDate: zod_1.z.string().optional().default(""),
});
const schema_js_1 = require("../../core/schema.js");
exports.default = ({ sse, middleware, options, }) => {
    router.get("/tree", ...middleware, async (req, res) => {
        const db = req.app.locals.db;
        const queryResult = traceTreeQuerySchema.safeParse(req.query);
        if (!queryResult.success) {
            res.status(400).json({
                error: "Invalid query parameters",
                details: queryResult.error.errors,
            });
            return;
        }
        const { page, pageSize, searchText, componentId, startDate, endDate } = queryResult.data;
        const offset = page * pageSize;
        if (!Number.isSafeInteger(offset) || offset > Number.MAX_SAFE_INTEGER) {
            res.status(400).json({
                error: "Page number too large, would cause overflow",
                details: { page, pageSize, calculatedOffset: offset },
            });
            return;
        }
        const rootFilter = (0, drizzle_orm_1.and)((0, drizzle_orm_1.isNull)(trace_js_1.Trace.parentId), (0, drizzle_orm_1.isNull)(trace_js_1.Trace.action));
        const count = await db.select({ count: (0, drizzle_orm_1.sql) `count(*)` }).from(trace_js_1.Trace).where(rootFilter).execute();
        const total = Number(count[0].count ?? 0);
        const searchFilter = (0, drizzle_orm_1.or)((0, drizzle_orm_1.like)(trace_js_1.Trace.attributes, `%${searchText}%`), (0, drizzle_orm_1.like)(trace_js_1.Trace.name, `%${searchText}%`), (0, drizzle_orm_1.like)(trace_js_1.Trace.id, `%${searchText}%`));
        let whereClause = searchText ? (0, drizzle_orm_1.and)(rootFilter, searchFilter) : rootFilter;
        if (startDate && endDate) {
            whereClause = (0, drizzle_orm_1.and)(whereClause, (0, drizzle_orm_1.between)(trace_js_1.Trace.startTime, new Date(startDate).getTime(), new Date(endDate).getTime()));
        }
        if (componentId) {
            whereClause = (0, drizzle_orm_1.and)(whereClause, (0, drizzle_orm_1.eq)(trace_js_1.Trace.componentId, componentId));
        }
        const rootCalls = await db
            .select({
            id: trace_js_1.Trace.id,
            rootId: trace_js_1.Trace.rootId,
            parentId: trace_js_1.Trace.parentId,
            name: trace_js_1.Trace.name,
            startTime: trace_js_1.Trace.startTime,
            endTime: trace_js_1.Trace.endTime,
            status: trace_js_1.Trace.status,
            attributes: (0, drizzle_orm_1.sql) `
          CASE 
            WHEN ${trace_js_1.Trace.attributes} IS NULL THEN JSON_OBJECT('input', '', 'output', '')
            ELSE JSON_OBJECT(
              'input', 
              CASE 
                WHEN JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.input') IS NOT NULL 
                THEN SUBSTR(CAST(JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.input') AS TEXT), 1, 150) ||
                CASE WHEN LENGTH(CAST(JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.input') AS TEXT)) > 150 THEN '...' ELSE '' END
                ELSE ''
              END,
              'output',
              CASE 
                WHEN JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.output') IS NOT NULL 
                THEN SUBSTR(CAST(JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.output') AS TEXT), 1, 150) ||
                CASE WHEN LENGTH(CAST(JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.output') AS TEXT)) > 150 THEN '...' ELSE '' END
                ELSE ''
              END
            )
          END
        `,
            userId: trace_js_1.Trace.userId,
            componentId: trace_js_1.Trace.componentId,
            token: trace_js_1.Trace.token,
            cost: trace_js_1.Trace.cost,
        })
            .from(trace_js_1.Trace)
            .where(whereClause)
            .orderBy((0, drizzle_orm_1.desc)(trace_js_1.Trace.startTime))
            .limit(pageSize)
            .offset(offset)
            .execute();
        const processedRootCalls = rootCalls.map((call) => {
            try {
                return {
                    ...call,
                    attributes: JSON.parse(call.attributes),
                };
            }
            catch {
                return call;
            }
        });
        res.json({
            total,
            page,
            pageSize,
            data: processedRootCalls,
        });
    });
    router.get("/tree/summary", ...middleware, async (req, res) => {
        const db = req.app.locals.db;
        const baseWhere = (0, drizzle_orm_1.and)((0, drizzle_orm_1.isNull)(trace_js_1.Trace.parentId), (0, drizzle_orm_1.isNull)(trace_js_1.Trace.action));
        const totalCountResult = await db
            .select({ totalCount: (0, drizzle_orm_1.sql) `COUNT(*)` })
            .from(trace_js_1.Trace)
            .where(baseWhere)
            .execute();
        const totalCount = totalCountResult[0]?.totalCount ?? 0;
        const statusResult = await db
            .select({
            successCount: (0, drizzle_orm_1.sql) `
          COUNT(
            CASE 
              WHEN JSON_EXTRACT(${trace_js_1.Trace.status}, '$.code') = 1 AND ${trace_js_1.Trace.endTime} > 0 
              THEN 1 
            END
          )
        `,
        })
            .from(trace_js_1.Trace)
            .where(baseWhere)
            .execute();
        const successCount = statusResult[0]?.successCount ?? 0;
        const failCount = totalCount - successCount;
        const TotalStatsResult = await db
            .select({
            totalToken: (0, drizzle_orm_1.sql) `COALESCE(SUM(CASE WHEN ${trace_js_1.Trace.endTime} > 0 THEN ${trace_js_1.Trace.token} ELSE 0 END), 0)`,
            totalCost: (0, drizzle_orm_1.sql) `COALESCE(SUM(CASE WHEN ${trace_js_1.Trace.endTime} > 0 THEN ${trace_js_1.Trace.cost} ELSE 0 END), 0)`,
            maxLatency: (0, drizzle_orm_1.sql) `COALESCE(MAX(CASE WHEN ${trace_js_1.Trace.endTime} > 0 THEN ${trace_js_1.Trace.endTime} - ${trace_js_1.Trace.startTime} ELSE NULL END), 0)`,
            minLatency: (0, drizzle_orm_1.sql) `COALESCE(MIN(CASE WHEN ${trace_js_1.Trace.endTime} > 0 THEN ${trace_js_1.Trace.endTime} - ${trace_js_1.Trace.startTime} ELSE NULL END), 0)`,
            avgLatency: (0, drizzle_orm_1.sql) `COALESCE(AVG(CASE WHEN ${trace_js_1.Trace.endTime} > 0 THEN ${trace_js_1.Trace.endTime} - ${trace_js_1.Trace.startTime} ELSE NULL END), 0)`,
            totalDuration: (0, drizzle_orm_1.sql) `COALESCE(SUM(CASE WHEN ${trace_js_1.Trace.endTime} > 0 THEN ${trace_js_1.Trace.endTime} - ${trace_js_1.Trace.startTime} ELSE 0 END), 0)`,
        })
            .from(trace_js_1.Trace)
            .where(baseWhere)
            .execute();
        const totalStats = TotalStatsResult[0];
        const llmWhere = (0, drizzle_orm_1.and)((0, drizzle_orm_1.sql) `${trace_js_1.Trace.attributes} IS NOT NULL`, (0, drizzle_orm_1.sql) `JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.output.model') IS NOT NULL`);
        const llmStatsResult = await db
            .select({
            llmTotalCount: (0, drizzle_orm_1.sql) `COUNT(*)`,
            llmSuccessCount: (0, drizzle_orm_1.sql) `
          COUNT(
            CASE 
              WHEN JSON_EXTRACT(${trace_js_1.Trace.status}, '$.code') = 1 AND ${trace_js_1.Trace.endTime} > 0 
              THEN 1 
            END
          )
        `,
            llmTotalDuration: (0, drizzle_orm_1.sql) `
          COALESCE(SUM(
            CASE WHEN ${trace_js_1.Trace.endTime} > 0 THEN ${trace_js_1.Trace.endTime} - ${trace_js_1.Trace.startTime} ELSE 0 END
          ), 0)
        `,
        })
            .from(trace_js_1.Trace)
            .where(llmWhere)
            .execute();
        res.json({
            totalCount,
            successCount,
            failCount,
            totalToken: totalStats?.totalToken ?? 0,
            totalCost: totalStats?.totalCost ?? 0,
            maxLatency: totalStats?.maxLatency ?? 0,
            minLatency: totalStats?.minLatency ?? 0,
            avgLatency: totalStats?.avgLatency ?? 0,
            totalDuration: totalStats?.totalDuration ?? 0,
            llmSuccessCount: llmStatsResult[0]?.llmSuccessCount ?? 0,
            llmTotalCount: llmStatsResult[0]?.llmTotalCount ?? 0,
            llmTotalDuration: llmStatsResult[0]?.llmTotalDuration ?? 0,
        });
    });
    router.get("/tree/components", ...middleware, async (req, res) => {
        const db = req.app.locals.db;
        const components = await db
            .select({ componentId: trace_js_1.Trace.componentId })
            .from(trace_js_1.Trace)
            .where((0, drizzle_orm_1.and)((0, drizzle_orm_1.isNotNull)(trace_js_1.Trace.componentId), (0, drizzle_orm_1.isNull)(trace_js_1.Trace.action)))
            .groupBy(trace_js_1.Trace.componentId)
            .execute();
        const componentIds = components.map((c) => c.componentId).filter(Boolean);
        res.json({
            data: componentIds,
            total: componentIds.length,
        });
    });
    router.get("/tree/stats", async (req, res) => {
        const db = req.app.locals.db;
        const rootFilter = (0, drizzle_orm_1.and)((0, drizzle_orm_1.or)((0, drizzle_orm_1.isNull)(trace_js_1.Trace.parentId), (0, drizzle_orm_1.eq)(trace_js_1.Trace.parentId, "")), (0, drizzle_orm_1.isNull)(trace_js_1.Trace.action));
        const [latestRoot] = (await db
            .select()
            .from(trace_js_1.Trace)
            .where(rootFilter)
            .orderBy((0, drizzle_orm_1.desc)(trace_js_1.Trace.startTime))
            .limit(1)
            .execute()) || [];
        const settingPath = (0, index_js_1.getGlobalSettingPath)();
        let settings = {
            lastTrace: { id: "", endTime: 0 },
        };
        if (!(0, node_fs_1.existsSync)(settingPath)) {
            await (0, promises_1.writeFile)(settingPath, (0, yaml_1.stringify)(settings));
        }
        else {
            settings = (0, yaml_1.parse)(await (0, promises_1.readFile)(settingPath, "utf8"));
        }
        const lastTraceChanged = latestRoot &&
            (settings.lastTrace?.id !== latestRoot.id ||
                settings.lastTrace?.endTime !== latestRoot.endTime);
        if (lastTraceChanged) {
            await (0, promises_1.writeFile)(settingPath, (0, yaml_1.stringify)({
                ...settings,
                lastTrace: {
                    id: latestRoot.id,
                    rootId: latestRoot.rootId,
                    startTime: latestRoot.startTime,
                    endTime: latestRoot.endTime,
                },
            }));
        }
        res.json({ code: 0, data: { lastTraceChanged } });
    });
    router.get("/tree/children/:id", async (req, res) => {
        const id = req.params.id;
        if (!id)
            throw new Error("id is required");
        const db = req.app.locals.db;
        const rootCalls = await db.select().from(trace_js_1.Trace).where((0, drizzle_orm_1.eq)(trace_js_1.Trace.id, id)).execute();
        if (rootCalls.length === 0)
            throw new Error(`Not found trace: ${id}`);
        const all = await db.select().from(trace_js_1.Trace).where((0, drizzle_orm_1.eq)(trace_js_1.Trace.id, id)).execute();
        res.json({ data: all[0] });
    });
    router.get("/tree/:id", async (req, res) => {
        const id = req.params.id;
        if (!id)
            throw new Error("id is required");
        const db = req.app.locals.db;
        const rootCalls = await db.select().from(trace_js_1.Trace).where((0, drizzle_orm_1.eq)(trace_js_1.Trace.id, id)).execute();
        if (rootCalls.length === 0)
            throw new Error(`Not found trace: ${id}`);
        const rootCallIds = rootCalls.map((r) => r.rootId).filter((id) => !!id);
        const all = await db
            .select({
            id: trace_js_1.Trace.id,
            rootId: trace_js_1.Trace.rootId,
            parentId: trace_js_1.Trace.parentId,
            name: trace_js_1.Trace.name,
            startTime: trace_js_1.Trace.startTime,
            endTime: trace_js_1.Trace.endTime,
            status: trace_js_1.Trace.status,
            attributes: (0, drizzle_orm_1.sql) `
          CASE
            WHEN JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.output.usage') IS NOT NULL THEN
              JSON_OBJECT(
                'output', JSON_OBJECT(
                  'usage', JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.output.usage'),
                  'model', JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.output.model')
                ),
                'agentTag', JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.agentTag')
              )
            ELSE JSON_OBJECT(
              'output', JSON_OBJECT(),
              'agentTag', JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.agentTag')
            )
          END
        `,
            userId: trace_js_1.Trace.userId,
            sessionId: trace_js_1.Trace.sessionId,
            componentId: trace_js_1.Trace.componentId,
            token: trace_js_1.Trace.token,
            cost: trace_js_1.Trace.cost,
        })
            .from(trace_js_1.Trace)
            .where((0, drizzle_orm_1.inArray)(trace_js_1.Trace.rootId, rootCallIds))
            .execute();
        const calls = new Map();
        all.forEach((call) => {
            const { attributes } = call;
            let _attributes = {};
            if (attributes) {
                try {
                    _attributes = JSON.parse(attributes);
                }
                catch (err) {
                    console.error(`parse attributes failed for trace ${call.id}:`, err);
                }
            }
            calls.set(call.id, { ...call, children: [], attributes: _attributes });
        });
        all.forEach((call) => {
            if (call.parentId) {
                const parent = calls.get(call.parentId);
                if (parent) {
                    parent.children.push(calls.get(call.id));
                }
            }
        });
        const trees = rootCalls.map((run) => calls.get(run.id));
        res.json({ data: trees[0] });
    });
    router.post("/tree", async (req, res) => {
        if (!req.body || req.body.length === 0) {
            throw new Error("req.body is empty");
        }
        const validatedTraces = schema_js_1.createTraceBatchSchema.parse(req.body);
        let live = false;
        const settingPath = (0, index_js_1.getGlobalSettingPath)();
        if (!(0, node_fs_1.existsSync)(settingPath)) {
            live = false;
        }
        else {
            const setting = (0, yaml_1.parse)(await (0, promises_1.readFile)(settingPath, "utf8"));
            live = setting.live;
        }
        const db = req.app.locals.db;
        for (const trace of validatedTraces) {
            try {
                const mapping = {
                    files: options?.formatOutputFiles,
                    images: options?.formatOutputImages,
                };
                for (const [key, formatter] of Object.entries(mapping)) {
                    const items = trace?.attributes?.output?.[key];
                    if (items?.length && formatter) {
                        trace.attributes.output[key] = await formatter(items);
                    }
                }
                await (0, util_js_1.insertTrace)(db, trace);
            }
            catch (err) {
                console.error(`upsert spans failed for trace ${trace.id}:`, err);
            }
        }
        if (live) {
            sse.send({ type: "event", data: {} });
        }
        res.json({ code: 0, message: "ok" });
    });
    router.delete("/tree", async (req, res) => {
        const db = req.app.locals.db;
        await db.delete(trace_js_1.Trace).execute();
        res.json({ code: 0, message: "ok" });
    });
    return router;
};
