app.py

import os
import hashlib
import json
from flask import Flask, render_template, request, jsonify, send_file
from werkzeug.utils import secure_filename
import torch
import torchvision.transforms as transforms
from torchvision import models
from PIL import Image
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import TfidfVectorizer
import pandas as pd
from datetime import datetime
import shutil
from pathlib import Path
import send2trash  # For Recycle Bin support

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 500 * 1024 * 1024  # 500MB max
app.config['THUMBNAIL_FOLDER'] = 'static/thumbnails'

# Create necessary directories
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(app.config['THUMBNAIL_FOLDER'], exist_ok=True)

# Initialize AI model for image analysis
class ImageFeatureExtractor:
    def __init__(self):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = models.mobilenet_v2(pretrained=True)
        self.model = torch.nn.Sequential(*list(self.model.children())[:-1])
        self.model = self.model.to(self.device)
        self.model.eval()
        
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
    
    def extract_features(self, image_path):
        try:
            image = Image.open(image_path).convert('RGB')
            image_tensor = self.transform(image).unsqueeze(0).to(self.device)
            
            with torch.no_grad():
                features = self.model(image_tensor)
            
            return features.cpu().numpy().flatten()
        except Exception as e:
            print(f"Error extracting features from {image_path}: {e}")
            return None

# Initialize text analyzer
class TextDuplicateAnalyzer:
    def __init__(self):
        self.vectorizer = TfidfVectorizer(max_features=1000, stop_words='english')
    
    def extract_features(self, file_path):
        try:
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                text = f.read()
            return text
        except Exception as e:
            print(f"Error reading text file {file_path}: {e}")
            return None

# Initialize components
image_extractor = ImageFeatureExtractor()
text_analyzer = TextDuplicateAnalyzer()

def get_file_hash(filepath):
    """Calculate MD5 hash of file"""
    hash_md5 = hashlib.md5()
    with open(filepath, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()

def get_file_info(filepath):
    """Get file information"""
    stat = os.stat(filepath)
    return {
        'name': os.path.basename(filepath),
        'path': filepath,
        'size': stat.st_size,
        'modified': datetime.fromtimestamp(stat.st_mtime).strftime('%Y-%m-%d %H:%M:%S'),
        'extension': os.path.splitext(filepath)[1].lower()
    }

def is_image_file(filename):
    """Check if file is an image"""
    image_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff'}
    return os.path.splitext(filename)[1].lower() in image_extensions

def is_text_file(filename):
    """Check if file is a text document"""
    text_extensions = {'.txt', '.doc', '.docx', '.pdf', '.rtf', '.md'}
    return os.path.splitext(filename)[1].lower() in text_extensions

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/scan', methods=['POST'])
def scan_duplicates():
    """Scan for duplicates using AI"""
    folder_path = request.json.get('folder_path', app.config['UPLOAD_FOLDER'])
    
    print(f"\n{'='*50}")
    print(f"🔍 SCANNING FOLDER: {folder_path}")
    print(f"{'='*50}")
    
    if not os.path.exists(folder_path):
        print(f"❌ ERROR: Folder does not exist: {folder_path}")
        return jsonify({'error': 'Folder does not exist'}), 400
    
    # Get all files
    files = []
    for root, dirs, filenames in os.walk(folder_path):
        for filename in filenames:
            filepath = os.path.join(root, filename)
            files.append(filepath)
    
    if not files:
        print("❌ No files found in folder")
        return jsonify({'error': 'No files found'}), 400
    
    # Separate files by type
    image_files = [f for f in files if is_image_file(f)]
    text_files = [f for f in files if is_text_file(f)]
    other_files = [f for f in files if not is_image_file(f) and not is_text_file(f)]
    
    results = {
        'exact_duplicates': [],
        'near_duplicates': [],
        'similar_images': [],
        'stats': {
            'total_files': len(files),
            'total_size': sum(os.path.getsize(f) for f in files),
            'image_files': len(image_files),
            'text_files': len(text_files),
            'other_files': len(other_files)
        }
    }
    
    # Find exact duplicates using hash
    hash_map = {}
    for filepath in files:
        file_hash = get_file_hash(filepath)
        if file_hash in hash_map:
            results['exact_duplicates'].append({
                'original': get_file_info(hash_map[file_hash]),
                'duplicate': get_file_info(filepath)
            })
        else:
            hash_map[file_hash] = filepath
    
    # Find near-duplicate images using AI
    if image_files:
        image_features = []
        valid_images = []
        
        for img_path in image_files:
            features = image_extractor.extract_features(img_path)
            if features is not None:
                image_features.append(features)
                valid_images.append(img_path)
        
        if len(valid_images) > 1:
            similarity_matrix = cosine_similarity(image_features)
            
            for i in range(len(valid_images)):
                for j in range(i+1, len(valid_images)):
                    similarity = similarity_matrix[i][j]
                    if 0.80 <= similarity < 0.95:
                        results['similar_images'].append({
                            'file1': get_file_info(valid_images[i]),
                            'file2': get_file_info(valid_images[j]),
                            'similarity': float(similarity)
                        })
    
    # Calculate potential savings
    duplicate_size = sum(d['duplicate']['size'] for d in results['exact_duplicates'])
    similar_size = sum(d['file2']['size'] for d in results['similar_images'])
    
    results['stats']['potential_savings'] = duplicate_size + similar_size
    results['stats']['savings_percentage'] = round(
        (results['stats']['potential_savings'] / results['stats']['total_size']) * 100, 2
    ) if results['stats']['total_size'] > 0 else 0
    
    return jsonify(results)

# ============== FIXED DELETE FUNCTION WITH RECYCLE BIN ==============
@app.route('/delete', methods=['POST'])
def delete_files():
    """Move selected duplicate files to Recycle Bin"""
    files_to_delete = request.json.get('files', [])
    
    deleted = []
    errors = []
    
    for filepath in files_to_delete:
        try:
            # Clean up the path - fix missing backslashes
            if ':' in filepath and not filepath.startswith('E:\\'):
                # Extract just the filename if path is corrupted
                filename = os.path.basename(filepath)
                correct_path = os.path.join('E:\\AI_DRIVEN_DUPLICATE_FINDER\\uploads', filename)
                
                if os.path.exists(correct_path):
                    filepath = correct_path
                else:
                    # Try to find the file in uploads folder
                    uploads_path = os.path.join('uploads', filename)
                    if os.path.exists(uploads_path):
                        filepath = uploads_path
            
            if os.path.exists(filepath):
                # Send to Recycle Bin
                send2trash.send2trash(filepath)
                deleted.append(filepath)
                print(f"✅ Moved to Recycle Bin: {filepath}")
            else:
                errors.append(f"File not found: {filepath}")
        except Exception as e:
            errors.append(str(e))
            print(f"❌ Error: {e}")
    
    return jsonify({
        'deleted': deleted,
        'errors': errors,
        'message': f"Successfully moved {len(deleted)} files to Recycle Bin"
    })

# ============== FIXED KEEP BEST FUNCTION WITH RECYCLE BIN ==============
@app.route('/keep-best', methods=['POST'])
def keep_best_files():
    """Keep the best version of duplicate files"""
    duplicate_groups = request.json.get('groups', [])
    
    kept = []
    deleted = []
    errors = []
    
    for group in duplicate_groups:
        try:
            # Sort files by criteria (keep the largest file)
            files = sorted(group['files'], 
                         key=lambda x: (
                             -os.path.getsize(x['path']) if os.path.exists(x['path']) else 0
                         ))
            
            # Keep the first file, delete the rest
            keeper = files[0]
            kept.append(keeper['path'])
            
            for duplicate in files[1:]:
                try:
                    dup_path = duplicate['path']
                    
                    # Fix path if needed
                    if not os.path.exists(dup_path):
                        filename = os.path.basename(dup_path)
                        dup_path = os.path.join('uploads', filename)
                    
                    if os.path.exists(dup_path):
                        send2trash.send2trash(dup_path)
                        deleted.append(dup_path)
                        print(f"✅ Moved duplicate to Recycle Bin: {dup_path}")
                    else:
                        errors.append(f"Duplicate not found: {dup_path}")
                except Exception as e:
                    errors.append(f"Error deleting {duplicate['path']}: {str(e)}")
                    
        except Exception as e:
            errors.append(f"Error processing group: {str(e)}")
    
    return jsonify({
        'kept': kept,
        'deleted': deleted,
        'errors': errors,
        'message': f"Kept {len(kept)} files, moved {len(deleted)} duplicates to Recycle Bin"
    })

@app.route('/export-report', methods=['POST'])
def export_report():
    """Export scan results to CSV"""
    results = request.json.get('results', {})
    
    report_path = f'reports/duplicate_report_{datetime.now().strftime("%Y%m%d_%H%M%S")}.csv'
    os.makedirs('reports', exist_ok=True)
    
    # Prepare data for CSV
    rows = []
    
    # Add exact duplicates
    for dup in results.get('exact_duplicates', []):
        rows.append({
            'Type': 'Exact Duplicate',
            'Original File': dup['original']['name'],
            'Duplicate File': dup['duplicate']['name'],
            'Original Path': dup['original']['path'],
            'Duplicate Path': dup['duplicate']['path'],
            'Size (bytes)': dup['duplicate']['size'],
            'Similarity': '100%'
        })
    
    # Add near duplicates
    for sim in results.get('similar_images', []):
        rows.append({
            'Type': 'Near Duplicate',
            'Original File': sim['file1']['name'],
            'Duplicate File': sim['file2']['name'],
            'Original Path': sim['file1']['path'],
            'Duplicate Path': sim['file2']['path'],
            'Size (bytes)': sim['file2']['size'],
            'Similarity': f"{sim['similarity']*100:.1f}%"
        })
    
    df = pd.DataFrame(rows)
    df.to_csv(report_path, index=False)
    
    return send_file(report_path, as_attachment=True, download_name='duplicate_report.csv')

if __name__ == '__main__':
    print("="*50)
    print("SanDisk AI Smart Duplicate Finder")
    print("="*50)
    print(f"PyTorch version: {torch.__version__}")
    print(f"Device: {'CUDA' if torch.cuda.is_available() else 'CPU'}")
    print(f"Upload folder: {app.config['UPLOAD_FOLDER']}")
    print("="*50)
    print("\n🚀 Server starting...")
    print("📱 Open http://localhost:5000 in your browser")
    print("="*50)
    
    # Clean up any existing .deleted files on startup
    try:
        uploads_dir = app.config['UPLOAD_FOLDER']
        for file in os.listdir(uploads_dir):
            if file.endswith('.deleted'):
                file_path = os.path.join(uploads_dir, file)
                send2trash.send2trash(file_path)
                print(f"🧹 Cleaned up: {file}")
    except:
        pass
    
    app.run(debug=True, port=5000)