"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 archiver_1 = __importDefault(require("archiver"));
const drizzle_orm_1 = require("drizzle-orm");
const express_1 = __importDefault(require("express"));
const unzipper_1 = __importDefault(require("unzipper"));
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 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(""),
    showImportedOnly: zod_1.z
        .string()
        .optional()
        .transform((val) => val === "true")
        .default("false"),
});
const schema_js_1 = require("../../core/schema.js");
exports.default = ({ sse, middleware, options, }) => {
    const router = express_1.default.Router();
    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, showImportedOnly } = 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 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}%`), (0, drizzle_orm_1.like)(trace_js_1.Trace.userId, `%${(searchText || "").replace("did:abt:", "")}%`), (0, drizzle_orm_1.like)(trace_js_1.Trace.remark, `%${searchText}%`));
        let whereClause = searchText ? (0, drizzle_orm_1.and)(rootFilter, searchFilter) : rootFilter;
        if (startDate && endDate) {
            const start = new Date(startDate);
            start.setHours(0, 0, 0, 0);
            const end = new Date(endDate);
            end.setHours(23, 59, 59, 999);
            whereClause = (0, drizzle_orm_1.and)(whereClause, (0, drizzle_orm_1.between)(trace_js_1.Trace.startTime, start.getTime(), end.getTime()));
        }
        if (componentId) {
            whereClause = (0, drizzle_orm_1.and)(whereClause, (0, drizzle_orm_1.eq)(trace_js_1.Trace.componentId, componentId));
        }
        if (showImportedOnly) {
            whereClause = (0, drizzle_orm_1.and)(whereClause, (0, drizzle_orm_1.eq)(trace_js_1.Trace.isImport, 1));
        }
        const [count, rootCalls] = await Promise.all([
            db.select({ count: (0, drizzle_orm_1.sql) `count(*)` }).from(trace_js_1.Trace).where(whereClause).execute(),
            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,
                inputPreview: (0, drizzle_orm_1.sql) `SUBSTR(COALESCE(CAST(JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.input') AS TEXT), ''), 1, 150)`,
                inputLength: (0, drizzle_orm_1.sql) `LENGTH(COALESCE(CAST(JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.input') AS TEXT), ''))`,
                outputPreview: (0, drizzle_orm_1.sql) `SUBSTR(COALESCE(CAST(JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.output') AS TEXT), ''), 1, 150)`,
                outputLength: (0, drizzle_orm_1.sql) `LENGTH(COALESCE(CAST(JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.output') AS TEXT), ''))`,
                metadata: (0, drizzle_orm_1.sql) `JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.metadata')`,
                userId: trace_js_1.Trace.userId,
                componentId: trace_js_1.Trace.componentId,
                token: trace_js_1.Trace.token,
                cost: trace_js_1.Trace.cost,
                remark: trace_js_1.Trace.remark,
            })
                .from(trace_js_1.Trace)
                .where(whereClause)
                .orderBy((0, drizzle_orm_1.desc)(trace_js_1.Trace.startTime))
                .limit(pageSize)
                .offset(offset)
                .execute(),
        ]);
        const total = Number(count[0].count ?? 0);
        const processedRootCalls = rootCalls.map((call) => {
            const { inputPreview, inputLength, outputPreview, outputLength, metadata, ...rest } = call;
            return {
                ...rest,
                attributes: {
                    input: inputPreview + (inputLength > 150 ? "..." : ""),
                    output: outputPreview + (outputLength > 150 ? "..." : ""),
                    metadata: metadata ? JSON.parse(metadata) : null,
                },
            };
        });
        res.json({
            total,
            page,
            pageSize,
            data: processedRootCalls,
        });
    });
    router.get("/tree/summary", ...middleware, async (req, res) => {
        const db = req.app.locals.db;
        res.setHeader("Cache-Control", "public, max-age=600");
        const [stats, llmStats] = await Promise.all([
            db
                .select({
                totalCount: (0, drizzle_orm_1.sql) `COUNT(*)`,
                successCount: (0, drizzle_orm_1.sql) `
          SUM(CASE
            WHEN JSON_EXTRACT(${trace_js_1.Trace.status}, '$.code') = 1 AND ${trace_js_1.Trace.endTime} > 0
            THEN 1 ELSE 0
          END)
        `,
                totalToken: (0, drizzle_orm_1.sql) `
          SUM(CASE
            WHEN ${trace_js_1.Trace.endTime} > 0
            THEN COALESCE(${trace_js_1.Trace.token}, 0) ELSE 0
          END)
        `,
                totalCost: (0, drizzle_orm_1.sql) `
          SUM(CASE
            WHEN ${trace_js_1.Trace.endTime} > 0
            THEN COALESCE(${trace_js_1.Trace.cost}, 0) ELSE 0
          END)
        `,
                maxLatency: (0, drizzle_orm_1.sql) `
          MAX(CASE
            WHEN ${trace_js_1.Trace.endTime} > 0
            THEN ${trace_js_1.Trace.endTime} - ${trace_js_1.Trace.startTime} ELSE NULL
          END)
        `,
                minLatency: (0, drizzle_orm_1.sql) `
          MIN(CASE
            WHEN ${trace_js_1.Trace.endTime} > 0
            THEN ${trace_js_1.Trace.endTime} - ${trace_js_1.Trace.startTime} ELSE NULL
          END)
        `,
                avgLatency: (0, drizzle_orm_1.sql) `
          AVG(CASE
            WHEN ${trace_js_1.Trace.endTime} > 0
            THEN ${trace_js_1.Trace.endTime} - ${trace_js_1.Trace.startTime} ELSE NULL
          END)
        `,
                totalDuration: (0, drizzle_orm_1.sql) `
          SUM(CASE
            WHEN ${trace_js_1.Trace.endTime} > 0
            THEN ${trace_js_1.Trace.endTime} - ${trace_js_1.Trace.startTime} ELSE 0
          END)
        `,
            })
                .from(trace_js_1.Trace)
                .where((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)))
                .execute(),
            db
                .select({
                llmTotalCount: (0, drizzle_orm_1.sql) `COUNT(*)`,
                llmSuccessCount: (0, drizzle_orm_1.sql) `
            SUM(CASE
              WHEN JSON_EXTRACT(${trace_js_1.Trace.status}, '$.code') = 1 AND ${trace_js_1.Trace.endTime} > 0
              THEN 1 ELSE 0
            END)
          `,
                llmTotalDuration: (0, drizzle_orm_1.sql) `
            SUM(CASE
              WHEN ${trace_js_1.Trace.endTime} > 0
              THEN ${trace_js_1.Trace.endTime} - ${trace_js_1.Trace.startTime} ELSE 0
            END)
          `,
            })
                .from(trace_js_1.Trace)
                .where((0, drizzle_orm_1.like)(trace_js_1.Trace.name, "%Model%"))
                .execute(),
        ]);
        const result = stats[0];
        const llm = llmStats[0];
        res.json({
            totalCount: result?.totalCount ?? 0,
            successCount: result?.successCount ?? 0,
            failCount: (result?.totalCount ?? 0) - (result?.successCount ?? 0),
            totalToken: result?.totalToken ?? 0,
            totalCost: result?.totalCost ?? 0,
            maxLatency: result?.maxLatency ?? 0,
            minLatency: result?.minLatency ?? 0,
            avgLatency: result?.avgLatency ?? 0,
            totalDuration: result?.totalDuration ?? 0,
            llmSuccessCount: llm?.llmSuccessCount ?? 0,
            llmTotalCount: llm?.llmTotalCount ?? 0,
            llmTotalDuration: llm?.llmTotalDuration ?? 0,
        });
    });
    router.get("/tree/components", ...middleware, async (req, res) => {
        const db = req.app.locals.db;
        res.set("Cache-Control", "public, max-age=600");
        const components = await db
            .selectDistinct({ 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)))
            .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 trace = await db.select().from(trace_js_1.Trace).where((0, drizzle_orm_1.eq)(trace_js_1.Trace.id, id)).execute();
        if (trace.length === 0)
            throw new Error(`Not found trace: ${id}`);
        res.json({ data: trace[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,
            outputUsage: (0, drizzle_orm_1.sql) `JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.output.usage')`,
            outputModel: (0, drizzle_orm_1.sql) `JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.output.model')`,
            outputSeconds: (0, drizzle_orm_1.sql) `JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.output.seconds')`,
            agentTag: (0, drizzle_orm_1.sql) `JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.agentTag')`,
            metadata: (0, drizzle_orm_1.sql) `JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.metadata')`,
            taskTitle: (0, drizzle_orm_1.sql) `JSON_EXTRACT(${trace_js_1.Trace.attributes}, '$.taskTitle')`,
            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 { outputUsage, outputModel, outputSeconds, agentTag, metadata, taskTitle, ...rest } = call;
            const safeJsonParse = (value) => {
                if (!value)
                    return null;
                try {
                    return JSON.parse(value);
                }
                catch {
                    return value;
                }
            };
            const attributes = {
                output: {
                    usage: safeJsonParse(outputUsage),
                    model: safeJsonParse(outputModel),
                    seconds: safeJsonParse(outputSeconds),
                },
                agentTag: safeJsonParse(agentTag),
                metadata: safeJsonParse(metadata),
                taskTitle: safeJsonParse(taskTitle),
            };
            calls.set(call.id, { ...rest, children: [], 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.patch("/tree/:id", async (req, res) => {
        const db = req.app.locals.db;
        const { id } = req.params;
        const { input, output, userContext, memories, status } = req.body;
        if (!id) {
            res.status(400).json({ error: "id is required" });
            return;
        }
        await (0, util_js_1.updateTrace)(db, id, { input, output, userContext, memories, status });
        res.json({ code: 0, message: "ok" });
    });
    router.delete("/tree", async (req, res) => {
        const db = req.app.locals.db;
        const ids = req.body?.ids;
        if (Array.isArray(ids) && ids.length > 0) {
            await db.delete(trace_js_1.Trace).where((0, drizzle_orm_1.inArray)(trace_js_1.Trace.id, ids)).execute();
            await db.clean?.();
            res.json({ code: 0, message: `${ids.length} traces deleted` });
            return;
        }
        await db.delete(trace_js_1.Trace).execute();
        await db.clean?.();
        res.json({ code: 0, message: "all traces deleted" });
    });
    router.post("/remark", ...middleware, async (req, res) => {
        const db = req.app.locals.db;
        const { id, remark } = req.body;
        if (!id) {
            res.status(400).json({ error: "id is required" });
            return;
        }
        if (remark && remark.length > 50) {
            res.status(400).json({ error: "Remark cannot exceed 50 characters" });
            return;
        }
        try {
            await db
                .update(trace_js_1.Trace)
                .set({ remark: remark || null })
                .where((0, drizzle_orm_1.eq)(trace_js_1.Trace.id, id))
                .execute();
            res.json({ code: 0, message: "remark updated successfully" });
        }
        catch (error) {
            console.error("Failed to update remark:", error);
            res.status(500).json({ error: "Failed to update remark" });
        }
    });
    router.post("/download", ...middleware, async (req, res) => {
        const db = req.app.locals.db;
        try {
            const ids = req.body?.ids;
            if (!ids || !Array.isArray(ids) || ids.length === 0) {
                res.status(400).json({ error: "ids array is required" });
                return;
            }
            const selectedTraces = await db.select().from(trace_js_1.Trace).where((0, drizzle_orm_1.inArray)(trace_js_1.Trace.id, ids)).execute();
            if (selectedTraces.length === 0) {
                res.status(404).json({ error: "No traces found for provided ids" });
                return;
            }
            const rootIds = selectedTraces.map((t) => t.rootId).filter((id) => !!id);
            const allTraces = rootIds.length > 0
                ? await db.select().from(trace_js_1.Trace).where((0, drizzle_orm_1.inArray)(trace_js_1.Trace.rootId, rootIds)).execute()
                : selectedTraces;
            const now = new Date().toISOString().replace(/[:.]/g, "-");
            const archive = (0, archiver_1.default)("zip", { zlib: { level: 9 } });
            const jsonContent = JSON.stringify(allTraces, null, 2);
            res.setHeader("Content-Type", "application/zip");
            res.setHeader("Content-Disposition", `attachment; filename=traces-${now}.zip`);
            archive.on("error", (err) => {
                throw err;
            });
            archive.pipe(res);
            archive.append(jsonContent, { name: `traces-${now}.json` });
            await archive.finalize();
        }
        catch (error) {
            console.error("Failed to download traces:", error);
            res.status(500).json({ error: "Failed to download traces" });
        }
    });
    router.post("/upload", ...middleware, async (req, res) => {
        const db = req.app.locals.db;
        try {
            let traces;
            const contentType = req.headers["content-type"];
            if (contentType?.includes("application/zip") ||
                contentType?.includes("application/x-zip-compressed")) {
                const chunks = [];
                req.on("data", (chunk) => {
                    chunks.push(chunk);
                });
                await new Promise((resolve, reject) => {
                    req.on("end", () => resolve());
                    req.on("error", reject);
                });
                const zipBuffer = Buffer.concat(chunks);
                const directory = await unzipper_1.default.Open.buffer(zipBuffer);
                const jsonFile = directory.files.find((file) => file.path.endsWith(".json"));
                if (!jsonFile) {
                    res.status(400).json({ error: "No JSON file found in zip archive" });
                    return;
                }
                const jsonContent = await jsonFile.buffer();
                traces = JSON.parse(jsonContent.toString("utf-8"));
            }
            else {
                const body = req.body;
                traces = body?.traces || body;
                if (!Array.isArray(traces) && traces?.traces) {
                    traces = traces.traces;
                }
            }
            if (!Array.isArray(traces)) {
                res.status(400).json({ error: "Invalid data format: traces array is required" });
                return;
            }
            const invalidTraces = [];
            for (const trace of traces) {
                if (!trace.id || !trace.rootId || !trace.name || trace.startTime === undefined) {
                    invalidTraces.push(trace.id || "unknown");
                }
            }
            if (invalidTraces.length > 0) {
                res.status(400).json({
                    error: "Invalid trace data: missing required fields",
                    invalidTraces,
                    details: "Each trace must have id, rootId, name, and startTime",
                });
                return;
            }
            const uploadIds = traces.map((t) => t.id);
            const existingTraces = await db
                .select({
                id: trace_js_1.Trace.id,
                rootId: trace_js_1.Trace.rootId,
                parentId: trace_js_1.Trace.parentId,
            })
                .from(trace_js_1.Trace)
                .where((0, drizzle_orm_1.inArray)(trace_js_1.Trace.id, uploadIds))
                .execute();
            const existingKeys = new Set(existingTraces.map((t) => `${t.id}|${t.rootId}|${t.parentId ?? ""}`));
            const tracesToInsert = traces.filter((t) => !existingKeys.has(`${t.id}|${t.rootId}|${t.parentId ?? ""}`));
            const skippedCount = traces.length - tracesToInsert.length;
            let successCount = 0;
            let failCount = 0;
            for (const trace of tracesToInsert) {
                trace.isImport = 1;
                try {
                    await (0, util_js_1.insertTrace)(db, trace);
                    successCount++;
                }
                catch (err) {
                    console.error(`Failed to insert trace ${trace.id}:`, err);
                    failCount++;
                }
            }
            res.json({
                code: 0,
                message: "Upload completed",
                successCount,
                failCount,
                skippedCount,
                total: traces.length,
            });
        }
        catch (error) {
            console.error("Failed to upload traces:", error);
            res.status(500).json({ error: "Failed to upload traces" });
        }
    });
    return router;
};
