import React, { useState, useRef, useEffect } from 'react';
import { 
  Box, 
  Paper, 
  TextField, 
  Button, 
  Typography, 
  CircularProgress,
  Divider,
  List,
  ListItem,
  ListItemText,
  Collapse,
  IconButton
} from '@mui/material';
import SendIcon from '@mui/icons-material/Send';
import ExpandMoreIcon from '@mui/icons-material/ExpandMore';
import ExpandLessIcon from '@mui/icons-material/ExpandLess';
import axios from 'axios';

const Message = ({ content, isUser, context }) => {
  const [expanded, setExpanded] = useState(false);
  
  return (
    <Box 
      sx={{ 
        display: 'flex', 
        flexDirection: 'column',
        mb: 2,
        alignItems: isUser ? 'flex-end' : 'flex-start',
      }}
    >
      <Paper 
        elevation={1} 
        sx={{ 
          p: 2, 
          maxWidth: '80%',
          backgroundColor: isUser ? 'primary.light' : 'background.paper',
          color: isUser ? 'primary.contrastText' : 'text.primary',
        }}
      >
        <Typography variant="body1">{content}</Typography>
      </Paper>
      
      {context && context.length > 0 && (
        <Box sx={{ width: '100%', mt: 1 }}>
          <Button 
            size="small" 
            endIcon={expanded ? <ExpandLessIcon /> : <ExpandMoreIcon />}
            onClick={() => setExpanded(!expanded)}
            sx={{ color: 'text.secondary' }}
          >
            {expanded ? 'Hide Context' : 'Show Context'}
          </Button>
          
          <Collapse in={expanded}>
            <Paper variant="outlined" sx={{ mt: 1, p: 1 }}>
              <Typography variant="subtitle2" gutterBottom>
                Context used for this answer:
              </Typography>
              <List dense>
                {context.map((item, index) => (
                  <ListItem key={index} divider={index < context.length - 1}>
                    <ListItemText 
                      primary={item[0].substring(0, 100) + (item[0].length > 100 ? '...' : '')} 
                      secondary={`Relevance: ${(item[1] * 100).toFixed(2)}%`}
                    />
                  </ListItem>
                ))}
              </List>
            </Paper>
          </Collapse>
        </Box>
      )}
    </Box>
  );
};

const ChatInterface = ({ fileName, isProcessing }) => {
  const [messages, setMessages] = useState([]);
  const [input, setInput] = useState('');
  const [loading, setLoading] = useState(false);
  const messagesEndRef = useRef(null);
  
  const scrollToBottom = () => {
    messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' });
  };
  
  useEffect(() => {
    scrollToBottom();
  }, [messages]);
  
  const handleSend = async () => {
    if (!input.trim()) return;
    
    const userMessage = input.trim();
    setInput('');
    setMessages(prev => [...prev, { content: userMessage, isUser: true }]);
    setLoading(true);
    
    try {
      const response = await axios.post('http://localhost:8000/ask', {
        question: userMessage
      });
      
      setMessages(prev => [
        ...prev, 
        { 
          content: response.data.response, 
          isUser: false,
          context: response.data.context
        }
      ]);
    } catch (error) {
      console.error('Error asking question:', error);
      setMessages(prev => [
        ...prev, 
        { 
          content: 'Sorry, there was an error processing your question. Please try again.',
          isUser: false
        }
      ]);
    } finally {
      setLoading(false);
    }
  };
  
  const handleKeyPress = (e) => {
    if (e.key === 'Enter' && !e.shiftKey) {
      e.preventDefault();
      handleSend();
    }
  };
  
  return (
    <Paper elevation={3} sx={{ p: 3, height: '70vh', display: 'flex', flexDirection: 'column' }}>
      <Typography variant="h6" gutterBottom>
        Chat about {fileName}
      </Typography>
      
      <Divider sx={{ mb: 2 }} />
      
      <Box sx={{ flexGrow: 1, overflow: 'auto', mb: 2 }}>
        {messages.length === 0 ? (
          <Box sx={{ display: 'flex', justifyContent: 'center', alignItems: 'center', height: '100%' }}>
            <Typography variant="body1" color="text.secondary">
              Ask a question about the document to get started
            </Typography>
          </Box>
        ) : (
          messages.map((message, index) => (
            <Message 
              key={index} 
              content={message.content} 
              isUser={message.isUser}
              context={message.context}
            />
          ))
        )}
        <div ref={messagesEndRef} />
      </Box>
      
      <Box sx={{ display: 'flex', alignItems: 'center' }}>
        <TextField
          fullWidth
          variant="outlined"
          placeholder="Ask a question about the document..."
          value={input}
          onChange={(e) => setInput(e.target.value)}
          onKeyPress={handleKeyPress}
          disabled={loading || isProcessing}
          multiline
          maxRows={3}
          sx={{ mr: 1 }}
        />
        <IconButton 
          color="primary" 
          onClick={handleSend} 
          disabled={loading || isProcessing || !input.trim()}
          sx={{ p: 1 }}
        >
          {loading ? <CircularProgress size={24} /> : <SendIcon />}
        </IconButton>
      </Box>
    </Paper>
  );
};

export default ChatInterface;