import { LogDto, TraceWithLogsDto } from "@montelo/browser-client";

export interface Row {
  type: "trace" | "log";
  traceId: string;
  id: string;
  name: string | null;
  startTime: string | null;
  latency: number | null;
  totalCost: number | null;
  totalTokens: number | null;
  model: string | null;
  subRows: Row[];
  depth: number;
}

export const formatTraces = (traces: TraceWithLogsDto[]): Row[] => {
  return traces.map((trace) => {
    const logMap: { [key: string]: Row[] } = {};

    trace.logs.forEach((log) => {
      const row = formatLog(log);
      if (log.parentLogId === null) {
        if (!logMap[trace.id]) {
          logMap[trace.id] = [];
        }
        logMap[trace.id].push(row);
      } else {
        if (!logMap[log.parentLogId]) {
          logMap[log.parentLogId] = [];
        }
        logMap[log.parentLogId].push(row);
      }
    });

    const assignSubRows = (parentRows: Row[], depth: number) => {
      parentRows.forEach((row) => {
        row.depth = depth;
        if (logMap[row.id]) {
          row.subRows = logMap[row.id];
          assignSubRows(row.subRows, depth + 1);
        }
      });
    };

    const topLevelRows = logMap[trace.id] || [];
    assignSubRows(topLevelRows, 2);

    return {
      type: "trace",
      traceId: trace.id,
      id: trace.id,
      name: trace.name,
      startTime: trace.startTime,
      latency: trace.duration,
      totalCost: trace.totalCost,
      totalTokens: trace.totalTokens,
      model: null,
      subRows: topLevelRows,
      depth: 1,
    };
  });
};

const formatLog = (log: LogDto): Row => {
  return {
    type: "log",
    traceId: log.traceId,
    id: log.id,
    name: log.name,
    startTime: log.startTime || log.createdAt,
    latency: log.duration,
    totalCost: log.totalCost,
    totalTokens: log.totalTokens,
    model: log.model,
    subRows: [],
    depth: 0,
  };
};
