From 13ea0b3a5ca5cbe27c1fe8ceac26185b2269f042 Mon Sep 17 00:00:00 2001
From: aaldebs99 <aaldebs99@gmail.com>
Date: Thu, 19 Jun 2025 21:28:16 +0000
Subject: [PATCH] fix(frontend): labels

---
 frontend/src/components/Upload/UploadZone.tsx |  84 +++++++++++-
 frontend/src/pages/DocumentsPage.tsx          | 125 ++++++++++++++++++
 2 files changed, 208 insertions(+), 1 deletion(-)

diff --git a/frontend/src/components/Upload/UploadZone.tsx b/frontend/src/components/Upload/UploadZone.tsx
index 5ba7046..f74c37a 100644
--- a/frontend/src/components/Upload/UploadZone.tsx
+++ b/frontend/src/components/Upload/UploadZone.tsx
@@ -1,4 +1,4 @@
-import React, { useState, useCallback } from 'react';
+import React, { useState, useCallback, useEffect } from 'react';
 import {
   Box,
   Card,
@@ -17,6 +17,7 @@ import {
   Paper,
   alpha,
   useTheme,
+  Divider,
 } from '@mui/material';
 import {
   CloudUpload as UploadIcon,
@@ -29,6 +30,8 @@ import {
 import { useDropzone, FileRejection, DropzoneOptions } from 'react-dropzone';
 import api from '../../services/api';
 import { useNotifications } from '../../contexts/NotificationContext';
+import LabelSelector from '../Labels/LabelSelector';
+import { type LabelData } from '../Labels/Label';
 
 interface UploadedDocument {
   id: string;
@@ -59,6 +62,42 @@ const UploadZone: React.FC<UploadZoneProps> = ({ onUploadComplete }) => {
   const [files, setFiles] = useState<FileItem[]>([]);
   const [uploading, setUploading] = useState<boolean>(false);
   const [error, setError] = useState<string>('');
+  const [selectedLabels, setSelectedLabels] = useState<LabelData[]>([]);
+  const [availableLabels, setAvailableLabels] = useState<LabelData[]>([]);
+  const [labelsLoading, setLabelsLoading] = useState<boolean>(false);
+
+  useEffect(() => {
+    fetchLabels();
+  }, []);
+
+  const fetchLabels = async () => {
+    try {
+      setLabelsLoading(true);
+      const response = await api.get('/labels?include_counts=false');
+      
+      if (response.status === 200 && Array.isArray(response.data)) {
+        setAvailableLabels(response.data);
+      } else {
+        console.error('Failed to fetch labels:', response);
+      }
+    } catch (error) {
+      console.error('Failed to fetch labels:', error);
+    } finally {
+      setLabelsLoading(false);
+    }
+  };
+
+  const handleCreateLabel = async (labelData: Omit<LabelData, 'id' | 'is_system' | 'created_at' | 'updated_at' | 'document_count' | 'source_count'>) => {
+    try {
+      const response = await api.post('/labels', labelData);
+      const newLabel = response.data;
+      setAvailableLabels(prev => [...prev, newLabel]);
+      return newLabel;
+    } catch (error) {
+      console.error('Failed to create label:', error);
+      throw error;
+    }
+  };
 
   const onDrop = useCallback((acceptedFiles: File[], rejectedFiles: FileRejection[]) => {
     setError('');
@@ -105,6 +144,12 @@ const UploadZone: React.FC<UploadZoneProps> = ({ onUploadComplete }) => {
   const uploadFile = async (fileItem: FileItem): Promise<void> => {
     const formData = new FormData();
     formData.append('file', fileItem.file);
+    
+    // Add selected labels to the form data
+    if (selectedLabels.length > 0) {
+      const labelIds = selectedLabels.map(label => label.id);
+      formData.append('label_ids', JSON.stringify(labelIds));
+    }
 
     try {
       setFiles(prev => prev.map(f => 
@@ -135,6 +180,17 @@ const UploadZone: React.FC<UploadZoneProps> = ({ onUploadComplete }) => {
           : f
       ));
 
+      // Assign labels to the uploaded document if any are selected
+      if (selectedLabels.length > 0) {
+        try {
+          const labelIds = selectedLabels.map(label => label.id);
+          await api.put(`/labels/documents/${response.data.id}`, { label_ids: labelIds });
+        } catch (error) {
+          console.warn('Failed to assign labels to document:', error);
+          // Don't fail the upload if label assignment fails
+        }
+      }
+
       if (onUploadComplete) {
         onUploadComplete(response.data);
       }
@@ -293,6 +349,32 @@ const UploadZone: React.FC<UploadZoneProps> = ({ onUploadComplete }) => {
         </Alert>
       )}
 
+      {/* Label Selection */}
+      <Card elevation={0} sx={{ mb: 3 }}>
+        <CardContent>
+          <Typography variant="h6" sx={{ fontWeight: 600, mb: 2 }}>
+            📋 Label Assignment
+          </Typography>
+          <Typography variant="body2" color="text.secondary" sx={{ mb: 2 }}>
+            Select labels to automatically assign to all uploaded documents
+          </Typography>
+          <LabelSelector
+            selectedLabels={selectedLabels}
+            availableLabels={availableLabels}
+            onLabelsChange={setSelectedLabels}
+            onCreateLabel={handleCreateLabel}
+            placeholder="Choose labels for your documents..."
+            size="medium"
+            disabled={labelsLoading}
+          />
+          {selectedLabels.length > 0 && (
+            <Typography variant="caption" color="text.secondary" sx={{ display: 'block', mt: 1 }}>
+              These labels will be applied to all uploaded documents
+            </Typography>
+          )}
+        </CardContent>
+      </Card>
+
       {/* File List */}
       {files.length > 0 && (
         <Card elevation={0}>
diff --git a/frontend/src/pages/DocumentsPage.tsx b/frontend/src/pages/DocumentsPage.tsx
index 257a6f6..90a4b95 100644
--- a/frontend/src/pages/DocumentsPage.tsx
+++ b/frontend/src/pages/DocumentsPage.tsx
@@ -25,6 +25,10 @@ import {
   FormControl,
   InputLabel,
   Select,
+  Dialog,
+  DialogTitle,
+  DialogContent,
+  DialogActions,
 } from '@mui/material';
 import Grid from '@mui/material/GridLegacy';
 import {
@@ -44,9 +48,13 @@ import {
   Visibility as ViewIcon,
   ChevronLeft as ChevronLeftIcon,
   ChevronRight as ChevronRightIcon,
+  Edit as EditIcon,
 } from '@mui/icons-material';
 import { documentService } from '../services/api';
 import DocumentThumbnail from '../components/DocumentThumbnail';
+import Label, { type LabelData } from '../components/Labels/Label';
+import LabelSelector from '../components/Labels/LabelSelector';
+import { useApi } from '../hooks/useApi';
 
 interface Document {
   id: string;
@@ -59,6 +67,7 @@ interface Document {
   ocr_status?: string;
   ocr_confidence?: number;
   tags: string[];
+  labels?: LabelData[];
 }
 
 interface PaginationInfo {
@@ -79,6 +88,7 @@ type SortOrder = 'asc' | 'desc';
 
 const DocumentsPage: React.FC = () => {
   const navigate = useNavigate();
+  const api = useApi();
   const [documents, setDocuments] = useState<Document[]>([]);
   const [pagination, setPagination] = useState<PaginationInfo>({ total: 0, limit: 20, offset: 0, has_more: false });
   const [loading, setLoading] = useState<boolean>(true);
@@ -89,6 +99,13 @@ const DocumentsPage: React.FC = () => {
   const [sortOrder, setSortOrder] = useState<SortOrder>('desc');
   const [ocrFilter, setOcrFilter] = useState<string>('');
   
+  // Labels state
+  const [availableLabels, setAvailableLabels] = useState<LabelData[]>([]);
+  const [labelsLoading, setLabelsLoading] = useState<boolean>(false);
+  const [labelEditDialogOpen, setLabelEditDialogOpen] = useState<boolean>(false);
+  const [editingDocumentId, setEditingDocumentId] = useState<string | null>(null);
+  const [editingDocumentLabels, setEditingDocumentLabels] = useState<LabelData[]>([]);
+  
   // Menu states
   const [sortMenuAnchor, setSortMenuAnchor] = useState<null | HTMLElement>(null);
   const [docMenuAnchor, setDocMenuAnchor] = useState<null | HTMLElement>(null);
@@ -96,6 +113,7 @@ const DocumentsPage: React.FC = () => {
 
   useEffect(() => {
     fetchDocuments();
+    fetchLabels();
   }, [pagination.limit, pagination.offset, ocrFilter]);
 
   const fetchDocuments = async (): Promise<void> => {
@@ -116,6 +134,63 @@ const DocumentsPage: React.FC = () => {
     }
   };
 
+  const fetchLabels = async (): Promise<void> => {
+    try {
+      setLabelsLoading(true);
+      const response = await api.get('/labels?include_counts=false');
+      
+      if (response.status === 200 && Array.isArray(response.data)) {
+        setAvailableLabels(response.data);
+      } else {
+        console.error('Failed to fetch labels:', response);
+      }
+    } catch (error) {
+      console.error('Failed to fetch labels:', error);
+    } finally {
+      setLabelsLoading(false);
+    }
+  };
+
+  const handleCreateLabel = async (labelData: Omit<LabelData, 'id' | 'is_system' | 'created_at' | 'updated_at' | 'document_count' | 'source_count'>) => {
+    try {
+      const response = await api.post('/labels', labelData);
+      const newLabel = response.data;
+      setAvailableLabels(prev => [...prev, newLabel]);
+      return newLabel;
+    } catch (error) {
+      console.error('Failed to create label:', error);
+      throw error;
+    }
+  };
+
+  const handleEditDocumentLabels = (doc: Document) => {
+    setEditingDocumentId(doc.id);
+    setEditingDocumentLabels(doc.labels || []);
+    setLabelEditDialogOpen(true);
+  };
+
+  const handleSaveDocumentLabels = async () => {
+    if (!editingDocumentId) return;
+
+    try {
+      const labelIds = editingDocumentLabels.map(label => label.id);
+      await api.put(`/labels/documents/${editingDocumentId}`, { label_ids: labelIds });
+      
+      // Update the document in the local state
+      setDocuments(prev => prev.map(doc => 
+        doc.id === editingDocumentId 
+          ? { ...doc, labels: editingDocumentLabels }
+          : doc
+      ));
+      
+      setLabelEditDialogOpen(false);
+      setEditingDocumentId(null);
+      setEditingDocumentLabels([]);
+    } catch (error) {
+      console.error('Failed to update document labels:', error);
+    }
+  };
+
   const handleDownload = async (doc: Document): Promise<void> => {
     try {
       const response = await documentService.download(doc.id);
@@ -396,6 +471,13 @@ const DocumentsPage: React.FC = () => {
           <ListItemIcon><ViewIcon fontSize="small" /></ListItemIcon>
           <ListItemText>View Details</ListItemText>
         </MenuItem>
+        <MenuItem onClick={() => { 
+          if (selectedDoc) handleEditDocumentLabels(selectedDoc); 
+          handleDocMenuClose(); 
+        }}>
+          <ListItemIcon><EditIcon fontSize="small" /></ListItemIcon>
+          <ListItemText>Edit Labels</ListItemText>
+        </MenuItem>
       </Menu>
 
       {/* Documents Grid/List */}
@@ -520,6 +602,27 @@ const DocumentsPage: React.FC = () => {
                         </Stack>
                       )}
                       
+                      {doc.labels && doc.labels.length > 0 && (
+                        <Stack direction="row" spacing={0.5} sx={{ mb: 1, flexWrap: 'wrap' }}>
+                          {doc.labels.slice(0, 3).map((label) => (
+                            <Label
+                              key={label.id}
+                              label={label}
+                              size="small"
+                              variant="filled"
+                            />
+                          ))}
+                          {doc.labels.length > 3 && (
+                            <Chip 
+                              label={`+${doc.labels.length - 3}`}
+                              size="small" 
+                              variant="outlined"
+                              sx={{ fontSize: '0.7rem', height: '20px' }}
+                            />
+                          )}
+                        </Stack>
+                      )}
+                      
                       <Typography variant="caption" color="text.secondary">
                         {formatDate(doc.created_at)}
                       </Typography>
@@ -558,6 +661,28 @@ const DocumentsPage: React.FC = () => {
         </Grid>
       )}
 
+      {/* Label Edit Dialog */}
+      <Dialog open={labelEditDialogOpen} onClose={() => setLabelEditDialogOpen(false)} maxWidth="sm" fullWidth>
+        <DialogTitle>Edit Document Labels</DialogTitle>
+        <DialogContent>
+          <Box sx={{ pt: 2 }}>
+            <LabelSelector
+              selectedLabels={editingDocumentLabels}
+              availableLabels={availableLabels}
+              onLabelsChange={setEditingDocumentLabels}
+              onCreateLabel={handleCreateLabel}
+              placeholder="Select labels for this document..."
+              size="medium"
+              disabled={labelsLoading}
+            />
+          </Box>
+        </DialogContent>
+        <DialogActions>
+          <Button onClick={() => setLabelEditDialogOpen(false)}>Cancel</Button>
+          <Button onClick={handleSaveDocumentLabels} variant="contained">Save</Button>
+        </DialogActions>
+      </Dialog>
+
       {/* Results count and pagination */}
       <Box sx={{ mt: 3 }}>
         <Box sx={{ textAlign: 'center', mb: 2 }}>