frontend/src/hooks/useJobRunner.ts

import { useEffect, useMemo, useRef, useState } from 'react';
import type { JobConfig, JobResult } from '../types';
import type { PipelineResponseDTO } from '../services/api';
import { MLBiasAPI } from '../services/api';

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

type UseJobRunnerReturn = {
  result: JobResult | null;
  resp: PipelineResponseDTO | undefined;
  loading: boolean;
  error?: string;

  start: (config: JobConfig) => Promise<void>;
  cancel: () => void;

  jobId: string | null;
  live: {
    health: HealthLike | null;
    steps: Record<string, boolean>;
    updatedAt: string | null;
    finetuneRunning: boolean;
    progressPercent: number;
  };

  url: typeof MLBiasAPI.resolvePath;
};

export function useJobRunner(): UseJobRunnerReturn {
  const [jobId, setJobId] = useState<string | null>(null);

  const [result, setResult] = useState<JobResult | null>(null);
  const [resp, setResp] = useState<PipelineResponseDTO | undefined>();

  const [loading, setLoading] = useState(false);
  const [error, setErr] = useState<string | undefined>();

  const [health, setHealth] = useState<HealthLike | null>(null);

  const pollRef = useRef<number | null>(null);
  const aliveRef = useRef<boolean>(false);

  const stopPolling = () => {
    if (pollRef.current) {
      window.clearInterval(pollRef.current);
      pollRef.current = null;
    }
    aliveRef.current = false;
  };

  const cancel = () => {
    stopPolling();
    setLoading(false);
  };

  const progressPercent = useMemo(() => {
    const s = (health?.steps as Record<string, boolean | string>) || {};
    const keys = Object.keys(s);
    if (keys.length === 0) return result?.progress ?? 0;
    let score = 0;
    keys.forEach((k) => {
      const v = s[k];
      if (v === true || v === 'done') score += 1;
      else if (v === 'doing') score += 0.5;
    });
    return Math.max(0, Math.min(100, Math.round((score / keys.length) * 100)));
  }, [health?.steps, result?.progress]);

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

  const pollOnce = async () => {
    try {
      const h = (await MLBiasAPI.checkHealth()) as HealthLike;
      setHealth((prev) => (JSON.stringify(prev) === JSON.stringify(h) ? prev : h));

      const steps = (h?.steps || {}) as Record<string, boolean | string>;
      const plotsDone =
        !!steps['6_plots_and_metrics'] ||
        (resp?.results as any)?.plots_ready ||
        ((resp?.results as any)?.plot_urls?.length ?? 0) > 0;

      const r4 = !!steps['4_rank_sampling_original'];
      const r5 = !!steps['5_rank_sampling_cf'];
      const samplingDone = r4 && r5;

      const genAvailable = !!h?.generation_results_available;
      const ftMaybeDone =
        !!steps['7_finetune'] ||
        (resp?.results as any)?.finetune_done ||
        (resp?.results as any)?.finetune?.completed;

      const declaredCompleted = h?.completed === true || h?.status === 'completed';

      if (declaredCompleted || plotsDone || samplingDone || (genAvailable && ftMaybeDone)) {
        stopPolling();
        setLoading(false);
      }
    } catch (e: any) {
      setErr((e && e.message) || String(e));
    }
  };

  const start = async (config: JobConfig) => {
    setLoading(true);
    setErr(undefined);

    const now = new Date().toISOString();
    const provisionalId = crypto.randomUUID();
    setResult({
      id: provisionalId,
      status: 'running',
      progress: 0,
      config,
      createdAt: now,
      updatedAt: now,
    });
    setResp(undefined);
    setHealth(null);

    try {
      const runResp: any = await MLBiasAPI.runPipeline(config);

      const jid: string | undefined =
        runResp?.jobId || runResp?.job_id || runResp?.results?.jobId || runResp?.results?.job_id;
      setJobId(jid || provisionalId);

      if (runResp?.results?.metrics) {
        const final = runResp as PipelineResponseDTO;
        const now2 = new Date().toISOString();

        setResp(final);
        setResult({
          id: jid || provisionalId,
          status: 'completed',
          progress: 100,
          config,
          createdAt: now,
          updatedAt: now2,
          completedAt: now2,
          metrics: {
            finalMeanDiff: final.results.metrics.finalMeanDiff,
            reductionPct: final.results.metrics.reductionPct ?? 0,
            stableCoverage: final.results.metrics.stableCoverage ?? 100,
          },
        });
        setLoading(false);
        return;
      }

      aliveRef.current = true;
      await pollOnce(); 
      if (aliveRef.current) {
        pollRef.current = window.setInterval(pollOnce, 1000);
      }
    } catch (e: any) {
      setErr(e?.message || String(e));
      setResult((prev) =>
        prev
          ? { ...prev, status: 'failed', progress: 100, updatedAt: new Date().toISOString() }
          : null
      );
      setLoading(false);
    }
  };

  useEffect(() => stopPolling, []);

  const url = MLBiasAPI.resolvePath;

  return {
    result,
    resp,
    loading,
    error,
    start,
    cancel,
    jobId,
    live: {
      health,
      steps: liveSteps,
      updatedAt: (health && (health.updated_at || health.timestamp)) || null,
      finetuneRunning: !!(health?.finetune_running || (resp as any)?.results?.finetune?.running),
      progressPercent,
    },
    url,
  };
}