import ndjsonStream from 'ndjson-readablestream';
import { useEffect, useRef, useState } from 'react';

interface UseStreamResult {
  data: string;
  isLoading: boolean;
  error: Error | null;
  stream: (options?: RequestInit) => Promise<void>;
  cancel: () => void;
}

export default function useStream(url: string): UseStreamResult {
  const [data, setData] = useState<string>('');
  const [isLoading, setIsLoading] = useState<boolean>(false);
  const [error, setError] = useState<Error | null>(null);
  const abortControllerRef = useRef<AbortController | null>(null);
  const isMountedRef = useRef<boolean>(true);

  // Track mounted state and cleanup on unmount.
  useEffect(() => {
    isMountedRef.current = true;
    return () => {
      isMountedRef.current = false;
      if (abortControllerRef.current) {
        abortControllerRef.current.abort();
        abortControllerRef.current = null;
      }
    };
  }, []);

  const cancel = () => {
    if (abortControllerRef.current) {
      abortControllerRef.current.abort();
      abortControllerRef.current = null;
    }
    if (isMountedRef.current) {
      setIsLoading(false);
      setError(null);
    }
  };

  const stream = async (options: RequestInit = {}): Promise<void> => {
    // Cancel any existing stream.
    cancel();

    // Create new AbortController.
    abortControllerRef.current = new AbortController();

    if (isMountedRef.current) {
      setIsLoading(true);
      setError(null);
      setData('');
    }

    try {
      const response = await fetch(url, {
        ...options,
        headers: {
          ...options.headers,
          Accept: 'application/x-ndjson',
        },
        signal: abortControllerRef.current.signal,
      });

      if (!response.body) {
        throw new Error('No response body found');
      }

      for await (const chunk of ndjsonStream(response.body)) {
        if (chunk.chunk && isMountedRef.current) {
          setData((prev) => prev + chunk.chunk);
        }
      }
    } catch (err) {
      // Don't set error if it was intentionally aborted.
      if (err instanceof Error && err.name === 'AbortError') {
        return;
      }
      if (isMountedRef.current) {
        setError(err instanceof Error ? err : new Error('Unknown error'));
      }
    } finally {
      if (isMountedRef.current) {
        setIsLoading(false);
      }
      abortControllerRef.current = null;
    }
  };

  return { data, isLoading, error, stream, cancel };
}
