utils/file_upload_handler.py

"""
Enhanced File Upload Handler for CAD2Program
Handles various file types, batch uploads, and file validation
"""

import os
import tempfile
import zipfile
import tarfile
import json
import shutil
from pathlib import Path
from typing import Dict, List, Tuple, Optional, Any, Union
import logging
from PIL import Image, ImageOps
import io
import base64
import mimetypes

logger = logging.getLogger(__name__)

class FileUploadHandler:
    """Enhanced file upload handler with validation and processing"""
    
    def __init__(self, upload_dir: str = "uploads", max_file_size: int = 50 * 1024 * 1024):
        self.upload_dir = Path(upload_dir)
        self.upload_dir.mkdir(exist_ok=True)
        self.max_file_size = max_file_size  # 50MB default
        
        # Supported file types
        self.supported_image_types = {'.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.webp'}
        self.supported_program_types = {'.py', '.txt'}
        self.supported_archive_types = {'.zip', '.tar', '.tar.gz', '.tgz'}
        self.supported_cad_types = {'.dxf', '.dwg', '.step', '.stp', '.stl', '.obj'}
        
    def validate_file(self, file_path: Union[str, Path]) -> Dict[str, Any]:
        """Validate uploaded file"""
        
        file_path = Path(file_path)
        
        validation_result = {
            "valid": False,
            "file_type": None,
            "size": 0,
            "errors": [],
            "warnings": []
        }
        
        try:
            # Check if file exists
            if not file_path.exists():
                validation_result["errors"].append("File does not exist")
                return validation_result
            
            # Check file size
            file_size = file_path.stat().st_size
            validation_result["size"] = file_size
            
            if file_size > self.max_file_size:
                validation_result["errors"].append(f"File too large: {file_size / 1024 / 1024:.1f}MB > {self.max_file_size / 1024 / 1024:.1f}MB")
                return validation_result
            
            if file_size == 0:
                validation_result["errors"].append("File is empty")
                return validation_result
            
            # Check file extension
            file_extension = file_path.suffix.lower()
            
            if file_extension in self.supported_image_types:
                validation_result["file_type"] = "image"
                # Validate image
                try:
                    with Image.open(file_path) as img:
                        validation_result["image_info"] = {
                            "format": img.format,
                            "mode": img.mode,
                            "size": img.size
                        }
                        
                        # Check image dimensions
                        if img.size[0] < 64 or img.size[1] < 64:
                            validation_result["warnings"].append("Image resolution is very low")
                        elif img.size[0] > 4096 or img.size[1] > 4096:
                            validation_result["warnings"].append("Image resolution is very high, will be resized")
                        
                except Exception as e:
                    validation_result["errors"].append(f"Invalid image file: {str(e)}")
                    return validation_result
            
            elif file_extension in self.supported_program_types:
                validation_result["file_type"] = "program"
                # Validate program file
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        content = f.read()
                        validation_result["program_info"] = {
                            "lines": len(content.split('\n')),
                            "characters": len(content),
                            "has_primitives": "add_primitive" in content or "Primitive" in content
                        }
                except Exception as e:
                    validation_result["errors"].append(f"Cannot read program file: {str(e)}")
                    return validation_result
            
            elif file_extension in self.supported_archive_types:
                validation_result["file_type"] = "archive"
                # Validate archive
                try:
                    if file_extension == '.zip':
                        with zipfile.ZipFile(file_path, 'r') as zf:
                            validation_result["archive_info"] = {
                                "format": "zip",
                                "files": len(zf.namelist()),
                                "file_list": zf.namelist()[:10]  # First 10 files
                            }
                    elif file_extension in {'.tar', '.tar.gz', '.tgz'}:
                        with tarfile.open(file_path, 'r') as tf:
                            validation_result["archive_info"] = {
                                "format": "tar",
                                "files": len(tf.getnames()),
                                "file_list": tf.getnames()[:10]  # First 10 files
                            }
                except Exception as e:
                    validation_result["errors"].append(f"Invalid archive file: {str(e)}")
                    return validation_result
            
            elif file_extension in self.supported_cad_types:
                validation_result["file_type"] = "cad"
                validation_result["cad_info"] = {
                    "format": file_extension[1:].upper(),
                    "note": "CAD file detected - may need special processing"
                }
            
            else:
                validation_result["errors"].append(f"Unsupported file type: {file_extension}")
                return validation_result
            
            # If we get here, file is valid
            validation_result["valid"] = True
            
        except Exception as e:
            validation_result["errors"].append(f"Validation error: {str(e)}")
        
        return validation_result
    
    def process_single_file(self, file_data: Any, filename: str) -> Dict[str, Any]:
        """Process a single uploaded file"""
        
        # Create unique filename to avoid conflicts
        timestamp = str(int(os.path.getmtime(filename)) if os.path.exists(filename) else 0)
        safe_filename = self._sanitize_filename(filename)
        unique_filename = f"{timestamp}_{safe_filename}"
        
        file_path = self.upload_dir / unique_filename
        
        try:
            # Save file
            if hasattr(file_data, 'read'):
                # File-like object
                with open(file_path, 'wb') as f:
                    f.write(file_data.read())
            elif isinstance(file_data, bytes):
                # Bytes data
                with open(file_path, 'wb') as f:
                    f.write(file_data)
            else:
                # Assume it's a path
                shutil.copy2(file_data, file_path)
            
            # Validate file
            validation = self.validate_file(file_path)
            
            if not validation["valid"]:
                # Clean up invalid file
                if file_path.exists():
                    file_path.unlink()
                return {
                    "success": False,
                    "filename": filename,
                    "errors": validation["errors"],
                    "validation": validation
                }
            
            # Process based on file type
            processed_data = self._process_by_type(file_path, validation)
            
            return {
                "success": True,
                "filename": filename,
                "saved_path": str(file_path),
                "validation": validation,
                "processed_data": processed_data
            }
            
        except Exception as e:
            logger.error(f"Error processing file {filename}: {e}")
            return {
                "success": False,
                "filename": filename,
                "errors": [f"Processing error: {str(e)}"]
            }
    
    def process_batch_upload(self, files: List[Tuple[Any, str]]) -> Dict[str, Any]:
        """Process multiple files at once"""
        
        results = {
            "successful": [],
            "failed": [],
            "summary": {
                "total": len(files),
                "success_count": 0,
                "error_count": 0,
                "images": 0,
                "programs": 0,
                "archives": 0
            }
        }
        
        for file_data, filename in files:
            result = self.process_single_file(file_data, filename)
            
            if result["success"]:
                results["successful"].append(result)
                results["summary"]["success_count"] += 1
                
                # Count file types
                file_type = result["validation"]["file_type"]
                if file_type in results["summary"]:
                    results["summary"][file_type + "s"] += 1
            else:
                results["failed"].append(result)
                results["summary"]["error_count"] += 1
        
        return results
    
    def process_archive(self, archive_path: Union[str, Path]) -> Dict[str, Any]:
        """Extract and process archive contents"""
        
        archive_path = Path(archive_path)
        extract_dir = self.upload_dir / f"extracted_{archive_path.stem}"
        extract_dir.mkdir(exist_ok=True)
        
        extracted_files = []
        
        try:
            # Extract archive
            if archive_path.suffix.lower() == '.zip':
                with zipfile.ZipFile(archive_path, 'r') as zf:
                    zf.extractall(extract_dir)
                    extracted_files = [extract_dir / name for name in zf.namelist()]
            elif archive_path.suffix.lower() in {'.tar', '.tar.gz', '.tgz'}:
                with tarfile.open(archive_path, 'r') as tf:
                    tf.extractall(extract_dir)
                    extracted_files = [extract_dir / name for name in tf.getnames()]
            
            # Process extracted files
            processed_results = []
            
            for extracted_file in extracted_files:
                if extracted_file.is_file():
                    validation = self.validate_file(extracted_file)
                    if validation["valid"]:
                        processed_data = self._process_by_type(extracted_file, validation)
                        processed_results.append({
                            "file_path": str(extracted_file),
                            "validation": validation,
                            "processed_data": processed_data
                        })
            
            return {
                "success": True,
                "extract_dir": str(extract_dir),
                "extracted_count": len(extracted_files),
                "processed_count": len(processed_results),
                "results": processed_results
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": f"Archive processing failed: {str(e)}"
            }
    
    def _process_by_type(self, file_path: Path, validation: Dict) -> Dict[str, Any]:
        """Process file based on its type"""
        
        file_type = validation["file_type"]
        processed_data = {"type": file_type}
        
        if file_type == "image":
            processed_data.update(self._process_image(file_path, validation))
        elif file_type == "program":
            processed_data.update(self._process_program(file_path, validation))
        elif file_type == "archive":
            processed_data.update(self.process_archive(file_path))
        elif file_type == "cad":
            processed_data.update(self._process_cad_file(file_path, validation))
        
        return processed_data
    
    def _process_image(self, file_path: Path, validation: Dict) -> Dict[str, Any]:
        """Process uploaded image file"""
        
        try:
            with Image.open(file_path) as img:
                # Convert to RGB if necessary
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                
                # Resize if too large
                max_size = 1024
                if img.size[0] > max_size or img.size[1] > max_size:
                    img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
                    
                    # Save resized image
                    resized_path = file_path.parent / f"resized_{file_path.name}"
                    img.save(resized_path, 'PNG')
                    
                    return {
                        "original_size": validation["image_info"]["size"],
                        "resized_path": str(resized_path),
                        "new_size": img.size,
                        "format": "PNG"
                    }
                
                return {
                    "size": img.size,
                    "format": img.format,
                    "mode": img.mode,
                    "processed": True
                }
                
        except Exception as e:
            return {"error": f"Image processing failed: {str(e)}"}
    
    def _process_program(self, file_path: Path, validation: Dict) -> Dict[str, Any]:
        """Process uploaded program file"""
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            # Analyze program content
            analysis = {
                "lines": len(content.split('\n')),
                "characters": len(content),
                "primitives": [],
                "imports": [],
                "functions": []
            }
            
            # Extract primitives
            for line in content.split('\n'):
                line = line.strip()
                if 'add_primitive' in line:
                    analysis["primitives"].append(line)
                elif line.startswith('import ') or line.startswith('from '):
                    analysis["imports"].append(line)
                elif line.startswith('def '):
                    analysis["functions"].append(line)
            
            return {
                "analysis": analysis,
                "content_preview": content[:500],  # First 500 characters
                "processed": True
            }
            
        except Exception as e:
            return {"error": f"Program processing failed: {str(e)}"}
    
    def _process_cad_file(self, file_path: Path, validation: Dict) -> Dict[str, Any]:
        """Process CAD files (placeholder for future implementation)"""
        
        return {
            "format": validation["cad_info"]["format"],
            "note": "CAD file processing not yet implemented",
            "future_feature": True
        }
    
    def create_dataset_from_uploads(self, upload_results: Dict) -> Dict[str, Any]:
        """Create a dataset from processed uploads"""
        
        dataset_entries = []
        
        for result in upload_results["successful"]:
            validation = result["validation"]
            processed = result["processed_data"]
            
            if validation["file_type"] == "image":
                # Look for matching program file
                image_path = Path(result["saved_path"])
                base_name = image_path.stem
                
                # Try to find matching program
                program_path = None
                for prog_result in upload_results["successful"]:
                    if prog_result["validation"]["file_type"] == "program":
                        prog_path = Path(prog_result["saved_path"])
                        if prog_path.stem == base_name:
                            program_path = prog_path
                            break
                
                if program_path:
                    # Read program content
                    with open(program_path, 'r') as f:
                        program_content = f.read()
                    
                    entry = {
                        "id": f"upload_{base_name}",
                        "image_path": str(image_path),
                        "program_path": str(program_path),
                        "program_text": program_content,
                        "source": "user_upload",
                        "validation": validation,
                        "processed_data": processed
                    }
                    
                    dataset_entries.append(entry)
        
        # Save dataset
        dataset_file = self.upload_dir / "created_dataset.json"
        with open(dataset_file, 'w') as f:
            json.dump(dataset_entries, f, indent=2)
        
        return {
            "dataset_file": str(dataset_file),
            "entries": len(dataset_entries),
            "pairs_found": len(dataset_entries)
        }
    
    def _sanitize_filename(self, filename: str) -> str:
        """Sanitize filename for safe storage"""
        
        # Remove path separators and dangerous characters
        safe_chars = set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._-")
        sanitized = "".join(c if c in safe_chars else "_" for c in filename)
        
        # Limit length
        if len(sanitized) > 100:
            name, ext = os.path.splitext(sanitized)
            sanitized = name[:100-len(ext)] + ext
        
        return sanitized
    
    def cleanup_old_uploads(self, days: int = 7):
        """Clean up old uploaded files"""
        
        import time
        cutoff_time = time.time() - (days * 24 * 60 * 60)
        
        cleaned_files = 0
        
        for file_path in self.upload_dir.rglob('*'):
            if file_path.is_file():
                try:
                    if file_path.stat().st_mtime < cutoff_time:
                        file_path.unlink()
                        cleaned_files += 1
                except Exception as e:
                    logger.warning(f"Failed to clean up {file_path}: {e}")
        
        return {"cleaned_files": cleaned_files}
    
    def get_upload_stats(self) -> Dict[str, Any]:
        """Get statistics about uploads"""
        
        stats = {
            "total_files": 0,
            "total_size": 0,
            "file_types": {},
            "recent_uploads": []
        }
        
        for file_path in self.upload_dir.rglob('*'):
            if file_path.is_file():
                stats["total_files"] += 1
                file_size = file_path.stat().st_size
                stats["total_size"] += file_size
                
                # Count by extension
                ext = file_path.suffix.lower()
                stats["file_types"][ext] = stats["file_types"].get(ext, 0) + 1
                
                # Track recent uploads (last 24 hours)
                import time
                if file_path.stat().st_mtime > time.time() - 86400:
                    stats["recent_uploads"].append({
                        "name": file_path.name,
                        "size": file_size,
                        "modified": file_path.stat().st_mtime
                    })
        
        return stats

# Gradio-specific upload handlers
def create_gradio_upload_interface():
    """Create Gradio interface components for file uploads"""
    
    import gradio as gr
    
    upload_handler = FileUploadHandler()
    
    def handle_gradio_upload(files):
        """Handle file upload from Gradio interface"""
        if not files:
            return "No files uploaded", ""
        
        # Process files
        file_tuples = []
        for file in files:
            file_tuples.append((file.file, file.name))
        
        results = upload_handler.process_batch_upload(file_tuples)
        
        # Format results for display
        summary = f"""
        ## Upload Results
        
        **Total Files:** {results['summary']['total']}
        **Successful:** {results['summary']['success_count']}
        **Failed:** {results['summary']['error_count']}
        
        **File Types:**
        - Images: {results['summary']['images']}
        - Programs: {results['summary']['programs']}
        - Archives: {results['summary']['archives']}
        """
        
        # Create dataset if we have matching pairs
        dataset_info = ""
        if results['summary']['success_count'] > 1:
            dataset = upload_handler.create_dataset_from_uploads(results)
            dataset_info = f"\n**Dataset Created:** {dataset['entries']} entries"
        
        return summary + dataset_info, str(results)
    
    # Define Gradio components
    file_upload = gr.File(
        label="Upload CAD Files",
        file_count="multiple",
        file_types=[".png", ".jpg", ".jpeg", ".py", ".txt", ".zip"]
    )
    
    upload_btn = gr.Button("Process Uploads", variant="primary")
    
    results_display = gr.Markdown(label="Upload Results")
    detailed_results = gr.JSON(label="Detailed Results", visible=False)
    
    upload_btn.click(
        fn=handle_gradio_upload,
        inputs=[file_upload],
        outputs=[results_display, detailed_results]
    )
    
    return file_upload, upload_btn, results_display, detailed_results

# Utility functions for integration with main app
def setup_upload_system(base_dir: str = "uploads") -> FileUploadHandler:
    """Setup upload system with proper directory structure"""
    
    handler = FileUploadHandler(base_dir)
    
    # Create subdirectories
    (Path(base_dir) / "images").mkdir(exist_ok=True)
    (Path(base_dir) / "programs").mkdir(exist_ok=True)
    (Path(base_dir) / "archives").mkdir(exist_ok=True)
    (Path(base_dir) / "processed").mkdir(exist_ok=True)
    
    return handler

if __name__ == "__main__":
    # Test upload handler
    handler = setup_upload_system("test_uploads")
    
    print("Upload handler created")
    print(f"Upload directory: {handler.upload_dir}")
    
    # Test file validation
    test_extensions = ['.png', '.py', '.zip', '.invalid']
    for ext in test_extensions:
        # Create dummy file for testing
        test_file = handler.upload_dir / f"test{ext}"
        test_file.write_text("test content")
        
        validation = handler.validate_file(test_file)
        print(f"{ext}: Valid={validation['valid']}, Type={validation.get('file_type', 'unknown')}")
        
        # Clean up
        test_file.unlink()
    
    # Get stats
    stats = handler.get_upload_stats()
    print(f"Upload stats: {stats}")