import React, { useState, useRef, useEffect } from 'react';
import { 
  ChakraProvider, 
  Box, 
  VStack, 
  HStack, 
  Text, 
  Input, 
  Button, 
  Flex, 
  Heading, 
  Container,
  useToast,
  Divider,
  Progress,
  extendTheme,
  Image
} from '@chakra-ui/react';
import axios from 'axios';
import { useDropzone } from 'react-dropzone';
import { FiSend, FiUpload } from 'react-icons/fi';
import ReactMarkdown from 'react-markdown';

// Star Wars theme
const starWarsTheme = extendTheme({
  colors: {
    brand: {
      100: '#ffe81f', // Star Wars yellow
      200: '#ffe81f',
      300: '#ffe81f',
      400: '#ffe81f',
      500: '#ffe81f',
      600: '#d6c119',
      700: '#a99a14',
      800: '#7c710f',
      900: '#4f480a',
    },
    imperial: {
      500: '#ff0000', // Empire red
    },
    rebel: {
      500: '#4bd5ee', // Rebel blue
    },
    dark: {
      500: '#000000', // Dark side
    },
    light: {
      500: '#ffffff', // Light side
    },
    space: {
      100: '#05050f',
      500: '#0a0a1f',
      900: '#000005',
    }
  },
  fonts: {
    heading: "'Star Jedi', 'Roboto', sans-serif",
    body: "'Roboto', sans-serif",
  },
  styles: {
    global: {
      body: {
        bg: 'space.500',
        color: 'light.500',
      },
    },
  },
});

// API URL - Using the browser's current hostname for backend access
const getAPIURL = () => {
  // If we're in development mode (running with npm start)
  if (process.env.NODE_ENV === 'development') {
    return 'http://localhost:8000';
  }
  
  // Get current protocol (http: or https:)
  const protocol = window.location.protocol;
  const hostname = window.location.hostname;
  
  // For Hugging Face deployment - use the same URL without specifying port
  if (hostname.includes('hf.space')) {
    return `${protocol}//${hostname}`;
  }
  
  // When running in other production environments, use port 8000
  return `${protocol}//${hostname}:8000`;
};

const API_URL = process.env.REACT_APP_API_URL || getAPIURL();

// Debug log
console.log('Using API URL:', API_URL);
console.log('Environment:', process.env.NODE_ENV);
console.log('Window location:', window.location.hostname);

// Add axios default timeout and error handling
axios.defaults.timeout = 120000; // 120 seconds
axios.interceptors.response.use(
  response => response,
  error => {
    console.error('Axios error:', error);
    // Log the specific details
    if (error.response) {
      // The request was made and the server responded with a status code
      // that falls out of the range of 2xx
      console.error('Error response data:', error.response.data);
      console.error('Error response status:', error.response.status);
      console.error('Error response headers:', error.response.headers);
    } else if (error.request) {
      // The request was made but no response was received
      console.error('Error request:', error.request);
      if (error.code === 'ECONNABORTED') {
        console.error('Request timed out after', axios.defaults.timeout, 'ms');
      }
    } else {
      // Something happened in setting up the request that triggered an Error
      console.error('Error message:', error.message);
    }
    return Promise.reject(error);
  }
);

function ChatMessage({ message, isUser, isStreaming }) {
  const [displayedText, setDisplayedText] = useState('');
  const [charIndex, setCharIndex] = useState(0);
  const messageRef = useRef(null);
  
  // Star Wars-style typewriter effect for streamed responses
  useEffect(() => {
    if (isUser || !isStreaming) {
      setDisplayedText(message);
      return;
    }
    
    // Reset if message changes
    if (charIndex === 0) {
      setDisplayedText('');
    }
    
    // Implement the typing effect with randomized speeds for Star Wars terminal feel
    if (charIndex < message.length) {
      const delay = Math.random() * 20 + 10; // Random delay between 10-30ms
      const timer = setTimeout(() => {
        setDisplayedText(prev => prev + message[charIndex]);
        setCharIndex(prevIndex => prevIndex + 1);
      }, delay);
      
      return () => clearTimeout(timer);
    }
  }, [message, charIndex, isUser, isStreaming]);
  
  // Auto-scroll to the bottom of the message as it streams
  useEffect(() => {
    if (isStreaming && messageRef.current) {
      messageRef.current.scrollIntoView({ behavior: 'smooth', block: 'end' });
    }
  }, [displayedText, isStreaming]);
  
  return (
    <Box 
      bg={isUser ? 'rebel.500' : 'imperial.500'} 
      p={3} 
      borderRadius="md" 
      borderWidth="1px"
      borderColor={isUser ? 'brand.500' : 'dark.500'}
      alignSelf={isUser ? 'flex-end' : 'flex-start'}
      maxW="80%"
      boxShadow="0 0 5px"
      color={isUser ? 'dark.500' : 'light.500'}
      ref={messageRef}
      position="relative"
    >
      <Text fontWeight="bold" fontSize="sm" mb={1}>
        {isUser ? 'Rebel Commander' : 'Jedi Archives'}
      </Text>
      {isStreaming && !isUser ? (
        <Box position="relative">
          <ReactMarkdown>{displayedText}</ReactMarkdown>
          {charIndex < message.length && (
            <Box 
              as="span" 
              display="inline-block"
              w="10px"
              h="16px"
              bg="brand.500"
              position="absolute"
              ml="2px"
              opacity={0.8}
              animation="blink 1s step-end infinite"
              sx={{
                '@keyframes blink': {
                  '0%, 100%': { opacity: 0 },
                  '50%': { opacity: 1 }
                }
              }}
            />
          )}
        </Box>
      ) : (
        <ReactMarkdown>{message}</ReactMarkdown>
      )}
    </Box>
  );
}

function FileUploader({ onFileUpload }) {
  const toast = useToast();
  const [isUploading, setIsUploading] = useState(false);
  const [uploadProgress, setUploadProgress] = useState(0);
  const [processingStatus, setProcessingStatus] = useState(null);
  const [processingProgress, setProcessingProgress] = useState(0);
  const [processingSteps, setProcessingSteps] = useState(0);
  
  const { getRootProps, getInputProps } = useDropzone({
    maxFiles: 1,
    maxSize: 5 * 1024 * 1024, // 5MB max size
    accept: {
      'text/plain': ['.txt'],
      'application/pdf': ['.pdf']
    },
    onDropRejected: (rejectedFiles) => {
      toast({
        title: 'Transmission rejected',
        description: rejectedFiles[0]?.errors[0]?.message || 'File rejected by the Empire',
        status: 'error',
        duration: 5000,
        isClosable: true,
      });
    },
    onDrop: async (acceptedFiles) => {
      if (acceptedFiles.length === 0) return;
      
      setIsUploading(true);
      setUploadProgress(0);
      const file = acceptedFiles[0];
      
      // Check file size
      if (file.size > 5 * 1024 * 1024) {
        toast({
          title: 'File too large for hyperdrive',
          description: 'Maximum file size is 5MB - even the Death Star plans were smaller',
          status: 'error',
          duration: 5000,
          isClosable: true,
        });
        setIsUploading(false);
        return;
      }
      
      const formData = new FormData();
      formData.append('file', file);
      
      try {
        // Either use the API_URL or direct backend based on environment
        const uploadUrl = `${API_URL}/upload/`;
        console.log('Uploading file to:', uploadUrl);
        
        const response = await axios.post(uploadUrl, formData, {
          headers: {
            'Content-Type': 'multipart/form-data',
          },
          timeout: 180000, // 3 minutes timeout for large files
          maxContentLength: Infinity,
          maxBodyLength: Infinity,
          onUploadProgress: (progressEvent) => {
            const percentCompleted = Math.round((progressEvent.loaded * 100) / progressEvent.total);
            setUploadProgress(percentCompleted);
          },
          // Add retry logic for network errors
          validateStatus: function (status) {
            return status >= 200 && status < 500; // Handle 4xx errors in our own logic
          }
        });
        
        console.log('Upload response:', response.data);
        setProcessingStatus('starting');
        
        // Start polling for document processing status
        const sessionId = response.data.session_id;
        const pollStatus = async () => {
          try {
            const statusUrl = `${API_URL}/session/${sessionId}/status`;
            console.log('Checking status at:', statusUrl);
            
            const statusResponse = await axios.get(statusUrl);
            console.log('Status response:', statusResponse.data);
            
            if (statusResponse.data.status === 'ready') {
              setProcessingStatus('complete');
              onFileUpload(sessionId, file.name);
              return;
            } else if (statusResponse.data.status === 'failed') {
              setProcessingStatus('failed');
              toast({
                title: 'Processing failed',
                description: 'There was a disturbance in the Force. Please try again with a different file.',
                status: 'error',
                duration: 7000,
                isClosable: true,
              });
              setIsUploading(false);
              return;
            }
            
            // Still processing, continue polling
            setProcessingStatus('processing');
            setTimeout(pollStatus, 3000);
          } catch (error) {
            console.error('Error checking status:', error);
            
            // Continue polling if there are non-critical errors
            if (error.code === 'ECONNABORTED') {
              // Request timed out
              toast({
                title: 'Status check timed out',
                description: 'Your document is being processed by the Jedi Council. Please be patient, this may take time.',
                status: 'warning',
                duration: 7000,
                isClosable: true,
              });
              setProcessingStatus('timeout');
              // Keep polling, but with a longer delay
              setTimeout(pollStatus, 10000);
            } else {
              // Other errors, but still try to continue polling
              setTimeout(pollStatus, 5000);
            }
          }
        };
        
        // Start polling
        setTimeout(pollStatus, 1000);
        
      } catch (error) {
        console.error('Error uploading file:', error);
        setProcessingStatus(null);
        
        let errorMessage = 'Network error - the Death Star has jammed our comms';
        
        if (error.response) {
          errorMessage = error.response.data?.detail || `Imperial error (${error.response.status})`;
        } else if (error.code === 'ECONNABORTED') {
          errorMessage = 'Request timed out. Even the Millennium Falcon would struggle with this file.';
        }
        
        toast({
          title: 'Upload failed',
          description: errorMessage,
          status: 'error',
          duration: 5000,
          isClosable: true,
        });
        
        setIsUploading(false);
      }
    }
  });
  
  // Move pollSessionStatus inside the component where it has access to the necessary variables
  const pollSessionStatus = async (sessionId, file, retries = 40, interval = 5000) => {
    // Increased retries from 30 to 40 for longer processing documents
    let currentRetry = 0;
    
    while (currentRetry < retries) {
      try {
        const statusUrl = `${API_URL}/session/${sessionId}/status`;
        console.log(`Checking status (attempt ${currentRetry + 1}/${retries}):`, statusUrl);
        
        const statusResponse = await axios.get(statusUrl, {
          timeout: 30000 // 30 second timeout for status checks
        });
        
        console.log('Status response:', statusResponse.data);
        
        if (statusResponse.data.status === 'ready') {
          setProcessingStatus('complete');
          setProcessingProgress(100);
          onFileUpload(sessionId, file.name);
          return;
        } else if (statusResponse.data.status === 'failed') {
          setProcessingStatus('failed');
          throw new Error('Processing failed on server');
        }
        
        // Still processing, update progress based on attempt number
        setProcessingStatus('processing');
        // Calculate progress - more rapid at start, slower towards end
        const progressIncrement = 75 / retries; // Max out at 75% during polling
        setProcessingProgress(Math.min(5 + (currentRetry * progressIncrement), 75));
        
        // Increment processing steps to show activity
        setProcessingSteps(prev => prev + 1);
        
        await new Promise(resolve => setTimeout(resolve, interval));
        currentRetry++;
        
        // Increase interval slightly for each retry to prevent overwhelming the server
        interval = Math.min(interval * 1.1, 15000); // Cap at 15 seconds
      } catch (error) {
        console.error('Error checking status:', error);
        
        // If we hit a timeout or network issue, wait a bit longer before retrying
        await new Promise(resolve => setTimeout(resolve, interval * 2));
        currentRetry++;
      }
    }
    
    // If we've exhausted all retries and still don't have a ready status
    throw new Error('Status polling timed out');
  };
  
  // Status message based on current processing state
  const getStatusMessage = () => {
    const steps = ['Analyzing text', 'Splitting document', 'Creating embeddings', 'Building vector database', 'Finalizing'];
    const currentStep = steps[processingSteps % steps.length];
    
    switch(processingStatus) {
      case 'starting':
        return 'Initiating hyperspace jump...';
      case 'uploading':
        return 'Sending document to the Jedi Archives...';
      case 'processing':
        return `${currentStep}... This may take several minutes.`;
      case 'timeout':
        return 'Document processing is taking longer than expected. Patience, young Padawan...';
      case 'failed':
        return 'Document processing failed. The dark side clouded this document.';
      case 'complete':
        return 'Your document has joined the Jedi Archives!';
      default:
        return '';
    }
  };
  
  return (
    <Box 
      {...getRootProps()} 
      border="2px dashed" 
      borderColor="brand.500" 
      borderRadius="md" 
      p={10} 
      textAlign="center"
      cursor="pointer"
      bg="space.100"
      _hover={{ bg: 'space.900', borderColor: 'rebel.500' }}
    >
      <input {...getInputProps()} />
      <VStack spacing={2}>
        <FiUpload size={30} color="#ffe81f" />
        <Text>Drop a holocron (PDF or text file) here, or click to select</Text>
        <Text fontSize="sm" color="brand.500">
          Max file size: 5MB - suitable for Death Star plans
        </Text>
        {isUploading && (
          <>
            <Text color="brand.500">Uploading to the Jedi Archives...</Text>
            <Progress 
              value={processingStatus === 'uploading' ? uploadProgress : processingProgress}
              size="sm" 
              colorScheme="yellow" 
              width="100%" 
              borderRadius="md"
            />
            {processingStatus && (
              <Text 
                color={processingStatus === 'failed' ? 'imperial.500' : 'brand.500'}
                fontSize="sm"
                mt={2}
              >
                {getStatusMessage()}
              </Text>
            )}
          </>
        )}
      </VStack>
    </Box>
  );
}

function App() {
  const [sessionId, setSessionId] = useState(null);
  const [fileName, setFileName] = useState(null);
  const [messages, setMessages] = useState([]);
  const [inputText, setInputText] = useState('');
  const [isProcessing, setIsProcessing] = useState(false);
  const [isDocProcessing, setIsDocProcessing] = useState(false);
  const [streamingIndex, setStreamingIndex] = useState(-1); // Track which message is streaming
  const messagesEndRef = useRef(null);
  const toast = useToast();
  
  const handleFileUpload = (newSessionId, name) => {
    setSessionId(newSessionId);
    setFileName(name);
    setIsDocProcessing(false);
    setMessages([
      { text: `"${name}" has been added to the Jedi Archives. What knowledge do you seek?`, isUser: false }
    ]);
    
    // Don't poll again - already handled in FileUploader
  };
  
  const handleSendMessage = async () => {
    if (!inputText.trim() || !sessionId || isDocProcessing) return;
    
    const userMessage = inputText;
    setInputText('');
    // Add user message right away
    setMessages(prev => [...prev, { text: userMessage, isUser: true }]);
    
    // Add empty response message that will be filled via streaming
    const messageIndex = messages.length + 1; // +1 since we just added user message
    setMessages(prev => [...prev, { text: '', isUser: false }]);
    setStreamingIndex(messageIndex);
    setIsProcessing(true);
    
    try {
      const queryUrl = `${API_URL}/query/`;
      console.log('Sending query to:', queryUrl);
      
      // We need to handle the streaming response from the backend
      const response = await fetch(queryUrl, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({
          session_id: sessionId,
          query: userMessage
        })
      });
      
      if (!response.ok) {
        throw new Error(`HTTP error! status: ${response.status}`);
      }
      
      // Get the response reader for streaming
      const reader = response.body.getReader();
      const decoder = new TextDecoder();
      let accumulatedResponse = '';
      
      try {
        // Process the stream as it comes in
        while (true) {
          const { value, done } = await reader.read();
          
          // If the stream is done, break out of the loop
          if (done) break;
          
          // Decode the chunk and add it to our response
          const chunk = decoder.decode(value, { stream: true });
          accumulatedResponse += chunk;
          
          // Update the current message with the accumulated response
          setMessages(prev => prev.map((msg, idx) => 
            idx === messageIndex ? { ...msg, text: accumulatedResponse } : msg
          ));
        }
      } catch (streamError) {
        console.error('Stream processing error:', streamError);
      } finally {
        // We're done streaming, clean up
        setStreamingIndex(-1);
        setIsProcessing(false);
      }
    } catch (error) {
      console.error('Error sending message:', error);
      setStreamingIndex(-1);
      
      // Handle specific errors
      if (error.response?.status === 409) {
        // Document still processing
        toast({
          title: 'Document still processing',
          description: 'The Jedi Council is still analyzing this document. Please wait a moment and try again.',
          status: 'warning',
          duration: 5000,
          isClosable: true,
        });
        
        setMessages(prev => [...prev.slice(0, -1), { 
          text: "The Jedi Council is still analyzing this document. Patience, young Padawan.", 
          isUser: false 
        }]);
      } else {
        // General error
        toast({
          title: 'Error',
          description: error.response?.data?.detail || 'A disturbance in the Force - make sure the backend is operational',
          status: 'error',
          duration: 5000,
          isClosable: true,
        });
        
        setMessages(prev => [...prev.slice(0, -1), { 
          text: "I find your lack of network connectivity disturbing. Please try again.", 
          isUser: false 
        }]);
      }
      setIsProcessing(false);
    }
  };
  
  // Scroll to the bottom of messages
  React.useEffect(() => {
    messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
  }, [messages]);
  
  // Handle Enter key press
  const handleKeyPress = (e) => {
    if (e.key === 'Enter' && !e.shiftKey) {
      e.preventDefault();
      handleSendMessage();
    }
  };
  
  return (
    <ChakraProvider theme={starWarsTheme}>
      <Box bg="space.500" minH="100vh" py={8}>
        <Container maxW="container.lg">
          <VStack spacing={6} align="stretch" h="100vh">
            <Box textAlign="center" mb={4}>
              <Heading 
                as="h1" 
                size="xl" 
                color="brand.500"
                textShadow="0 0 10px #ffe81f"
                letterSpacing="2px"
              >
                Jedi Archives Chat
              </Heading>
              <Text color="light.500" mt={2}>The galaxy's knowledge at your fingertips</Text>
            </Box>
            
            {!sessionId ? (
              <FileUploader onFileUpload={handleFileUpload} />
            ) : (
              <>
                <Flex justify="space-between" align="center">
                  <Text fontWeight="bold" color="brand.500">
                    Current holocron: {fileName} {isDocProcessing && "(Jedi Council analyzing...)"}
                  </Text>
                  <Button 
                    size="sm" 
                    colorScheme="yellow"
                    variant="outline"
                    onClick={() => {
                      setSessionId(null);
                      setFileName(null);
                      setMessages([]);
                      setIsDocProcessing(false);
                    }}
                  >
                    Access different holocron
                  </Button>
                </Flex>
                
                <Divider borderColor="brand.500" />
                
                <Box 
                  flex="1" 
                  overflowY="auto" 
                  p={4} 
                  bg="space.100" 
                  borderRadius="md"
                  borderWidth="1px"
                  borderColor="brand.500"
                  boxShadow="0 0 15px #ffe81f22"
                  minH="300px"
                >
                  <VStack spacing={4} align="stretch">
                    {messages.map((msg, idx) => (
                      <ChatMessage 
                        key={idx} 
                        message={msg.text} 
                        isUser={msg.isUser}
                        isStreaming={idx === streamingIndex}
                      />
                    ))}
                    {isDocProcessing && (
                      <Box textAlign="center" p={4}>
                        <Progress 
                          size="xs" 
                          isIndeterminate 
                          colorScheme="yellow"
                          width="80%"
                          mx="auto"
                        />
                        <Text mt={2} color="brand.500">
                          The Force is strong with this document... Processing in progress
                        </Text>
                      </Box>
                    )}
                    <div ref={messagesEndRef} />
                  </VStack>
                </Box>
                
                <HStack>
                  <Input
                    placeholder={isDocProcessing 
                      ? "Waiting for the Jedi Council to complete analysis..." 
                      : "What knowledge do you seek from the holocron?"}
                    value={inputText}
                    onChange={(e) => setInputText(e.target.value)}
                    onKeyPress={handleKeyPress}
                    disabled={isProcessing || isDocProcessing}
                    bg="space.100"
                    color="light.500"
                    borderColor="brand.500"
                    _hover={{ borderColor: "rebel.500" }}
                    _focus={{ borderColor: "rebel.500", boxShadow: "0 0 0 1px #4bd5ee" }}
                  />
                  <Button
                    colorScheme="yellow"
                    isLoading={isProcessing}
                    onClick={handleSendMessage}
                    disabled={!inputText.trim() || isProcessing || isDocProcessing}
                    leftIcon={<FiSend />}
                    _hover={{ bg: "rebel.500", color: "dark.500" }}
                  >
                    Send
                  </Button>
                </HStack>
              </>
            )}
          </VStack>
        </Container>
      </Box>
    </ChakraProvider>
  );
}

export default App;