import { useState, useEffect, useRef, useCallback } from 'react';
import type { ChatMessage } from '../../../Chat';
import { randomId } from 'helpers/random-id';

type Phase = 'ready' | 'running' | 'done';

type StepDef = {
  task: string;
  foreach?: string;
  concurrency?: number;
};

type StepStatus = {
  taskName: string;
  iteration?: { current: number; total: number; label: string };
  status: 'pending' | 'running' | 'complete' | 'skipped';
  cost?: { inputTokens: number; outputTokens: number; totalUSD: number };
};

export type IterationStatus = {
  label: string;
  status: 'pending' | 'running' | 'complete' | 'error';
  error?: string;
  cost?: { inputTokens: number; outputTokens: number; totalUSD: number };
};

type ParallelStep = {
  stepIndex: number;
  taskName: string;
  concurrency: number;
  iterations: IterationStatus[];
};

type ServerMessage =
  | { jobId: string; type: 'pipeline:init'; steps: StepDef[] }
  | { jobId: string; type: 'step:start'; stepIndex: number; taskName: string; iteration?: { current: number; total: number; label: string } }
  | { jobId: string; type: 'step:complete'; stepIndex: number; cost?: { inputTokens: number; outputTokens: number; totalUSD: number } }
  | { jobId: string; type: 'step:skip'; stepIndex: number; label: string; reason: string }
  | { jobId: string; type: 'step:parallel'; stepIndex: number; taskName: string; iterations: string[]; concurrency: number }
  | { jobId: string; type: 'step:waiting'; stepIndex: number; iterationLabel?: string; elapsed: number }
  | { jobId: string; type: 'iteration:start'; stepIndex: number; label: string }
  | { jobId: string; type: 'iteration:complete'; stepIndex: number; label: string; cost?: { inputTokens: number; outputTokens: number; totalUSD: number } }
  | { jobId: string; type: 'iteration:error'; stepIndex: number; label: string; error: string }
  | { jobId: string; type: 'assistant:delta'; text: string; iterationLabel?: string }
  | { jobId: string; type: 'assistant:text'; text: string; iterationLabel?: string }
  | { jobId: string; type: 'tool:start'; toolCallId: string; toolName: string; toolInput: Record<string, unknown>; iterationLabel?: string }
  | { jobId: string; type: 'tool:result'; toolCallId: string; output: string; isError: boolean; iterationLabel?: string }
  | { jobId: string; type: 'pipeline:complete'; totalCost: { inputTokens: number; outputTokens: number; totalUSD: number } }
  | { jobId: string; type: 'error'; message: string }
  | { jobId: string; type: 'stopped' }
  | { type: 'job:created'; jobId: string }
  | { type: 'job:state'; jobId: string; status: string; progress: unknown; cost: unknown }
  | { type: 'error'; message: string };

export function usePipelineRunner() {
  const [phase, setPhase] = useState<Phase>('ready');
  const [isConnected, setIsConnected] = useState(false);
  const [jobId, setJobId] = useState<string | null>(null);
  const [steps, setSteps] = useState<StepDef[]>([]);
  const [currentStep, setCurrentStep] = useState<StepStatus | null>(null);
  const [parallelStep, setParallelStep] = useState<ParallelStep | null>(null);
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  const [streamingText, setStreamingText] = useState('');
  const [totalCost, setTotalCost] = useState<{ inputTokens: number; outputTokens: number; totalUSD: number } | null>(null);
  const [runningCost, setRunningCost] = useState({ inputTokens: 0, outputTokens: 0, totalUSD: 0 });
  const [hasError, setHasError] = useState(false);
  const [skippedItems, setSkippedItems] = useState<Array<{ label: string; reason: string }>>([]);
  const [elapsed, setElapsed] = useState(0);
  const [waitingStatus, setWaitingStatus] = useState<{ stepIndex: number; elapsed: number; iterationLabel?: string } | null>(null);
  const wsRef = useRef<WebSocket | null>(null);
  const streamBufferRef = useRef('');
  const startTimeRef = useRef<number>(0);
  const timerRef = useRef<ReturnType<typeof setInterval> | null>(null);
  const jobIdRef = useRef<string | null>(null);
  const inParallelRef = useRef(false);

  const flushStream = useCallback(() => {
    const text = streamBufferRef.current;
    if (text) {
      setMessages((prev) => [...prev, { role: 'assistant', id: randomId(), text }]);
      streamBufferRef.current = '';
      setStreamingText('');
    }
  }, []);

  const stopTimer = useCallback(() => {
    if (timerRef.current) { clearInterval(timerRef.current); timerRef.current = null; }
  }, []);

  const addCost = useCallback((cost: { inputTokens: number; outputTokens: number; totalUSD: number }) => {
    setRunningCost((prev) => ({
      inputTokens: prev.inputTokens + cost.inputTokens,
      outputTokens: prev.outputTokens + cost.outputTokens,
      totalUSD: prev.totalUSD + cost.totalUSD,
    }));
  }, []);

  const handleEvent = useCallback((msg: ServerMessage) => {
    // Filter events by jobId (ignore events from other jobs)
    if ('jobId' in msg && msg.jobId && jobIdRef.current && msg.jobId !== jobIdRef.current) return;

    switch (msg.type) {
      case 'job:created':
        jobIdRef.current = msg.jobId;
        setJobId(msg.jobId);
        break;

      case 'job:state':
        // Reconnection to a completed/failed job
        if (msg.status === 'completed' || msg.status === 'failed' || msg.status === 'stopped' || msg.status === 'interrupted') {
          setPhase('done');
          if (msg.cost) setTotalCost(msg.cost as { inputTokens: number; outputTokens: number; totalUSD: number });
          if (msg.status === 'failed' || msg.status === 'interrupted') setHasError(true);
          stopTimer();
        }
        break;

      case 'pipeline:init':
        setSteps(msg.steps);
        break;

      case 'step:start':
        flushStream();
        setMessages([]);
        setParallelStep(null);
        setWaitingStatus(null);
        inParallelRef.current = false;
        setCurrentStep({
          taskName: msg.taskName,
          iteration: msg.iteration,
          status: 'running',
        });
        break;

      case 'step:complete':
        flushStream();
        setWaitingStatus(null);
        setCurrentStep((prev) => prev ? { ...prev, status: 'complete', cost: msg.cost } : null);
        if (msg.cost) addCost(msg.cost);
        break;

      case 'step:skip':
        setSkippedItems((prev) => [...prev, { label: msg.label, reason: msg.reason }]);
        break;

      case 'step:waiting':
        setWaitingStatus({ stepIndex: msg.stepIndex, elapsed: msg.elapsed, iterationLabel: msg.iterationLabel });
        break;

      case 'step:parallel':
        flushStream();
        setMessages([]);
        setCurrentStep(null);
        inParallelRef.current = true;
        setParallelStep({
          stepIndex: msg.stepIndex,
          taskName: msg.taskName,
          concurrency: msg.concurrency,
          iterations: msg.iterations.map((label) => ({ label, status: 'pending' })),
        });
        break;

      case 'iteration:start':
        setParallelStep((prev) => {
          if (!prev) return prev;
          return {
            ...prev,
            iterations: prev.iterations.map((it) =>
              it.label === msg.label ? { ...it, status: 'running' } : it,
            ),
          };
        });
        break;

      case 'iteration:complete':
        setParallelStep((prev) => {
          if (!prev) return prev;
          return {
            ...prev,
            iterations: prev.iterations.map((it) =>
              it.label === msg.label ? { ...it, status: 'complete', cost: msg.cost } : it,
            ),
          };
        });
        if (msg.cost) addCost(msg.cost);
        break;

      case 'iteration:error':
        setParallelStep((prev) => {
          if (!prev) return prev;
          return {
            ...prev,
            iterations: prev.iterations.map((it) =>
              it.label === msg.label ? { ...it, status: 'error', error: msg.error } : it,
            ),
          };
        });
        break;

      case 'assistant:delta':
        // Skip messages from parallel sub-agents (shown in iteration grid instead)
        if (inParallelRef.current && msg.iterationLabel) break;
        setWaitingStatus(null);
        streamBufferRef.current += msg.text;
        setStreamingText(streamBufferRef.current);
        break;

      case 'assistant:text': {
        if (inParallelRef.current && msg.iterationLabel) break;
        const text = msg.text || streamBufferRef.current;
        if (text) {
          setMessages((prev) => [...prev, { role: 'assistant', id: randomId(), text }]);
        }
        streamBufferRef.current = '';
        setStreamingText('');
        break;
      }

      case 'tool:start':
        if (inParallelRef.current && msg.iterationLabel) break;
        setWaitingStatus(null);
        flushStream();
        setMessages((prev) => [
          ...prev,
          {
            role: 'tool' as const,
            id: randomId(),
            toolCallId: msg.toolCallId,
            toolName: msg.toolName,
            toolInput: msg.toolInput,
            output: undefined,
            isError: false,
          },
        ]);
        break;

      case 'tool:result':
        if (inParallelRef.current && msg.iterationLabel) break;
        setMessages((prev) =>
          prev.map((m) =>
            m.role === 'tool' && 'toolCallId' in m && m.toolCallId === msg.toolCallId
              ? { ...m, output: msg.output, isError: msg.isError }
              : m,
          ),
        );
        break;

      case 'pipeline:complete':
        flushStream();
        setTotalCost(msg.totalCost);
        setPhase('done');
        stopTimer();
        break;

      case 'error':
        flushStream();
        setMessages((prev) => [...prev, { role: 'error' as const, id: randomId(), text: msg.message }]);
        setHasError(true);
        setPhase('done');
        stopTimer();
        break;

      case 'stopped':
        flushStream();
        setPhase('done');
        stopTimer();
        break;
    }
  }, [flushStream, stopTimer, addCost]);

  // The socket below is opened once and must stay open, so its listener is registered once too — and would
  // hold the first render's `handleEvent` forever. That closure carries `flushStream`'s captured
  // `streamingText`, so a re-render mid-run would leave every later event folding into a stale buffer.
  // Reading through a ref keeps the socket long-lived and the handler current.
  const handleEventRef = useRef(handleEvent);
  handleEventRef.current = handleEvent;

  useEffect(() => {
    const token = localStorage.getItem('BEARER_TOKEN');
    if (!token) return;

    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
    const url = `${protocol}//${window.location.host}/api/tasks/pipeline/ws?token=${token}`;
    const ws = new WebSocket(url);
    wsRef.current = ws;

    ws.addEventListener('open', () => {
      setIsConnected(true);
      if (jobIdRef.current) {
        ws.send(JSON.stringify({ type: 'attach', jobId: jobIdRef.current }));
      }
    });
    ws.addEventListener('close', () => setIsConnected(false));

    ws.addEventListener('message', (ev) => {
      try {
        const msg = JSON.parse(ev.data) as ServerMessage;
        handleEventRef.current(msg);
      } catch {
        // ignore
      }
    });

    return () => {
      ws.close();
      wsRef.current = null;
      stopTimer();
    };
  }, []);

  const run = useCallback((taskDirName: string, inputs: Record<string, string>, cwd?: string, model?: string, startAt?: number) => {
    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) return;

    setPhase('running');
    setMessages([]);
    setStreamingText('');
    setTotalCost(null);
    setRunningCost({ inputTokens: 0, outputTokens: 0, totalUSD: 0 });
    setHasError(false);
    setSkippedItems([]);
    setCurrentStep(null);
    setParallelStep(null);
    setWaitingStatus(null);
    setElapsed(0);
    setJobId(null);
    jobIdRef.current = null;
    streamBufferRef.current = '';

    startTimeRef.current = Date.now();
    if (timerRef.current) clearInterval(timerRef.current);
    timerRef.current = setInterval(() => {
      setElapsed(Math.floor((Date.now() - startTimeRef.current) / 1000));
    }, 1000);

    wsRef.current.send(JSON.stringify({ type: 'run', taskDirName, inputs, cwd, model, startAt }));
  }, []);

  const stop = useCallback(() => {
    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN || !jobIdRef.current) return;
    wsRef.current.send(JSON.stringify({ type: 'stop', jobId: jobIdRef.current }));
  }, []);

  return {
    phase, isConnected, jobId, steps, currentStep, parallelStep, messages, streamingText,
    totalCost, runningCost, hasError, skippedItems, elapsed, waitingStatus, run, stop,
  };
}
