import React, { useState, useRef } 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;
  
  // When running in production, use the same host with the backend port
  // This works because we're exposing the backend port in docker-compose
  // If port is 7860, use the same port (Hugging Face scenario)
  const currentPort = window.location.port;
  if (currentPort === '7860') {
    return `${protocol}//${window.location.hostname}:${currentPort}`;
  } else {
    return `${protocol}//${window.location.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 }) {
  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'}
    >
      <Text fontWeight="bold" fontSize="sm" mb={1}>
        {isUser ? 'Rebel Commander' : 'Jedi Archives'}
      </Text>
      <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 { 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',
          },
          onUploadProgress: (progressEvent) => {
            const percentCompleted = Math.round((progressEvent.loaded * 100) / progressEvent.total);
            setUploadProgress(percentCompleted);
          }
        });
        
        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);
      }
    }
  });
  
  // Status message based on current processing state
  const getStatusMessage = () => {
    switch(processingStatus) {
      case 'starting':
        return 'Initiating hyperspace jump...';
      case 'processing':
        return 'The Force is analyzing your document... 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={uploadProgress} 
              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 messagesEndRef = useRef(null);
  const toast = useToast();
  
  const handleFileUpload = (newSessionId, name) => {
    setSessionId(newSessionId);
    setFileName(name);
    setIsDocProcessing(true);
    setMessages([
      { text: `Processing ${name}. May the Force be with you...`, isUser: false }
    ]);
    
    // Poll for document processing status
    const checkStatus = async () => {
      try {
        const response = await axios.get(`${API_URL}/session/${newSessionId}/status`);
        console.log('Status response:', response.data);
        
        if (response.data.status === 'ready') {
          setIsDocProcessing(false);
          setMessages([
            { text: `"${name}" has been added to the Jedi Archives. What knowledge do you seek?`, isUser: false }
          ]);
          return;
        }
        
        // Continue polling if still processing
        if (response.data.status === 'processing') {
          setTimeout(checkStatus, 2000);
        }
      } catch (error) {
        console.error('Error checking status:', error);
        // Continue polling even if there's an error
        setTimeout(checkStatus, 3000);
      }
    };
    
    checkStatus();
  };
  
  const handleSendMessage = async () => {
    if (!inputText.trim() || !sessionId || isDocProcessing) return;
    
    const userMessage = inputText;
    setInputText('');
    setMessages(prev => [...prev, { text: userMessage, isUser: true }]);
    setIsProcessing(true);
    
    try {
      // Either use the API_URL or direct backend based on environment
      const queryUrl = `${API_URL}/query/`;
      console.log('Sending query to:', queryUrl);
      
      const response = await axios.post(queryUrl, {
        session_id: sessionId,
        query: userMessage
      });
      
      console.log('Query response:', response.data);
      setMessages(prev => [...prev, { text: response.data.response, isUser: false }]);
    } catch (error) {
      console.error('Error sending message:', error);
      
      // 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, { 
          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, { 
          text: "I find your lack of network connectivity disturbing. Please try again.", 
          isUser: false 
        }]);
      }
    } finally {
      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}
                      />
                    ))}
                    {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;