// src/screens/ChatScreen.tsx
import React, { useState, useCallback, useRef, useEffect } from 'react';
import {
  View,
  StyleSheet,
  FlatList,
  KeyboardAvoidingView,
  Platform,
  TextInput as RNTextInput, // Keep if needed for specific props, otherwise can remove
  NativeSyntheticEvent,
  TextInputKeyPressEventData,
  ActivityIndicator // Import ActivityIndicator
} from 'react-native';
import { Text, useTheme, TextInput, IconButton } from 'react-native-paper';
import { SafeAreaView } from 'react-native-safe-area-context';
import apiClient from '../api/client'; // Import the apiClient
import { useRoute, RouteProp } from '@react-navigation/native'; // Import useRoute and RouteProp

// Define the structure for a message
interface Message {
  id: string;
  text: string;
  sender: 'user' | 'ai';
  timestamp: Date;
}

// Define the expected structure for the API response from /nlp/process-command
interface NlpResponse {
  responses: string[];
}

// Define the expected structure for the API response from /nlp/history
interface ChatHistoryResponse {
  id: number;
  sender: 'user' | 'ai';
  text: string;
  timestamp: string; // Backend sends ISO string
}

// Define the type for the navigation route parameters
type RootStackParamList = {
  Chat: { // Assuming 'Chat' is the name of the route for this screen
    initialQuestion?: string; // Make initialQuestion optional
  };
  // Add other routes here if needed
};

type ChatScreenRouteProp = RouteProp<RootStackParamList, 'Chat'>;

const ChatScreen = () => {
  const theme = useTheme();
  const route = useRoute<ChatScreenRouteProp>(); // Get route params
  const initialQuestion = route.params?.initialQuestion; // Extract initialQuestion

  const [messages, setMessages] = useState<Message[]>([]);
  const [inputText, setInputText] = useState('');
  const [isLoading, setIsLoading] = useState(false); // Loading state for sending messages
  const [isHistoryLoading, setIsHistoryLoading] = useState(true); // Loading state for initial history fetch
  const flatListRef = useRef<FlatList>(null);

  // --- Function to send a message to the backend --- (Extracted logic)
  const sendMessageToApi = useCallback(async (textToSend: string) => {
    if (!textToSend) return; // Don't send empty messages

    const userMessage: Message = {
      id: Date.now().toString() + '-user', // Temporary frontend ID
      text: textToSend,
      sender: 'user',
      timestamp: new Date(),
    };

    // Add user message optimistically
    setMessages(prevMessages => [...prevMessages, userMessage]);
    setIsLoading(true);
    // --- Call Backend API ---
    try {
      console.log(`[ChatScreen] Sending to /nlp/process-command: ${textToSend}`);
      const response = await apiClient.post<NlpResponse>('/nlp/process-command', { user_input: textToSend });
      console.log("[ChatScreen] Received response:", response.data);

      const aiResponses: Message[] = [];
      if (response.data && Array.isArray(response.data.responses) && response.data.responses.length > 0) {
         response.data.responses.forEach((responseText, index) => {
            aiResponses.push({
              id: `${Date.now()}-ai-${index}`, // Temporary frontend ID
              text: responseText || "...", // Handle potentially empty strings
              sender: 'ai',
              timestamp: new Date(),
            });
         });
      } else {
        console.warn("[ChatScreen] Received invalid or empty responses array:", response.data);
        aiResponses.push({
          id: Date.now().toString() + '-ai-fallback',
          text: "Sorry, I couldn't process that properly.",
          sender: 'ai',
          timestamp: new Date(),
        });
      }

      // Add all AI responses to the state
      setMessages(prevMessages => [...prevMessages, ...aiResponses]);

    } catch (error: any) {
      console.error("Failed to get AI response:", error.response?.data || error.message || error);
      const errorResponse: Message = {
        id: Date.now().toString() + '-error',
        text: 'Sorry, I encountered an error trying to reach MAIA.',
        sender: 'ai',
        timestamp: new Date(),
      };
       setMessages(prevMessages => [...prevMessages, errorResponse]);
    } finally {
      setIsLoading(false);
    }
    // --- End API Call ---
  }, []); // Dependencies are correct

  // --- Load messages from backend API on mount & handle initial question ---
  useEffect(() => {
    let isMounted = true; // Flag to prevent state updates on unmounted component

    const loadHistoryAndSendInitial = async () => {
      console.log("[ChatScreen] Component mounted. Loading history...");
      setIsHistoryLoading(true);
      let historyLoadedSuccessfully = false;
      try {
        const response = await apiClient.get<ChatHistoryResponse[]>('/nlp/history');
        if (isMounted) {
            console.log("[ChatScreen] Received history:", response.data);
            if (response.data && Array.isArray(response.data)) {
              const historyMessages = response.data.map((msg) => ({
                id: msg.id.toString(),
                text: msg.text,
                sender: msg.sender,
                timestamp: new Date(msg.timestamp),
              }));
              setMessages(historyMessages);
              historyLoadedSuccessfully = true; // Mark history as loaded
              // ADDED BACK: Scroll after initial history is set, with a delay
              setTimeout(() => flatListRef.current?.scrollToEnd({ animated: false }), 100);
            } else {
              console.warn("[ChatScreen] Received invalid history data:", response.data);
              setMessages([]);
            }
        }
      } catch (error: any) {
        if (isMounted) {
            console.error("Failed to load chat history:", error.response?.data || error.message || error);
            setMessages([]); // Clear messages on error
        }
      } finally {
        if (isMounted) {
          setIsHistoryLoading(false);
          console.log("[ChatScreen] History loading finished. History loaded:", historyLoadedSuccessfully);
          // Send initial question *after* history load attempt, if provided
          if (initialQuestion) {
            console.log("[ChatScreen] Initial question provided:", initialQuestion);
            // Use functional update form of setMessages to get the latest state
            setMessages(currentMessages => {
                const lastMessageText = currentMessages[currentMessages.length - 1]?.text;
                if (lastMessageText !== initialQuestion) {
                    console.log("[ChatScreen] Sending initial question now.");
                    // Use a timeout to ensure history state update is processed before sending
                    setTimeout(() => sendMessageToApi(initialQuestion), 0);
                } else {
                    console.log("[ChatScreen] Initial question seems to match last history message, not sending again.");
                }
                return currentMessages; // Return the state unchanged if not sending
            });
          } else {
             console.log("[ChatScreen] No initial question provided.");
          }
        }
      }
    };

    loadHistoryAndSendInitial();

    return () => {
      isMounted = false; // Cleanup function to set flag on unmount
      console.log("[ChatScreen] Component unmounted.");
    };
  // Dependencies are correct
  // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [initialQuestion, sendMessageToApi]);

  // Function to handle sending a message via input button
  const handleSend = useCallback(async () => {
    const trimmedText = inputText.trim();
    if (!trimmedText || isLoading) return;

    setInputText(''); // Clear input immediately
    await sendMessageToApi(trimmedText); // Use the extracted function

  }, [inputText, isLoading, sendMessageToApi]);

  // Function to handle Enter key press for sending
  const handleKeyPress = (e: NativeSyntheticEvent<TextInputKeyPressEventData>) => {
    // Check if Enter is pressed without Shift key
    if (e.nativeEvent.key === 'Enter' && !(e.nativeEvent as any).shiftKey) {
      e.preventDefault(); // Prevent default behavior (like newline)
      handleSend(); // Trigger send action
    }
  };

  // Render individual message item
  const renderMessage = ({ item }: { item: Message }) => {
    const isUser = item.sender === 'user';
    return (
      <View style={[styles.messageBubble, isUser ? styles.userBubble : styles.aiBubble]}>
        <Text style={{ color: isUser ? theme.colors.onPrimary : theme.colors.onSurface }}>
          {item.text}
        </Text>
        {/* Optional: Add timestamp */}
        {/* <Text style={[styles.timestamp, { color: isUser ? theme.colors.onPrimary : theme.colors.onSurfaceVariant }]}>
          {item.timestamp.toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' })}
        </Text> */}
      </View>
    );
  };

  const styles = StyleSheet.create({
    container: { // For SafeAreaView
      flex: 1,
      backgroundColor: theme.colors.background,
    },
    loadingContainer: { // Centering container for loading indicator
        flex: 1,
        justifyContent: 'center',
        alignItems: 'center',
        backgroundColor: theme.colors.background,
    },
    keyboardAvoidingContainer: { // Style for KAV
        flex: 1,
    },
    listContainer: { // Container for the list, should take up available space
      flex: 1,
    },
    messageList: { // Padding for the list content itself
      paddingHorizontal: 10,
      paddingVertical: 10,
    },
    inputContainer: { // Input container should stick to the bottom
      flexDirection: 'row',
      alignItems: 'center', // Align items vertically in the center
      paddingHorizontal: 8, // Add horizontal padding
      paddingVertical: 8,   // Add some vertical padding
      borderTopWidth: StyleSheet.hairlineWidth,
      borderTopColor: theme.colors.outlineVariant,
      backgroundColor: theme.colors.surface, // Use surface color for input area
    },
    textInput: {
      flex: 1, // Take available horizontal space
      marginRight: 8,
      backgroundColor: theme.colors.surface, // Match container background
      // Remove explicit padding if mode="outlined" handles it well
      maxHeight: 100, // Optional: prevent input from getting too tall with multiline
    },
    sendButton: {
       margin: 0, // Remove default margins if IconButton has them
    },
    messageBubble: {
      maxWidth: '80%',
      padding: 12, // Slightly larger padding
      borderRadius: 18, // More rounded corners
      marginBottom: 10,
    },
    userBubble: {
      alignSelf: 'flex-end',
      backgroundColor: theme.colors.primary,
      borderBottomRightRadius: 5, // Keep the chat bubble tail effect
    },
    aiBubble: {
      alignSelf: 'flex-start',
      backgroundColor: theme.colors.surfaceVariant,
      borderBottomLeftRadius: 5, // Keep the chat bubble tail effect
    },
    timestamp: {
        fontSize: 10,
        marginTop: 4,
        alignSelf: 'flex-end',
        opacity: 0.7,
    }
  });

  // Show a loading indicator while history loads
  if (isHistoryLoading) {
    return (
      <SafeAreaView style={styles.loadingContainer} edges={['bottom', 'left', 'right']}>
        <ActivityIndicator animating={true} size="large" color={theme.colors.primary} />
        <Text style={{ marginTop: 10, color: theme.colors.onBackground }}>Loading chat history...</Text>
      </SafeAreaView>
    );
  }

  return (
    <SafeAreaView style={styles.container} edges={['bottom', 'left', 'right']}>
      <KeyboardAvoidingView
        style={styles.keyboardAvoidingContainer}
        behavior={Platform.OS === "ios" ? "padding" : "height"} // Use height for Android if padding causes issues
        keyboardVerticalOffset={Platform.OS === "ios" ? 60 : 0} // Adjust as needed
      >
          {/* List container takes available space */}
          <View style={styles.listContainer}>
            <FlatList
              ref={flatListRef}
              data={messages}
              renderItem={renderMessage}
              keyExtractor={(item) => item.id}
              contentContainerStyle={styles.messageList}
              // ADDED: Scroll to end when content size changes
              onContentSizeChange={() => flatListRef.current?.scrollToEnd({ animated: false })}
            />
          </View>

          {/* Input container is last, outside the list's flex */}
          <View style={styles.inputContainer}>
            <TextInput
              style={styles.textInput}
              value={inputText}
              onChangeText={setInputText}
              placeholder="Ask MAIA..."
              mode="outlined"
              multiline
              onKeyPress={handleKeyPress} // Use onKeyPress for web/desktop-like Enter behavior
              blurOnSubmit={false} // Keep false for multiline + send button
              disabled={isLoading} // Disable input while AI is responding
              outlineStyle={{ borderRadius: 20 }} // Make input more rounded
              dense // Reduce vertical padding
            />
            <IconButton
              icon="send"
              size={24}
              onPress={handleSend}
              disabled={!inputText.trim() || isLoading}
              mode="contained"
              iconColor={theme.colors.onPrimary}
              containerColor={theme.colors.primary}
              style={styles.sendButton}
              animated // Add subtle animation
            />
          </View>
      </KeyboardAvoidingView>
    </SafeAreaView>
  );
};

export default ChatScreen;