frontend/src/lib/trend-data-generator.ts

import { Trace } from "@/types";

export interface TrendDataPoint {
  time: string;
  value: number;
  value2?: number; // Optional second value for dual-line charts
  date: Date;
}

export interface TrendData {
  tracesAndGraphs: TrendDataPoint[]; // Combined traces + graphs
  entitiesAndRelations: TrendDataPoint[]; // Combined entities + relations
  failures: TrendDataPoint[]; // New failures trend
  // Legacy single values for backwards compatibility
  traces: TrendDataPoint[];
  agentGraphs: TrendDataPoint[];
  entities: TrendDataPoint[];
  relations: TrendDataPoint[];
}

function createRealTrendData(
  traces: Trace[],
  valueExtractor: (traces: Trace[], date: Date) => number,
  valueExtractor2?: (traces: Trace[], date: Date) => number
): TrendDataPoint[] {
  const data: TrendDataPoint[] = [];
  const now = new Date();

  // Get all upload dates from traces
  const traceDates = traces
    .map((trace) =>
      trace.upload_timestamp ? new Date(trace.upload_timestamp) : null
    )
    .filter((date): date is Date => date !== null)
    .sort((a, b) => a.getTime() - b.getTime());

  if (traceDates.length === 0) {
    // If no traces, return empty data
    return [];
  }

  // Create data points for the last 7 days
  for (let i = 6; i >= 0; i--) {
    const date = new Date(now);
    date.setDate(date.getDate() - i);
    date.setHours(23, 59, 59, 999); // End of day for cumulative counting

    const value = valueExtractor(traces, date);
    const value2 = valueExtractor2 ? valueExtractor2(traces, date) : undefined;

    data.push({
      time: date.toLocaleDateString("en-US", {
        month: "short",
        day: "numeric",
      }),
      value: value,
      value2: value2,
      date: date,
    });
  }

  return data;
}

function createCombinedRealTrendData(
  traces: Trace[],
  valueExtractor1: (traces: Trace[], date: Date) => number,
  valueExtractor2: (traces: Trace[], date: Date) => number
): TrendDataPoint[] {
  return createRealTrendData(traces, valueExtractor1, valueExtractor2);
}

// Helper function to extract failure count from traces
export function calculateFailureCount(traces: Trace[]): number {
  let totalFailures = 0;

  traces.forEach((trace) => {
    trace.knowledge_graphs?.forEach((kg: any) => {
      if (kg.graph_data) {
        try {
          const graphData =
            typeof kg.graph_data === "string"
              ? JSON.parse(kg.graph_data)
              : kg.graph_data;

          // Count failures
          if (graphData.failures && Array.isArray(graphData.failures)) {
            totalFailures += graphData.failures.length;
          }
        } catch (error) {
          console.warn("Error parsing graph_data for failure count:", error);
        }
      }
    });
  });

  return totalFailures;
}

// Value extractors for different metrics - CUMULATIVE COUNTS
const getTracesUpToDate = (traces: Trace[], date: Date): number => {
  return traces.filter((trace) => {
    if (!trace.upload_timestamp) return false;
    const uploadDate = new Date(trace.upload_timestamp);
    return uploadDate <= date; // Cumulative: count all traces up to this date
  }).length;
};

const getKnowledgeGraphsUpToDate = (traces: Trace[], date: Date): number => {
  return traces.reduce((count, trace) => {
    const finalKGs =
      trace.knowledge_graphs?.filter((kg) => {
        if (!kg.created_at) return false;
        const creationDate = new Date(kg.created_at);
        const isBeforeOrOnDate = creationDate <= date; // Cumulative
        const isFinal =
          kg.is_final === true ||
          (kg.window_index === null && kg.window_total !== null);
        return isBeforeOrOnDate && isFinal;
      }) || [];
    return count + finalKGs.length;
  }, 0);
};

const getEntitiesUpToDate = (traces: Trace[], date: Date): number => {
  return traces.reduce((count, trace) => {
    const entityCount =
      trace.knowledge_graphs?.reduce((kgTotal, kg) => {
        if (!kg.created_at) return kgTotal;
        const creationDate = new Date(kg.created_at);
        const isBeforeOrOnDate = creationDate <= date; // Cumulative
        return isBeforeOrOnDate ? kgTotal + (kg.entity_count || 0) : kgTotal;
      }, 0) || 0;
    return count + entityCount;
  }, 0);
};

const getRelationsUpToDate = (traces: Trace[], date: Date): number => {
  return traces.reduce((count, trace) => {
    const relationCount =
      trace.knowledge_graphs?.reduce((kgTotal, kg) => {
        if (!kg.created_at) return kgTotal;
        const creationDate = new Date(kg.created_at);
        const isBeforeOrOnDate = creationDate <= date; // Cumulative
        return isBeforeOrOnDate ? kgTotal + (kg.relation_count || 0) : kgTotal;
      }, 0) || 0;
    return count + relationCount;
  }, 0);
};

const getFailuresUpToDate = (traces: Trace[], date: Date): number => {
  return traces.reduce((count, trace) => {
    let traceFailures = 0;
    trace.knowledge_graphs?.forEach((kg) => {
      if (!kg.created_at) return;
      const creationDate = new Date(kg.created_at);
      const isBeforeOrOnDate = creationDate <= date; // Cumulative

      if (isBeforeOrOnDate && kg.graph_data) {
        try {
          const graphData =
            typeof kg.graph_data === "string"
              ? JSON.parse(kg.graph_data)
              : kg.graph_data;

          if (graphData.failures && Array.isArray(graphData.failures)) {
            traceFailures += graphData.failures.length;
          }
        } catch (error) {
          console.warn("Error parsing graph_data for failure count:", error);
        }
      }
    });
    return count + traceFailures;
  }, 0);
};

export function generateTrendData(traces: Trace[]): TrendData {
  return {
    // New combined data for merged cards
    tracesAndGraphs: createCombinedRealTrendData(
      traces,
      getTracesUpToDate,
      getKnowledgeGraphsUpToDate
    ),
    entitiesAndRelations: createCombinedRealTrendData(
      traces,
      getEntitiesUpToDate,
      getRelationsUpToDate
    ),
    failures: createRealTrendData(traces, getFailuresUpToDate),

    // Legacy single values for backwards compatibility
    traces: createRealTrendData(traces, getTracesUpToDate),
    agentGraphs: createRealTrendData(traces, getKnowledgeGraphsUpToDate),
    entities: createRealTrendData(traces, getEntitiesUpToDate),
    relations: createRealTrendData(traces, getRelationsUpToDate),
  };
}