khata-ui/src/FetchRequest/components/PipelineStepper.tsx

import React, { useMemo } from "react";
import {
  Box,
  Paper,
  Typography,
  Stepper,
  Step,
  StepLabel,
  LinearProgress,
} from "@mui/material";
import CheckCircleIcon from "@mui/icons-material/CheckCircle";
import ErrorIcon from "@mui/icons-material/Error";
import WarningAmberIcon from "@mui/icons-material/WarningAmber";
import CircularProgress from "@mui/material/CircularProgress";
import type { FetchRequestStatus, SSEEvent, ProgressMessage } from "../../features/fetch-requests";

const STEP_LABELS = ["Extract", "Raw Expense", "Enrich", "Save"];

function computeProgressPercent(
  status: FetchRequestStatus,
  liveCount: number,
  seenSteps: Set<string>,
  stepStats: Record<string, number>,
  txnBlockCount: number,
  txnDictCount: number,
): number {
  if (status === "pending") return 0;
  if (status === "completed") return 100;

  let pct = 0;

  if (seenSteps.has("raw_lines") || seenSteps.has("txn_blocks")) pct += 10;

  if (txnBlockCount > 0) {
    const current = Math.max(liveCount, stepStats.txn_dicts ?? 0);
    pct += Math.min(1, current / txnBlockCount) * 20;
  }

  if (txnDictCount > 0) {
    pct += Math.min(1, (stepStats.enrich_count ?? 0) / txnDictCount) * 50;
    pct += Math.min(1, (stepStats.save_count ?? 0) / txnDictCount) * 20;
  }

  return Math.round(Math.min(100, pct));
}

function computeActiveStep(status: FetchRequestStatus, seenSteps: Set<string>): number {
  if (status === "completed") return STEP_LABELS.length;

  if (seenSteps.has("save_expenses/completed") || seenSteps.has("complete/completed")) return STEP_LABELS.length;
  if (seenSteps.has("save_expenses") || seenSteps.has("complete")) return 3;

  if (seenSteps.has("enrich/completed")) return 3;
  if (seenSteps.has("enrich")) return 2;

  if (seenSteps.has("txn_dicts/completed") || status === "raw_expenses_done") return 2;
  if (seenSteps.has("txn_dicts")) return 1;

  if (seenSteps.has("txn_blocks/completed")) return 1;
  if (seenSteps.has("raw_lines") || seenSteps.has("txn_blocks")) return 0;

  if (status === "processing" || status === "paused") return 0;

  return -1;
}

function computeStepMessages(
  fetchRequest: any,
  stepStats: Record<string, number>,
  liveParsedCount: number,
  txnBlockCount: number,
): Record<number, string> {
  const msgs: Record<number, string> = {};
  const source = fetchRequest?.source;

  const rawLineCount = stepStats.raw_lines ?? (source?.raw_lines?.length ?? 0);
  if (rawLineCount) msgs[0] = `${rawLineCount}`;

  const sourceDictCount = source?.txn_dict_count ?? source?.txn_dicts_count ?? 0;
  const dictLive = liveParsedCount ?? stepStats.txn_dicts ?? 0;
  const dictCurrent = Math.max(dictLive, sourceDictCount);
  if (dictCurrent && txnBlockCount) msgs[1] = `${dictCurrent}/${txnBlockCount}`;
  else if (dictCurrent) msgs[1] = `${dictCurrent}`;

  const txnDictDenom = stepStats.txn_dicts ?? sourceDictCount;
  if (stepStats.enrich_count && txnDictDenom) msgs[2] = `${stepStats.enrich_count}/${txnDictDenom}`;
  else if (stepStats.enrich_count) msgs[2] = `${stepStats.enrich_count}`;

  if (stepStats.save_count && txnDictDenom) msgs[3] = `${stepStats.save_count}/${txnDictDenom}`;
  else if (stepStats.save_count) msgs[3] = `${stepStats.save_count}`;

  return msgs;
}

interface PipelineStepperProps {
  fetchRequest: any;
  sseEvents: SSEEvent[];
  stepStats: Record<string, number>;
  liveParsedCount: number;
}

export function PipelineStepper({ fetchRequest, sseEvents, stepStats, liveParsedCount }: PipelineStepperProps) {
  const status = (fetchRequest?.status ?? "pending") as FetchRequestStatus;

  const seenSteps = useMemo(() => {
    const steps = new Set<string>();
    for (const evt of sseEvents) {
      steps.add(evt.step);
      if (evt.status === "completed") steps.add(`${evt.step}/completed`);
      if (evt.status === "failed") steps.add(`${evt.step}/failed`);
      if (evt.status === "started") steps.add(`${evt.step}/started`);
      if (evt.status === "progress") steps.add(`${evt.step}/progress`);
    }
    return steps;
  }, [sseEvents]);

  const txnBlockCount = useMemo(() => {
    const blocks = fetchRequest?.source?.txn_blocks;
    if (!blocks) return 0;
    return Object.values(blocks).reduce(
      (sum: number, list: any) => sum + (Array.isArray(list) ? list.length : 0),
      0,
    );
  }, [fetchRequest]);

  const txnDictCount = useMemo(() => {
    const source = fetchRequest?.source;
    if (stepStats.txn_dicts && stepStats.txn_dicts > 0) return stepStats.txn_dicts;
    return source?.txn_dict_count ?? source?.txn_dicts_count ?? 0;
  }, [fetchRequest, stepStats]);

  const displayParsedCount = useMemo(() => {
    if (liveParsedCount && liveParsedCount > 0) return liveParsedCount;
    const source = fetchRequest?.source;
    const persistedCount = source?.txn_dict_count ?? source?.txn_dicts_count ?? 0;
    if (persistedCount > 0) return persistedCount;
    const dicts = source?.txn_dicts;
    if (Array.isArray(dicts) && dicts.length > 0) return dicts.length;
    return 0;
  }, [liveParsedCount, fetchRequest]);

  const activeStep = computeActiveStep(status, seenSteps);
  const progressPercent = computeProgressPercent(status, displayParsedCount, seenSteps, stepStats, txnBlockCount, txnDictCount);
  const stepMessages = computeStepMessages(fetchRequest, stepStats, liveParsedCount, txnBlockCount);

  return (
    <Paper sx={{ p: 3, borderRadius: 4, mb: 3 }} variant="outlined">
      <Typography variant="subtitle1" fontWeight={600} gutterBottom>
        Pipeline Progress
      </Typography>
      <Stepper activeStep={activeStep} alternativeLabel>
        {STEP_LABELS.map((label, index) => {
          const isCompleted = index < activeStep;
          const isActive = index === activeStep;
          const isPaused = status === "paused" && isActive;
          const isFailed = status === "failed" && isActive;

          let icon: React.ReactNode;
          if (isCompleted) {
            icon = <CheckCircleIcon sx={{ color: "success.main" }} />;
          } else if (isFailed) {
            icon = <ErrorIcon sx={{ color: "error.main" }} />;
          } else if (isPaused) {
            icon = <WarningAmberIcon sx={{ color: "warning.main" }} />;
          } else if (isActive) {
            icon = <CircularProgress size={20} />;
          } else {
            icon = <Typography variant="caption" color="text.disabled">{index + 1}</Typography>;
          }

          return (
            <Step key={label}>
              <StepLabel
                StepIconComponent={() => <Box sx={{ display: "flex", alignItems: "center" }}>{icon}</Box>}
              >
                <Typography variant="body2" fontWeight={600}>{label}</Typography>
                {stepMessages[index] && (
                  <Typography variant="caption" color="text.secondary" sx={{ display: "block", lineHeight: 1.2 }}>
                    {stepMessages[index]}
                  </Typography>
                )}
              </StepLabel>
            </Step>
          );
        })}
      </Stepper>

      <Box sx={{ mt: 3, mb: 1 }}>
        <Box sx={{ display: "flex", alignItems: "center", gap: 2, mb: 0.5 }}>
          <Typography variant="caption" color="text.secondary">Overall Progress</Typography>
          {["processing", "paused"].includes(status) && displayParsedCount > 0 && (
            <Typography variant="caption" fontWeight={600} color="info.main">
              Validated: {displayParsedCount} transactions
            </Typography>
          )}
        </Box>
        <LinearProgress
          variant="determinate"
          value={progressPercent}
          color={status === "failed" ? "error" : status === "completed" ? "success" : "primary"}
          sx={{ borderRadius: 1, height: 8, transition: "width 0.3s ease" }}
        />
        <Typography variant="caption" color="text.secondary" sx={{ mt: 0.25, display: "block" }}>
          {progressPercent}%
        </Typography>
      </Box>
    </Paper>
  );
}