frontend/src/components/PipelineProgress.tsx

import { useEffect, useMemo, useRef, useState } from 'react';
import { CheckCircle2, Loader2, Database, Brain, Sparkles, Rocket, LineChart } from 'lucide-react';
import { MLBiasAPI } from '../services/api';
import { useJobRunner } from '../hooks/JobRunnerProvider';

type Health = {
  status?: string;
  timestamp?: string;
  loaded_models?: string[];
  dataset_loaded?: boolean;
  generation_results_available?: boolean;
  finetune_running?: boolean;
  steps?: Record<string, boolean | 'todo' | 'doing' | 'done'>;
};

type StepKey =
  | 'Activate Task'
  | 'Load Dataset'
  | 'Load Model'
  | 'Generate and Score'
  | 'Counterfactual'
  | 'Sampling'
  | 'Plot and Output'
  | 'Finetune';

type StepState = 'todo' | 'doing' | 'done';

export default function PipelineProgress() {
  const { result, resp } = useJobRunner();

  const [health, setHealth] = useState<Health | null>(null);
  const pollRef = useRef<number | null>(null);

  useEffect(() => {
    const poll = async () => {
      try {
        const h = (await MLBiasAPI.checkHealth()) as Health;
        setHealth(prev => (JSON.stringify(prev) === JSON.stringify(h) ? prev : h));
      } catch {
      }
    };
    void poll();
    pollRef.current = window.setInterval(poll, 1000);
    return () => {
      if (pollRef.current) window.clearInterval(pollRef.current);
    };
  }, []);

  const [elapsed, setElapsed] = useState<number>(0);
  const timerRef = useRef<number | null>(null);
  useEffect(() => {
    const startedAt = Date.now();
    timerRef.current = window.setInterval(() => {
      setElapsed(Math.floor((Date.now() - startedAt) / 1000));
    }, 1000);
    return () => {
      if (timerRef.current) window.clearInterval(timerRef.current);
    };
  }, []);

  const modelName = result?.config?.languageModel || '';

  const wantFT = Boolean(
    result?.config?.enableFineTuning ?? (resp?.results as any)?.config_used?.enableFineTuning
  );

  const backendSteps = useMemo(() => {
    const fromResp = ((resp?.results as any)?.steps || {}) as Record<string, boolean>;
    const fromHealth = ((health?.steps || {}) as Record<string, boolean | 'todo' | 'doing' | 'done'>);
    const merged: Record<string, boolean> = { ...fromResp };
    Object.keys(fromHealth).forEach(k => {
      const v = (fromHealth as any)[k];
      merged[k] = v === true || v === 'doing' || v === 'done';
    });
    return merged;
  }, [health?.steps, resp?.results]);

  const resultsAny = (resp?.results ?? {}) as any;

  const inferred = useMemo(() => {
    const hasData = Boolean(health?.dataset_loaded);
    const hasModel = Boolean(health?.loaded_models && health.loaded_models.length > 0);

    const genDone = Boolean(
      backendSteps['3_generate_and_eval'] ||
      health?.generation_results_available ||
      resultsAny.generation_done
    );

    const r4Flag = Boolean(
      backendSteps['4_rank_sampling_original'] ||
      resultsAny.rank_sampling_original_done ||
      resultsAny.rank_sampling?.original_done
    );

    const r5Flag = Boolean(
      backendSteps['5_rank_sampling_cf'] ||
      resultsAny.rank_sampling_cf_done ||
      resultsAny.rank_sampling?.cf_done
    );

    const plotsFlag = Boolean(
      backendSteps['6_plots_and_metrics'] ||
      resultsAny.plot_urls ||
      resultsAny.plots_ready ||
      (resultsAny.plots &&
        (resultsAny.plots.original_sentiment || resultsAny.plots.counterfactual_sentiment))
    );

    const ftDoneFlag = Boolean(
      backendSteps['7_finetune'] === true ||
      resultsAny.finetune_done ||
      resultsAny.finetune?.completed ||
      resultsAny.finetune?.saved_model_path
    );

    const ftRunning = Boolean(resultsAny.finetune?.running || (health as any)?.finetune_running);

    const noStepSignals =
      Object.keys(backendSteps || {}).length === 0 &&
      !resultsAny.rank_sampling_original_done &&
      !resultsAny.rank_sampling_cf_done &&
      !resultsAny.plots_ready &&
      !resultsAny.finetune_done;

    const cfByTime   = noStepSignals && genDone && elapsed > 30;
    const rsByTime   = noStepSignals && genDone && elapsed > 45;
    const plotsByTime= noStepSignals && genDone && elapsed > 70;

    const cfDone = Boolean(
      backendSteps['3_5_counterfactual'] ||
      resultsAny.counterfactual_done ||
      resultsAny.counterfactual_results ||
      r4Flag || r5Flag || plotsFlag || ftDoneFlag ||
      cfByTime
    );

    const r4 = r4Flag || rsByTime;
    const r5 = r5Flag || rsByTime;
    const plots = plotsFlag || plotsByTime;
    const ftDone = ftDoneFlag;

    return { hasData, hasModel, genDone, cfDone, r4, r5, plots, ftDone, ftRunning };
  }, [backendSteps, health, resultsAny, elapsed]);


  const rawSteps = useMemo<Record<StepKey, StepState>>(() => {
    const states: Record<StepKey, StepState> = {
      'Activate Task': 'todo',
      'Load Dataset': 'todo',
      'Load Model': 'todo',
      'Generate and Score': 'todo',
      'Counterfactual': 'todo',
      'Sampling': 'todo',
      'Plot and Output': 'todo',
      'Finetune': 'todo',
    };

    if (result?.status === 'running') {
      states['Activate Task'] = 'doing';
    }
    if (inferred.hasData) {
      states['Activate Task'] = 'done';
      states['Load Dataset'] = 'done';
    }

    if (inferred.hasModel) {
      states['Load Model'] = 'done';
    } else if (inferred.hasData) {
      states['Load Model'] = 'doing';
    }

    if (inferred.genDone) {
      states['Generate and Score'] = 'done';
    } else if (inferred.hasModel) {
      states['Generate and Score'] = 'doing';
    }

    if (inferred.cfDone) {
      states['Counterfactual'] = 'done';
    } else if (states['Generate and Score'] === 'done') {
      states['Counterfactual'] = 'doing';
    }

    const shouldStartSampling =
      inferred.r4 || inferred.r5 ||
      states['Counterfactual'] === 'done' ||
      (states['Generate and Score'] === 'done' && elapsed > 20);

    if (inferred.r4 && inferred.r5) {
      states['Sampling'] = 'done';
    } else if (shouldStartSampling) {
      states['Sampling'] = 'doing';
    }

    const shouldStartPlotting =
      inferred.plots ||
      states['Sampling'] === 'done' ||
      (states['Sampling'] === 'doing' && elapsed > 40);

    if (inferred.plots) {
      states['Plot and Output'] = 'done';
    } else if (shouldStartPlotting) {
      states['Plot and Output'] = 'doing';
    }

    if (wantFT) {
      if (inferred.ftDone) states['Finetune'] = 'done';
      else if (inferred.ftRunning || states['Plot and Output'] === 'done')
        states['Finetune'] = 'doing';
      else states['Finetune'] = 'todo';
    } else {
      states['Finetune'] = 'todo';
    }

    return states;
  }, [elapsed, inferred, wantFT, result?.status]);

  const STUCK_TIMEOUT = 30; // 秒
  const [enteredAt, setEnteredAt] = useState<Record<StepKey, number>>({} as any);
  const [forcedDone, setForcedDone] = useState<Record<StepKey, boolean>>({} as any);

  useEffect(() => {
    const next: Record<StepKey, number> = { ...enteredAt } as any;
    (Object.keys(rawSteps) as StepKey[]).forEach((k) => {
      if (rawSteps[k] === 'doing' && !next[k]) next[k] = Date.now();
      if (rawSteps[k] !== 'doing' && next[k]) delete next[k];
    });
    if (JSON.stringify(next) !== JSON.stringify(enteredAt)) setEnteredAt(next);
  }, [rawSteps]);

  useEffect(() => {
    const now = Date.now();
    const k: StepKey = 'Counterfactual';
    if (rawSteps[k] === 'doing' && enteredAt[k] && now - enteredAt[k] > STUCK_TIMEOUT * 1000) {
      if (!forcedDone[k]) setForcedDone(prev => ({ ...prev, [k]: true }));
    }
  }, [enteredAt, rawSteps, forcedDone]);

  const steps = useMemo(() => {
    const s = { ...rawSteps } as Record<StepKey, StepState>;
    (Object.keys(forcedDone) as StepKey[]).forEach((k) => {
      if (forcedDone[k]) s[k] = 'done';
    });
    return s;
  }, [rawSteps, forcedDone]);

  const ft = resultsAny?.finetune || {};
  const downloadPath: string | undefined =
    ft.download_url || ft.model_url || ft.saved_model_path || resultsAny?.finetune_model_url;

  const downloadHref = downloadPath ? MLBiasAPI.resolvePath(downloadPath) : undefined;

  const baseSteps = [
    { key: 'Activate Task', icon: Rocket },
    { key: 'Load Dataset', icon: Database },
    { key: 'Load Model', icon: Brain },
    { key: 'Generate and Score', icon: Sparkles },
    { key: 'Counterfactual', icon: Sparkles },
    { key: 'Sampling', icon: LineChart },
    { key: 'Plot and Output', icon: LineChart },
  ] as const;

  const stepList = wantFT
    ? [...baseSteps, { key: 'Finetune', icon: Rocket } as const]
    : baseSteps;

  const completedCount = stepList.reduce(
    (acc, s) => acc + (steps[s.key as StepKey] === 'done' ? 1 : 0),
    0
  );
  const doingCount = stepList.reduce(
    (acc, s) => acc + (steps[s.key as StepKey] === 'doing' ? 1 : 0),
    0
  );
  const percent = Math.min(
    100,
    Math.round(((completedCount + doingCount * 0.5) / stepList.length) * 100)
  );

  const hasStuckStep =
    Object.values(steps).some((state) => state === 'doing') &&
    elapsed > 60 &&
    completedCount < stepList.length - 1;

  return (
    <div className="relative overflow-hidden rounded-2xl border border-indigo-200/50 bg-white/70 backdrop-blur">
      <div className="absolute inset-0 -z-10 bg-[radial-gradient(800px_300px_at_20%_-20%,rgba(99,102,241,0.15),transparent_60%),radial-gradient(800px_300px_at_120%_0%,rgba(244,114,182,0.15),transparent_60%)]" />
      <div className="p-6">
        <div className="flex items-center justify-between mb-4">
          <div>
            <h3 className="text-lg font-semibold text-slate-900">Pipeline Running</h3>
            <p className="text-sm text-slate-600">
              Model: <span className="font-medium text-slate-800">{modelName || '（未指定）'}</span>
            </p>
            {hasStuckStep && (
              <p className="text-xs text-amber-600 mt-1">⚠️ Some steps may run slowly and are automatically attempted to proceed.</p>
            )}
          </div>
          <div className="flex items-center gap-2 text-slate-600">
            <Loader2 className="w-5 h-5 animate-spin" />
            <span className="text-sm">Executed {elapsed}s</span>
          </div>
        </div>

        <div className="w-full h-3 rounded-full bg-slate-200 overflow-hidden">
          <div
            className="h-3 bg-gradient-to-r from-indigo-500 via-violet-500 to-fuchsia-500 transition-all duration-500"
            style={{ width: `${percent}%` }}
          />
        </div>

        <ol className="mt-6 grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 gap-3">
          {stepList.map(({ key, icon: Icon }) => {
            const state = steps[key as StepKey];
            const isDone = state === 'done';
            const isDoing = state === 'doing';
            const startTs = enteredAt[key as StepKey];
            const isStuck = isDoing && startTs && (Date.now() - startTs) / 1000 > STUCK_TIMEOUT;

            return (
              <li
                key={key}
                className={`flex items-center gap-3 rounded-xl border p-3 transition-all duration-300 ${
                  isDone
                    ? 'border-emerald-200 bg-emerald-50'
                    : isDoing
                    ? isStuck
                      ? 'border-amber-300 bg-amber-100'
                      : 'border-amber-200 bg-amber-50'
                    : 'border-slate-200 bg-white/70'
                }`}
              >
                <div
                  className={`rounded-lg p-2 transition-colors ${
                    isDone
                      ? 'bg-emerald-100 text-emerald-700'
                      : isDoing
                      ? isStuck
                        ? 'bg-amber-200 text-amber-800'
                        : 'bg-amber-100 text-amber-700'
                      : 'bg-slate-100 text-slate-600'
                  }`}
                >
                  {isDone ? (
                    <CheckCircle2 className="w-5 h-5" />
                  ) : isDoing ? (
                    <Loader2 className="w-5 h-5 animate-spin" />
                  ) : (
                    <Icon className="w-5 h-5" />
                  )}
                </div>
                <div className="flex-1">
                  <div className="text-sm font-medium text-slate-900">{key}</div>
                  <div className="text-xs text-slate-600">
                    {isDone ? 'Finished' : isDoing ? (isStuck ? 'Running...' : 'Running…') : 'Waiting'}
                  </div>
                </div>
              </li>
            );
          })}
        </ol>

        {/* 微調完成 → 顯示下載模型 */}
        {wantFT && inferred.ftDone && downloadHref && (
          <div className="mt-6">
            <a
              href={downloadHref}
              className="inline-flex items-center rounded-xl border bg-white px-4 py-2 text-sm font-medium text-slate-800 hover:bg-slate-50"
              target="_blank"
              rel="noopener noreferrer"
            >
              <Rocket className="w-4 h-4 mr-2" />
              Download Finetuned Model
            </a>
            {ft?.saved_model_path && (
              <p className="mt-2 text-xs text-slate-500 break-all">
                Model path: {String(ft.saved_model_path)}
              </p>
            )}
          </div>
        )}
      </div>
    </div>
  );
}