src/App.jsx

import React, { useState, useRef, useEffect, lazy, Suspense } from 'react';
import {
  Box,
  Container,
  Heading,
  Text,
  Input,
  Button,
  VStack,
  SimpleGrid,
  Image,
  FormControl,
  FormLabel,
  useToast,
  Progress,
  Badge,
  Flex,
  Divider,
  Spinner,
  Textarea,
  Tooltip,
  InputGroup,
  InputRightElement,
  RadioGroup,
  Radio,
  HStack,
} from '@chakra-ui/react';
import axios from 'axios';

// Lazy load heavier components that aren't needed immediately
const JSZip = lazy(() => import('jszip'));
const { saveAs } = lazy(() => import('file-saver'));

// Function to get filename without extension
const getFilenameWithoutExtension = filename => {
  return filename.substring(0, filename.lastIndexOf("."));
}

// Function to upload image to Cloudinary
const uploadToCloudinary = async (file) => {
  const formData = new FormData();
  formData.append('file', file);
  formData.append('upload_preset', import.meta.env.VITE_REACT_APP_CLOUDINARY_UPLOAD_PRESET);
  
  try {
    const response = await axios.post(
      import.meta.env.VITE_REACT_APP_CLOUDINARY_UPLOAD_URL,
      formData
    );
    
    if (response.status === 200) {
      return response.data.secure_url;
    } else {
      throw new Error('Failed to upload to Cloudinary');
    }
  } catch (error) {
    console.error('Error uploading to Cloudinary:', error);
    throw error;
  }
};

// Extract Glif ID from URL
const extractGlifId = (input) => {
  // If it's already just an ID (e.g. cm7xrsp4t000qlf0cuvxrxf5t), return it
  if (/^[a-zA-Z0-9]{24,}$/.test(input.trim())) {
    return input.trim();
  }
  
  // Try to extract ID from URL like https://glif.app/@R4Z0R1337/glifs/cm7xrsp4t000qlf0cuvxrxf5t
  const urlMatch = input.match(/glifs?\/([a-zA-Z0-9]{24,})/i);
  if (urlMatch && urlMatch[1]) {
    return urlMatch[1];
  }
  
  return null;
};

const API_KEY_STORAGE_KEY = 'glif-lora-api-key';
const CUSTOM_GLIF_ID_STORAGE_KEY = 'glif-lora-custom-id';
const TRIGGER_WORD_STORAGE_KEY = 'glif-lora-trigger-word';
const TRIGGER_PLACEMENT_STORAGE_KEY = 'glif-lora-trigger-placement';
const DEFAULT_GLIF_ID = 'cm7yya7850000la0ckalxpix2';

const App = () => {
  const [apiKey, setApiKey] = useState('');
  const [customGlifInput, setCustomGlifInput] = useState('');
  const [customGlifId, setCustomGlifId] = useState('');
  const [triggerWord, setTriggerWord] = useState('');
  const [triggerPlacement, setTriggerPlacement] = useState('none'); // 'none', 'before', or 'after'
  const [uploadedImages, setUploadedImages] = useState([]);
  const [processing, setProcessing] = useState(false);
  const [downloading, setDownloading] = useState(false);
  const [progress, setProgress] = useState(0);
  const [editingImageId, setEditingImageId] = useState(null);
  const fileInputRef = useRef(null);
  const toast = useToast();

  // Load API key, custom Glif ID, and trigger word settings from localStorage on component mount
  useEffect(() => {
    const savedApiKey = localStorage.getItem(API_KEY_STORAGE_KEY);
    const savedCustomGlifId = localStorage.getItem(CUSTOM_GLIF_ID_STORAGE_KEY);
    const savedTriggerWord = localStorage.getItem(TRIGGER_WORD_STORAGE_KEY);
    const savedTriggerPlacement = localStorage.getItem(TRIGGER_PLACEMENT_STORAGE_KEY);
    
    if (savedApiKey) {
      setApiKey(savedApiKey);
    }
    
    if (savedCustomGlifId) {
      setCustomGlifId(savedCustomGlifId);
      setCustomGlifInput(savedCustomGlifId);
    }

    if (savedTriggerWord) {
      setTriggerWord(savedTriggerWord);
    }

    if (savedTriggerPlacement) {
      setTriggerPlacement(savedTriggerPlacement);
    }
  }, []);

  // Update localStorage when API key changes
  const handleApiKeyChange = (e) => {
    const newApiKey = e.target.value;
    setApiKey(newApiKey);
    localStorage.setItem(API_KEY_STORAGE_KEY, newApiKey);
  };
  
  // Handle custom Glif ID/URL input
  const handleCustomGlifInputChange = (e) => {
    setCustomGlifInput(e.target.value);
  };
  
  // Save and validate custom Glif ID
  const saveCustomGlifId = () => {
    if (!customGlifInput.trim()) {
      // If input is empty, clear the custom ID
      setCustomGlifId('');
      localStorage.removeItem(CUSTOM_GLIF_ID_STORAGE_KEY);
      toast({
        title: 'Custom Glif ID Cleared',
        description: "Using default Glif for autocaption.",
        status: 'info',
        duration: 3000,
        isClosable: true,
      });
      return;
    }
    
    const extractedId = extractGlifId(customGlifInput);
    
    if (extractedId) {
      setCustomGlifId(extractedId);
      localStorage.setItem(CUSTOM_GLIF_ID_STORAGE_KEY, extractedId);
      toast({
        title: 'Custom Glif ID Saved',
        description: `Using Glif ID: ${extractedId}`,
        status: 'success',
        duration: 3000,
        isClosable: true,
      });
    } else {
      toast({
        title: 'Invalid Glif URL or ID',
        description: "Please enter a valid Glif URL or ID.",
        status: 'error',
        duration: 3000,
        isClosable: true,
      });
    }
  };
  
  // Clear custom Glif ID
  const clearCustomGlifId = () => {
    setCustomGlifInput('');
    setCustomGlifId('');
    localStorage.removeItem(CUSTOM_GLIF_ID_STORAGE_KEY);
    toast({
      title: 'Custom Glif ID Cleared',
      description: "Using default Glif for autocaption.",
      status: 'info',
      duration: 3000,
      isClosable: true,
    });
  };

  // Handle trigger word change
  const handleTriggerWordChange = (e) => {
    const newTriggerWord = e.target.value;
    setTriggerWord(newTriggerWord);
    localStorage.setItem(TRIGGER_WORD_STORAGE_KEY, newTriggerWord);
  };

  // Handle trigger placement change
  const handleTriggerPlacementChange = (value) => {
    setTriggerPlacement(value);
    localStorage.setItem(TRIGGER_PLACEMENT_STORAGE_KEY, value);
  };

  // Apply trigger word to caption
  const applyTriggerWord = (caption) => {
    if (!triggerWord.trim() || triggerPlacement === 'none') {
      return caption;
    }

    if (triggerPlacement === 'before') {
      return `${triggerWord.trim()}, ${caption}`;
    } else if (triggerPlacement === 'after') {
      return `${caption}, ${triggerWord.trim()}`;
    }

    return caption;
  };

  const handleFileChange = async (e) => {
    const files = Array.from(e.target.files);
    
    if (files.length === 0) return;

    // Initialize new images with placeholder data
    const newImages = files.map(file => ({
      id: `img-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`,
      file: file,
      name: file.name,
      preview: URL.createObjectURL(file),
      status: 'pending',
      caption: '',
      editedCaption: '',
      isEdited: false,
      error: null
    }));

    setUploadedImages([...uploadedImages, ...newImages]);
  };

  const processImages = async () => {
    if (!apiKey.trim()) {
      toast({
        title: 'API Key Required',
        description: "Please enter your Glif API key.",
        status: 'error',
        duration: 3000,
        isClosable: true,
      });
      return;
    }

    if (uploadedImages.length === 0) {
      toast({
        title: 'No Images',
        description: "Please upload at least one image.",
        status: 'warning',
        duration: 3000,
        isClosable: true,
      });
      return;
    }

    setProcessing(true);
    const pendingImages = uploadedImages.filter(img => img.status === 'pending');
    
    for (let i = 0; i < pendingImages.length; i++) {
      const image = pendingImages[i];
      
      setProgress(Math.floor((i / pendingImages.length) * 100));
      
      // Update image status to processing
      setUploadedImages(prev => prev.map(img => 
        img.id === image.id ? { ...img, status: 'processing' } : img
      ));
      
      try {
        // First upload image to Cloudinary
        const cloudinaryUrl = await uploadToCloudinary(image.file);
        
        // Call the GLIF API with the Cloudinary URL
        const response = await axios.post(`https://simple-api.glif.app/${customGlifId || DEFAULT_GLIF_ID}`, {
          image: cloudinaryUrl
        }, {
          headers: {
            'Authorization': `Bearer ${apiKey}`,
            'Content-Type': 'application/json'
          }
        });
        
        let generatedCaption = response.data.output || 'No caption generated';
        
        // Apply trigger word if specified
        generatedCaption = applyTriggerWord(generatedCaption);
        
        // Update image with caption
        setUploadedImages(prev => prev.map(img => 
          img.id === image.id ? { 
            ...img, 
            status: 'completed', 
            caption: generatedCaption,
            editedCaption: generatedCaption, // Initialize edited caption with the original
            isEdited: false
          } : img
        ));
        
      } catch (error) {
        console.error(`Error processing image ${image.name}:`, error);
        
        // Update image with error
        setUploadedImages(prev => prev.map(img => 
          img.id === image.id ? { 
            ...img, 
            status: 'error', 
            error: error.response?.data?.error || error.message || 'Unknown error occurred' 
          } : img
        ));
      }
    }
    
    setProgress(100);
    setProcessing(false);
    
    toast({
      title: 'Processing Complete',
      description: "All images have been processed.",
      status: 'success',
      duration: 5000,
      isClosable: true,
    });
  };

  const downloadCaptions = async () => {
    const completedImages = uploadedImages.filter(img => img.status === 'completed' && img.caption);
    
    if (completedImages.length === 0) {
      toast({
        title: 'No Captions Available',
        description: "There are no completed captions to download.",
        status: 'warning',
        duration: 3000,
        isClosable: true,
      });
      return;
    }

    setDownloading(true);
    
    try {
      // Dynamically import JSZip and file-saver only when needed
      const [JSZip, { saveAs }] = await Promise.all([
        import('jszip'),
        import('file-saver')
      ]);
      
      // Create a new zip file
      const zip = new JSZip.default();
      
      // Add text files to the zip with the same filenames as the images but .txt extension
      completedImages.forEach(image => {
        const fileName = getFilenameWithoutExtension(image.name) + ".txt";
        const captionToSave = image.isEdited ? image.editedCaption : image.caption;
        zip.file(fileName, captionToSave);
      });
      
      // Generate the zip file
      const zipBlob = await zip.generateAsync({ type: 'blob' });
      
      // Save the zip file
      saveAs(zipBlob, "glif-captions.zip");
      
      toast({
        title: 'Download Complete',
        description: `Successfully packaged ${completedImages.length} captions.`,
        status: 'success',
        duration: 3000,
        isClosable: true,
      });
    } catch (error) {
      console.error("Error creating zip file:", error);
      toast({
        title: 'Download Failed',
        description: "Failed to create zip file. " + error.message,
        status: 'error',
        duration: 3000,
        isClosable: true,
      });
    } finally {
      setDownloading(false);
    }
  };

  const startEditingCaption = (id) => {
    setEditingImageId(id);
  };

  const saveEditedCaption = (id, newCaption) => {
    setUploadedImages(prev => prev.map(img => 
      img.id === id ? { 
        ...img, 
        editedCaption: newCaption,
        isEdited: newCaption !== img.caption 
      } : img
    ));
    setEditingImageId(null);
    
    toast({
      title: 'Caption Saved',
      description: "Your edited caption has been saved.",
      status: 'success',
      duration: 2000,
      isClosable: true,
    });
  };

  const cancelEditingCaption = () => {
    setEditingImageId(null);
  };

  const handleEditCaptionChange = (id, newValue) => {
    setUploadedImages(prev => prev.map(img => 
      img.id === id ? { ...img, editedCaption: newValue } : img
    ));
  };

  const resetCaption = (id) => {
    setUploadedImages(prev => prev.map(img => 
      img.id === id ? { 
        ...img, 
        editedCaption: img.caption,
        isEdited: false 
      } : img
    ));
    
    toast({
      title: 'Caption Reset',
      description: "Caption has been reset to the original generated text.",
      status: 'info',
      duration: 2000,
      isClosable: true,
    });
  };

  const removeImage = (id) => {
    setUploadedImages(prev => prev.filter(img => img.id !== id));
  };

  const clearAll = () => {
    setUploadedImages([]);
    if (fileInputRef.current) fileInputRef.current.value = '';
  };

  const getStatusBadge = (status) => {
    switch (status) {
      case 'pending':
        return <Badge colorScheme="gray">Pending</Badge>;
      case 'processing':
        return <Badge colorScheme="blue">Processing</Badge>;
      case 'completed':
        return <Badge colorScheme="green">Completed</Badge>;
      case 'error':
        return <Badge colorScheme="red">Error</Badge>;
      default:
        return null;
    }
  };

  // Clear API Key function
  const clearApiKey = () => {
    setApiKey('');
    localStorage.removeItem(API_KEY_STORAGE_KEY);
    toast({
      title: 'API Key Cleared',
      description: "Your API key has been removed from browser storage.",
      status: 'info',
      duration: 3000,
      isClosable: true,
    });
  };

  // Check if we have any completed captions
  const hasCaptions = uploadedImages.some(img => img.status === 'completed' && img.caption);

  return (
    <Suspense fallback={
      <Box display="flex" justifyContent="center" alignItems="center" height="100vh">
        <Spinner size="xl" />
      </Box>
    }>
      <Container maxW="container.xl" py={10}>
        <VStack spacing={8} align="stretch">
          <Box textAlign="center">
            <Heading as="h1" size="2xl" mb={2} bgGradient="linear(to-r, cyan.400, blue.500, purple.600)" bgClip="text">
              Glif LoRA Autocaption
            </Heading>
            <Text fontSize="lg" color="gray.400">
              Generate captions for your images using Glif API
            </Text>
          </Box>

          <Divider />

          <Box>
            <FormControl mb={4}>
              <FormLabel>API Key</FormLabel>
              <Flex>
                <Input 
                  type="password"
                  placeholder="Enter your Glif API key" 
                  value={apiKey}
                  onChange={handleApiKeyChange}
                  mr={2}
                />
                <Button 
                  size="md"
                  onClick={clearApiKey}
                  variant="outline"
                  colorScheme="red"
                >
                  Clear
                </Button>
              </Flex>
              <Text fontSize="xs" color="gray.500" mt={1}>
                Your API key is securely saved in your browser for future use.
              </Text>
            </FormControl>
            
            <FormControl mb={4}>
              <FormLabel>Custom Glif ID or URL</FormLabel>
              <InputGroup>
                <Input 
                  type="text"
                  placeholder="Enter custom Glif ID or URL (optional)" 
                  value={customGlifInput}
                  onChange={handleCustomGlifInputChange}
                />
                <InputRightElement width="auto" mr={2}>
                  <Button size="sm" onClick={saveCustomGlifId} mr={2}>
                    Save
                  </Button>
                  <Button size="sm" onClick={clearCustomGlifId} variant="outline" colorScheme="red">
                    Clear
                  </Button>
                </InputRightElement>
              </InputGroup>
              <Text fontSize="xs" color="gray.500" mt={1}>
                Enter a custom Glif ID or URL to use for autocaption, follow the format here: https://glif.app/@saqib/glifs/cm7yya7850000la0ckalxpix2. Leave empty to use the default Glif.
              </Text>
            </FormControl>
            
            <FormControl mb={4}>
              <FormLabel>Trigger Word</FormLabel>
              <Input 
                type="text"
                placeholder="Enter trigger word (optional)" 
                value={triggerWord}
                onChange={handleTriggerWordChange}
              />
              <Text fontSize="xs" color="gray.500" mt={1}>
                Add a trigger word to all captions (e.g., "masterpiece", "best quality").
              </Text>
            </FormControl>

            <FormControl mb={4}>
              <FormLabel>Trigger Word Placement</FormLabel>
              <RadioGroup onChange={handleTriggerPlacementChange} value={triggerPlacement}>
                <HStack spacing={4}>
                  <Radio value="none">No Trigger</Radio>
                  <Radio value="before">Before Caption</Radio>
                  <Radio value="after">After Caption</Radio>
                </HStack>
              </RadioGroup>
              <Text fontSize="xs" color="gray.500" mt={1}>
                Choose where to place the trigger word in relation to the generated caption.
              </Text>
            </FormControl>
            
            <Flex flexWrap="wrap" gap={4} mb={6} justify="center">
              <Button
                onClick={() => fileInputRef.current?.click()}
                variant="solid"
                size="md"
                disabled={processing}
              >
                Upload Images
              </Button>
              <input
                type="file"
                multiple
                accept="image/*"
                onChange={handleFileChange}
                style={{ display: 'none' }}
                ref={fileInputRef}
              />
              
              <Button 
                onClick={processImages}
                colorScheme="blue" 
                isLoading={processing}
                loadingText="Processing"
                disabled={processing || uploadedImages.filter(img => img.status === 'pending').length === 0}
              >
                Process Images
              </Button>
              
              <Button 
                onClick={downloadCaptions}
                colorScheme="teal"
                isLoading={downloading}
                loadingText="Downloading"
                disabled={downloading || !hasCaptions}
              >
                Download Captions
              </Button>
              
              <Button 
                onClick={clearAll}
                variant="outline" 
                disabled={processing || uploadedImages.length === 0}
              >
                Clear All
              </Button>
            </Flex>

            {processing && (
              <Progress 
                value={progress} 
                size="sm" 
                colorScheme="blue" 
                hasStripe 
                mb={4}
                borderRadius="md"
              />
            )}
          </Box>

          {uploadedImages.length > 0 ? (
            <SimpleGrid columns={{ base: 1, md: 2, lg: 3 }} spacing={6}>
              {uploadedImages.map((image) => (
                <Box 
                  key={image.id} 
                  borderWidth="1px" 
                  borderRadius="lg" 
                  overflow="hidden"
                  bg="gray.900"
                  position="relative"
                >
                  <Box position="absolute" top={2} right={2} zIndex={1}>
                    {getStatusBadge(image.status)}
                    {image.isEdited && (
                      <Badge ml={2} colorScheme="purple">Edited</Badge>
                    )}
                  </Box>
                  <Image 
                    src={image.preview} 
                    alt={image.name}
                    w="100%" 
                    maxH="300px"
                    objectFit="contain"
                    bg="gray.800"
                    p={2}
                  />
                  <Box p={4}>
                    <Flex justify="space-between" align="center" mb={2}>
                      <Text fontWeight="semibold" isTruncated maxW="70%">{image.name}</Text>
                      <Button 
                        size="xs" 
                        onClick={() => removeImage(image.id)}
                        disabled={processing}
                      >
                        Remove
                      </Button>
                    </Flex>
                    
                    {image.status === 'completed' && (
                      <>
                        {editingImageId === image.id ? (
                          <Box mt={2}>
                            <Textarea 
                              value={image.editedCaption}
                              onChange={(e) => handleEditCaptionChange(image.id, e.target.value)}
                              size="sm"
                              mb={2}
                            />
                            <Flex justify="flex-end" gap={2}>
                              <Button size="xs" onClick={cancelEditingCaption}>
                                Cancel
                              </Button>
                              <Button 
                                size="xs" 
                                colorScheme="blue" 
                                onClick={() => saveEditedCaption(image.id, image.editedCaption)}
                              >
                                Save
                              </Button>
                            </Flex>
                          </Box>
                        ) : (
                          <Box mt={2} p={2} bg="gray.800" borderRadius="md">
                            <Text fontSize="sm">
                              {image.isEdited ? image.editedCaption : image.caption}
                            </Text>
                            <Flex justify="flex-end" mt={2} gap={2}>
                              {image.isEdited && (
                                <Tooltip label="Reset to original caption">
                                  <Button 
                                    size="xs" 
                                    colorScheme="orange" 
                                    onClick={() => resetCaption(image.id)}
                                  >
                                    Reset
                                  </Button>
                                </Tooltip>
                              )}
                              <Button 
                                size="xs" 
                                colorScheme="teal" 
                                onClick={() => startEditingCaption(image.id)}
                              >
                                Edit
                              </Button>
                            </Flex>
                          </Box>
                        )}
                      </>
                    )}
                    
                    {image.status === 'error' && (
                      <Box mt={2} p={2} bg="red.900" borderRadius="md">
                        <Text fontSize="sm" color="red.200">{image.error}</Text>
                      </Box>
                    )}
                  </Box>
                </Box>
              ))}
            </SimpleGrid>
          ) : (
            <Box 
              textAlign="center" 
              p={10} 
              borderWidth="1px" 
              borderRadius="lg" 
              borderStyle="dashed"
            >
              <Text color="gray.500">Upload images to get started</Text>
            </Box>
          )}
          
          <Divider mt={6} />
          
          <Box as="footer" textAlign="center">
            <Text fontSize="sm" color="gray.500">
              Glif LoRA Autocaption • {new Date().getFullYear()}
            </Text>
          </Box>
        </VStack>
      </Container>
    </Suspense>
  );
};

export default App;