import { existsSync } from "node:fs";
import { readFile, writeFile } from "node:fs/promises";
import archiver from "archiver";
import { and, between, desc, eq, inArray, isNotNull, isNull, like, or, sql } from "drizzle-orm";
import express from "express";
import unzipper from "unzipper";
import { parse, stringify } from "yaml";
import { z } from "zod";
import { insertTrace } from "../../core/util.js";
import { Trace } from "../models/trace.js";
import { getGlobalSettingPath } from "../utils/index.js";
const traceTreeQuerySchema = z.object({
    page: z.coerce
        .number()
        .int()
        .min(0)
        .catch(() => 0)
        .default(0),
    pageSize: z.coerce
        .number()
        .int()
        .min(1)
        .catch(() => 10)
        .default(10),
    searchText: z.string().optional().default(""),
    componentId: z.string().optional().default(""),
    startDate: z.string().optional().default(""),
    endDate: z.string().optional().default(""),
    showImportedOnly: z
        .string()
        .optional()
        .transform((val) => val === "true")
        .default("false"),
});
import { createTraceBatchSchema } from "../../core/schema.js";
export default ({ sse, middleware, options, }) => {
    const router = express.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 = and(isNull(Trace.parentId), isNull(Trace.action));
        const searchFilter = or(like(Trace.attributes, `%${searchText}%`), like(Trace.name, `%${searchText}%`), like(Trace.id, `%${searchText}%`), like(Trace.userId, `%${(searchText || "").replace("did:abt:", "")}%`), like(Trace.remark, `%${searchText}%`));
        let whereClause = searchText ? 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 = and(whereClause, between(Trace.startTime, start.getTime(), end.getTime()));
        }
        if (componentId) {
            whereClause = and(whereClause, eq(Trace.componentId, componentId));
        }
        if (showImportedOnly) {
            whereClause = and(whereClause, eq(Trace.isImport, 1));
        }
        const count = await db
            .select({ count: sql `count(*)` })
            .from(Trace)
            .where(whereClause)
            .execute();
        const total = Number(count[0].count ?? 0);
        const rootCalls = await db
            .select({
            id: Trace.id,
            rootId: Trace.rootId,
            parentId: Trace.parentId,
            name: Trace.name,
            startTime: Trace.startTime,
            endTime: Trace.endTime,
            status: Trace.status,
            attributes: sql `
          CASE
            WHEN ${Trace.attributes} IS NULL THEN JSON_OBJECT('input', '', 'output', '', 'metadata', NULL)
            ELSE JSON_OBJECT(
              'input',
              CASE
                WHEN JSON_EXTRACT(${Trace.attributes}, '$.input') IS NOT NULL
                THEN SUBSTR(CAST(JSON_EXTRACT(${Trace.attributes}, '$.input') AS TEXT), 1, 150) ||
                CASE WHEN LENGTH(CAST(JSON_EXTRACT(${Trace.attributes}, '$.input') AS TEXT)) > 150 THEN '...' ELSE '' END
                ELSE ''
              END,
              'output',
              CASE
                WHEN JSON_EXTRACT(${Trace.attributes}, '$.output') IS NOT NULL
                THEN SUBSTR(CAST(JSON_EXTRACT(${Trace.attributes}, '$.output') AS TEXT), 1, 150) ||
                CASE WHEN LENGTH(CAST(JSON_EXTRACT(${Trace.attributes}, '$.output') AS TEXT)) > 150 THEN '...' ELSE '' END
                ELSE ''
              END,
              'metadata',
              JSON_EXTRACT(${Trace.attributes}, '$.metadata')
            )
          END
        `,
            userId: Trace.userId,
            componentId: Trace.componentId,
            token: Trace.token,
            cost: Trace.cost,
            remark: Trace.remark,
        })
            .from(Trace)
            .where(whereClause)
            .orderBy(desc(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 = and(isNull(Trace.parentId), isNull(Trace.action));
        const totalCountResult = await db
            .select({ totalCount: sql `COUNT(*)` })
            .from(Trace)
            .where(baseWhere)
            .execute();
        const totalCount = totalCountResult[0]?.totalCount ?? 0;
        const statusResult = await db
            .select({
            successCount: sql `
          COUNT(
            CASE
              WHEN JSON_EXTRACT(${Trace.status}, '$.code') = 1 AND ${Trace.endTime} > 0
              THEN 1
            END
          )
        `,
        })
            .from(Trace)
            .where(baseWhere)
            .execute();
        const successCount = statusResult[0]?.successCount ?? 0;
        const failCount = totalCount - successCount;
        const TotalStatsResult = await db
            .select({
            totalToken: sql `COALESCE(SUM(CASE WHEN ${Trace.endTime} > 0 THEN ${Trace.token} ELSE 0 END), 0)`,
            totalCost: sql `COALESCE(SUM(CASE WHEN ${Trace.endTime} > 0 THEN ${Trace.cost} ELSE 0 END), 0)`,
            maxLatency: sql `COALESCE(MAX(CASE WHEN ${Trace.endTime} > 0 THEN ${Trace.endTime} - ${Trace.startTime} ELSE NULL END), 0)`,
            minLatency: sql `COALESCE(MIN(CASE WHEN ${Trace.endTime} > 0 THEN ${Trace.endTime} - ${Trace.startTime} ELSE NULL END), 0)`,
            avgLatency: sql `COALESCE(AVG(CASE WHEN ${Trace.endTime} > 0 THEN ${Trace.endTime} - ${Trace.startTime} ELSE NULL END), 0)`,
            totalDuration: sql `COALESCE(SUM(CASE WHEN ${Trace.endTime} > 0 THEN ${Trace.endTime} - ${Trace.startTime} ELSE 0 END), 0)`,
        })
            .from(Trace)
            .where(baseWhere)
            .execute();
        const totalStats = TotalStatsResult[0];
        const llmWhere = and(sql `${Trace.attributes} IS NOT NULL`, sql `JSON_EXTRACT(${Trace.attributes}, '$.output.model') IS NOT NULL`);
        const llmStatsResult = await db
            .select({
            llmTotalCount: sql `COUNT(*)`,
            llmSuccessCount: sql `
          COUNT(
            CASE
              WHEN JSON_EXTRACT(${Trace.status}, '$.code') = 1 AND ${Trace.endTime} > 0
              THEN 1
            END
          )
        `,
            llmTotalDuration: sql `
          COALESCE(SUM(
            CASE WHEN ${Trace.endTime} > 0 THEN ${Trace.endTime} - ${Trace.startTime} ELSE 0 END
          ), 0)
        `,
        })
            .from(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.componentId })
            .from(Trace)
            .where(and(isNotNull(Trace.componentId), isNull(Trace.action)))
            .groupBy(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 = and(or(isNull(Trace.parentId), eq(Trace.parentId, "")), isNull(Trace.action));
        const [latestRoot] = (await db
            .select()
            .from(Trace)
            .where(rootFilter)
            .orderBy(desc(Trace.startTime))
            .limit(1)
            .execute()) || [];
        const settingPath = getGlobalSettingPath();
        let settings = {
            lastTrace: { id: "", endTime: 0 },
        };
        if (!existsSync(settingPath)) {
            await writeFile(settingPath, stringify(settings));
        }
        else {
            settings = parse(await readFile(settingPath, "utf8"));
        }
        const lastTraceChanged = latestRoot &&
            (settings.lastTrace?.id !== latestRoot.id ||
                settings.lastTrace?.endTime !== latestRoot.endTime);
        if (lastTraceChanged) {
            await writeFile(settingPath, 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).where(eq(Trace.id, id)).execute();
        if (rootCalls.length === 0)
            throw new Error(`Not found trace: ${id}`);
        const all = await db.select().from(Trace).where(eq(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).where(eq(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.id,
            rootId: Trace.rootId,
            parentId: Trace.parentId,
            name: Trace.name,
            startTime: Trace.startTime,
            endTime: Trace.endTime,
            status: Trace.status,
            attributes: sql `
          CASE
            WHEN JSON_EXTRACT(${Trace.attributes}, '$.output.usage') IS NOT NULL THEN
              JSON_OBJECT(
                'output', JSON_OBJECT(
                  'usage', JSON_EXTRACT(${Trace.attributes}, '$.output.usage'),
                  'model', JSON_EXTRACT(${Trace.attributes}, '$.output.model'),
                  'seconds', JSON_EXTRACT(${Trace.attributes}, '$.output.seconds')
                ),
                'agentTag', JSON_EXTRACT(${Trace.attributes}, '$.agentTag'),
                'metadata', JSON_EXTRACT(${Trace.attributes}, '$.metadata')
              )
            ELSE JSON_OBJECT(
              'output', JSON_OBJECT(),
              'agentTag', JSON_EXTRACT(${Trace.attributes}, '$.agentTag'),
              'metadata', JSON_EXTRACT(${Trace.attributes}, '$.metadata')
            )
          END
        `,
            userId: Trace.userId,
            sessionId: Trace.sessionId,
            componentId: Trace.componentId,
            token: Trace.token,
            cost: Trace.cost,
        })
            .from(Trace)
            .where(inArray(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 = createTraceBatchSchema.parse(req.body);
        let live = false;
        const settingPath = getGlobalSettingPath();
        if (!existsSync(settingPath)) {
            live = false;
        }
        else {
            const setting = parse(await 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 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;
        const ids = req.body?.ids;
        if (Array.isArray(ids) && ids.length > 0) {
            await db.delete(Trace).where(inArray(Trace.id, ids)).execute();
            res.json({ code: 0, message: `${ids.length} traces deleted` });
            return;
        }
        await db.delete(Trace).execute();
        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)
                .set({ remark: remark || null })
                .where(eq(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).where(inArray(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).where(inArray(Trace.rootId, rootIds)).execute()
                : selectedTraces;
            const now = new Date().toISOString().replace(/[:.]/g, "-");
            const archive = archiver("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.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.id,
                rootId: Trace.rootId,
                parentId: Trace.parentId,
            })
                .from(Trace)
                .where(inArray(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 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;
};
