"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.TraceRepository = void 0;
const db_js_1 = require("../db/db.js");
const schema_js_1 = require("../db/schema.js");
const drizzle_orm_1 = require("drizzle-orm");
const uuid_1 = require("uuid");
class TraceRepository {
    // Trace operations
    async createTrace(data, tsx = null) {
        const drizzle = tsx ? tsx : (await db_js_1.db.getDrizzle());
        const id = (0, uuid_1.v4)();
        const createdAt = new Date().toISOString();
        const traceData = {
            id,
            projectId: data.projectId,
            name: data.name,
            startTime: data.startTime,
            endTime: data.endTime,
            status: 'running', // Default to 'running' if not provided in dat,
            metadata: JSON.stringify(data.metadata || {}),
            createdAt,
        };
        let createdSpans = [];
        const res = await drizzle.transaction(async (tsx1) => {
            await tsx1.insert(schema_js_1.traces).values(traceData);
            if (data.spans && data.spans.length > 0) {
                for (const span of data.spans) {
                    const newSpan = await this.createSpan({ ...span, traceId: id }, tsx1);
                    createdSpans.push(newSpan);
                }
            }
            return {
                ...traceData,
                metadata: data.metadata || {},
                spans: createdSpans // Add the spans to the response
            };
        });
        return res;
    }
    async getTraceById(id) {
        const drizzle = await db_js_1.db.getDrizzle();
        const result = await drizzle.select().from(schema_js_1.traces).where((0, drizzle_orm_1.eq)(schema_js_1.traces.id, id)).get();
        if (!result)
            return undefined;
        return {
            id: result.id,
            projectId: result.projectId,
            name: result.name,
            startTime: result.startTime,
            endTime: result.endTime || undefined,
            duration: result.duration || undefined,
            status: result.status,
            metadata: JSON.parse(result.metadata),
            createdAt: result.createdAt,
        };
    }
    async getTracesByProjectId(projectId, options = {}) {
        const drizzle = await db_js_1.db.getDrizzle();
        const { limit = 50, sortDirection = 'desc', cursor } = options;
        const fetchLimitPlusOne = limit + 1; // Fetch one extra to determine hasNextPage
        // Build conditions
        const conditions = [(0, drizzle_orm_1.eq)(schema_js_1.traces.projectId, projectId)];
        // Add cursor conditions for pagination
        if (cursor) {
            const cursorTimestamp = cursor.timestamp;
            const cursorId = cursor.id;
            if (sortDirection === 'desc') {
                conditions.push((0, drizzle_orm_1.or)((0, drizzle_orm_1.lt)(schema_js_1.traces.startTime, cursorTimestamp), (0, drizzle_orm_1.and)((0, drizzle_orm_1.eq)(schema_js_1.traces.startTime, cursorTimestamp), (0, drizzle_orm_1.lt)(schema_js_1.traces.id, cursorId))));
            }
            else { // asc
                conditions.push((0, drizzle_orm_1.or)((0, drizzle_orm_1.gt)(schema_js_1.traces.startTime, cursorTimestamp), (0, drizzle_orm_1.and)((0, drizzle_orm_1.eq)(schema_js_1.traces.startTime, cursorTimestamp), (0, drizzle_orm_1.gt)(schema_js_1.traces.id, cursorId))));
            }
        }
        const orderBy = sortDirection === 'desc'
            ? [(0, drizzle_orm_1.desc)(schema_js_1.traces.startTime), (0, drizzle_orm_1.desc)(schema_js_1.traces.id)]
            : [(0, drizzle_orm_1.asc)(schema_js_1.traces.startTime), (0, drizzle_orm_1.asc)(schema_js_1.traces.id)];
        // First, get the traces with pagination
        const traceResults = await drizzle
            .select()
            .from(schema_js_1.traces)
            .where((0, drizzle_orm_1.and)(...conditions))
            .orderBy(...orderBy)
            .limit(fetchLimitPlusOne)
            .all();
        // Determine if there's a next page
        const hasNextPage = traceResults.length > limit;
        // Return the first `limit` traces
        const tracesToReturn = hasNextPage ? traceResults.slice(0, limit) : traceResults;
        // Get all spans for the returned traces in a single query
        const traceIds = tracesToReturn.map(trace => trace.id);
        const spansResults = traceIds.length > 0 ? await drizzle
            .select()
            .from(schema_js_1.spans)
            .where((0, drizzle_orm_1.inArray)(schema_js_1.spans.traceId, traceIds))
            .orderBy((0, drizzle_orm_1.asc)(schema_js_1.spans.startTime))
            .all() : [];
        // Convert raw spans to SpanData format
        const allSpans = spansResults.map(span => ({
            id: span.id,
            traceId: span.traceId,
            parentSpanId: span.parentSpanId || undefined,
            name: span.name,
            startTime: span.startTime,
            endTime: span.endTime || undefined,
            status: span.status,
            metadata: JSON.parse(span.metadata),
            createdAt: span.createdAt,
            error: span.error ? JSON.parse(span.error) : undefined,
        }));
        // Group spans by trace ID and build hierarchy for each trace
        const spansByTraceId = allSpans.reduce((acc, span) => {
            if (!acc[span.traceId]) {
                acc[span.traceId] = [];
            }
            acc[span.traceId].push(span);
            return acc;
        }, {});
        // Build hierarchy for each trace's spans
        Object.keys(spansByTraceId).forEach(traceId => {
            spansByTraceId[traceId] = this.buildSpanHierarchy(spansByTraceId[traceId]);
        });
        return {
            traces: tracesToReturn.map(result => ({
                id: result.id,
                projectId: result.projectId,
                name: result.name,
                startTime: result.startTime,
                endTime: result.endTime || undefined,
                duration: result.duration || undefined,
                status: result.status,
                metadata: JSON.parse(result.metadata),
                createdAt: result.createdAt,
                spans: spansByTraceId[result.id] || [],
            })),
            hasNextPage,
        };
    }
    async updateTrace(id, data) {
        const drizzle = await db_js_1.db.getDrizzle();
        const updateData = {};
        if (data.endTime !== undefined)
            updateData.endTime = data.endTime;
        if (data.duration !== undefined)
            updateData.duration = data.duration;
        if (data.status !== undefined)
            updateData.status = data.status;
        if (data.metadata !== undefined)
            updateData.metadata = JSON.stringify(data.metadata);
        await drizzle.update(schema_js_1.traces).set(updateData).where((0, drizzle_orm_1.eq)(schema_js_1.traces.id, id));
        return this.getTraceById(id);
    }
    async deleteTrace(id) {
        const drizzle = await db_js_1.db.getDrizzle();
        await drizzle.delete(schema_js_1.traces).where((0, drizzle_orm_1.eq)(schema_js_1.traces.id, id));
    }
    buildSpanHierarchy(spans) {
        // Create a map for quick lookup
        const spanMap = new Map();
        const rootSpans = [];
        // Initialize all spans in the map with empty spans array
        spans.forEach(span => {
            spanMap.set(span.id, { ...span, spans: [] });
        });
        // Build the hierarchy
        spans.forEach(span => {
            const spanWithChildren = spanMap.get(span.id);
            if (span.parentSpanId) {
                // This span has a parent, add it to parent's spans array
                const parent = spanMap.get(span.parentSpanId);
                if (parent) {
                    parent.spans.push(spanWithChildren);
                }
                else {
                    // Parent not found, treat as root
                    rootSpans.push(spanWithChildren);
                }
            }
            else {
                // This is a root span
                rootSpans.push(spanWithChildren);
            }
        });
        return rootSpans;
    }
    async bulkTraceInsert(tracesData) {
        const drizzle = await db_js_1.db.getDrizzle();
        return await drizzle.transaction(async (tsx) => {
            const createdTraces = [];
            for (const traceData of tracesData) {
                const createdTrace = await this.createTrace(traceData, tsx);
                createdTraces.push(createdTrace);
            }
            return createdTraces;
        });
    }
    // Span operations
    async createSpan(data, tsx = null) {
        const drizzle = tsx ? tsx : (await db_js_1.db.getDrizzle());
        const createdAt = new Date().toISOString();
        const spanData = {
            id: data.id,
            traceId: data.traceId,
            parentSpanId: data.parentSpanId || null,
            name: data.name,
            startTime: data.startTime,
            endTime: data.endTime,
            error: data.error ? JSON.stringify(data.error) : null,
            metadata: JSON.stringify(data.metadata || {}),
            createdAt,
            status: data.status ?? "running"
        };
        await drizzle.insert(schema_js_1.spans).values({
            ...spanData,
        });
        return {
            ...spanData,
            parentSpanId: spanData.parentSpanId || undefined,
            metadata: data.metadata || {},
            error: data.error
        };
    }
    async getSpanById(id) {
        const drizzle = await db_js_1.db.getDrizzle();
        const result = await drizzle.select().from(schema_js_1.spans).where((0, drizzle_orm_1.eq)(schema_js_1.spans.id, id)).get();
        if (!result)
            return undefined;
        return {
            id: result.id,
            traceId: result.traceId,
            parentSpanId: result.parentSpanId || undefined,
            name: result.name,
            startTime: result.startTime,
            endTime: result.endTime || undefined,
            status: result.status,
            metadata: JSON.parse(result.metadata),
            createdAt: result.createdAt,
            error: result.error ? JSON.parse(result.error) : undefined,
        };
    }
    async getSpansByTraceId(traceId, options = {}) {
        const drizzle = await db_js_1.db.getDrizzle();
        const { limit = 100, sortDirection = 'asc', cursor } = options;
        const fetchLimitPlusOne = limit + 1; // Fetch one extra to determine hasNextPage
        // Build conditions
        const conditions = [(0, drizzle_orm_1.eq)(schema_js_1.spans.traceId, traceId)];
        // Add cursor conditions for pagination
        if (cursor) {
            const cursorTimestamp = cursor.timestamp;
            const cursorId = cursor.id;
            if (sortDirection === 'desc') {
                conditions.push((0, drizzle_orm_1.or)((0, drizzle_orm_1.lt)(schema_js_1.spans.startTime, cursorTimestamp), (0, drizzle_orm_1.and)((0, drizzle_orm_1.eq)(schema_js_1.spans.startTime, cursorTimestamp), (0, drizzle_orm_1.lt)(schema_js_1.spans.id, cursorId))));
            }
            else { // asc
                conditions.push((0, drizzle_orm_1.or)((0, drizzle_orm_1.gt)(schema_js_1.spans.startTime, cursorTimestamp), (0, drizzle_orm_1.and)((0, drizzle_orm_1.eq)(schema_js_1.spans.startTime, cursorTimestamp), (0, drizzle_orm_1.gt)(schema_js_1.spans.id, cursorId))));
            }
        }
        const orderBy = sortDirection === 'desc'
            ? [(0, drizzle_orm_1.desc)(schema_js_1.spans.startTime), (0, drizzle_orm_1.desc)(schema_js_1.spans.id)]
            : [(0, drizzle_orm_1.asc)(schema_js_1.spans.startTime), (0, drizzle_orm_1.asc)(schema_js_1.spans.id)];
        const results = await drizzle
            .select()
            .from(schema_js_1.spans)
            .where((0, drizzle_orm_1.and)(...conditions))
            .orderBy(...orderBy)
            .limit(fetchLimitPlusOne)
            .all();
        // Determine if there's a next page
        const hasNextPage = results.length > limit;
        // Return the first `limit` spans
        const spansToReturn = hasNextPage ? results.slice(0, limit) : results;
        return {
            spans: spansToReturn.map(result => ({
                id: result.id,
                traceId: result.traceId,
                parentSpanId: result.parentSpanId || undefined,
                name: result.name,
                startTime: result.startTime,
                endTime: result.endTime || undefined,
                status: result.status,
                metadata: JSON.parse(result.metadata),
                createdAt: result.createdAt,
                error: result.error ? JSON.parse(result.error) : undefined,
            })),
            hasNextPage,
        };
    }
    async getSpansByParentId(parentSpanId) {
        const drizzle = await db_js_1.db.getDrizzle();
        const results = await drizzle
            .select()
            .from(schema_js_1.spans)
            .where((0, drizzle_orm_1.eq)(schema_js_1.spans.parentSpanId, parentSpanId))
            .orderBy((0, drizzle_orm_1.asc)(schema_js_1.spans.startTime))
            .all();
        return results.map(result => ({
            id: result.id,
            traceId: result.traceId,
            parentSpanId: result.parentSpanId || undefined,
            name: result.name,
            startTime: result.startTime,
            endTime: result.endTime || undefined,
            status: result.status,
            metadata: JSON.parse(result.metadata),
            createdAt: result.createdAt,
            error: result.error ? JSON.parse(result.error) : null,
        }));
    }
    async updateSpan(id, data) {
        const drizzle = await db_js_1.db.getDrizzle();
        const updateData = {};
        if (data.endTime !== undefined)
            updateData.endTime = data.endTime;
        if (data.duration !== undefined)
            updateData.duration = data.duration;
        if (data.status !== undefined)
            updateData.status = data.status;
        if (data.metadata !== undefined)
            updateData.metadata = JSON.stringify(data.metadata);
        if (data.error !== undefined)
            updateData.error = data.error ? JSON.stringify(data.error) : null;
        await drizzle.update(schema_js_1.spans).set(updateData).where((0, drizzle_orm_1.eq)(schema_js_1.spans.id, id));
        return this.getSpanById(id);
    }
    async deleteSpan(id) {
        const drizzle = await db_js_1.db.getDrizzle();
        await drizzle.delete(schema_js_1.spans).where((0, drizzle_orm_1.eq)(schema_js_1.spans.id, id));
    }
    // Utility methods
    async getTraceWithSpans(traceId) {
        const trace = await this.getTraceById(traceId);
        if (!trace)
            return undefined;
        const spansResult = await this.getSpansByTraceId(traceId);
        return { trace, spans: spansResult.spans };
    }
    async clearProjectTraces(projectId) {
        const drizzle = await db_js_1.db.getDrizzle();
        await drizzle.delete(schema_js_1.traces).where((0, drizzle_orm_1.eq)(schema_js_1.traces.projectId, projectId));
    }
}
exports.TraceRepository = TraceRepository;
//# sourceMappingURL=traceRepository.js.map