import React, { useState, useEffect, useCallback, useRef } from 'react';
import { useNeura } from './NeuraProvider';
import { CompletionsRequestOptions, NeuraStream as NeuraStreamType, StreamChunk, ProcessingError } from '../types';

interface NeuraStreamProps {
  options: CompletionsRequestOptions;
  onChunk?: (chunk: StreamChunk) => void;
  onComplete?: (fullText: string) => void;
  onError?: (error: Error | ProcessingError) => void;
  children: (state: {
    chunks: StreamChunk[];
    fullText: string;
    loading: boolean;
    error: Error | ProcessingError | null;
    restart: () => void;
  }) => React.ReactNode;
}

export const NeuraStream: React.FC<NeuraStreamProps> = ({
  options,
  onChunk,
  onComplete,
  onError,
  children,
}) => {
  const { sdk } = useNeura();
  const [chunks, setChunks] = useState<StreamChunk[]>([]);
  const [fullText, setFullText] = useState<string>('');
  const [loading, setLoading] = useState<boolean>(true);
  const [error, setError] = useState<Error | ProcessingError | null>(null);
  
  // Use a ref to keep track of the stream
  const streamRef = useRef<NeuraStreamType | null>(null);

  const startStream = useCallback(() => {
    // Reset state
    setChunks([]);
    setFullText('');
    setLoading(true);
    setError(null);
    
    // Clean up previous stream if it exists
    if (streamRef.current) {
      streamRef.current.removeAllListeners();
    }
    
    // Create a new stream
    const stream = sdk.createCompletionStream(options);
    streamRef.current = stream;
    
    let accumulatedText = '';
    
    stream.on('chunk', (chunk: StreamChunk) => {
      setChunks((prev) => [...prev, chunk]);
      accumulatedText += chunk.chunk;
      setFullText(accumulatedText);
      onChunk?.(chunk);
    });
    
    stream.on('done', () => {
      setLoading(false);
      onComplete?.(accumulatedText);
    });
    
    // Fix: Update the error handler to accept both Error and ProcessingError
    stream.on('error', (err: Error | ProcessingError) => {
      setError(err);
      setLoading(false);
      onError?.(err);
    });
  }, [sdk, options, onChunk, onComplete, onError]);

  // Start streaming when the component mounts
  useEffect(() => {
    startStream();
    
    // Clean up the stream when the component unmounts
    return () => {
      if (streamRef.current) {
        streamRef.current.removeAllListeners();
        streamRef.current = null;
      }
    };
  }, [startStream]);

  return <>{children({ chunks, fullText, loading, error, restart: startStream })}</>;
};