t.py

#!/usr/bin/env python3
"""
Ultra-fast file transfer system optimized for large datasets
Fixed version with proper connection handling and no threading conflicts
"""

import os
import socket
import threading
import hashlib
import time
import json
import zlib
import struct
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
from queue import Queue
import argparse

class FastFileTransfer:
    def __init__(self, chunk_size=1024*1024, max_workers=8, compress=True):
        self.chunk_size = chunk_size  # 1MB chunks
        self.max_workers = max_workers
        self.compress = compress
        self.stats = {
            'files_sent': 0,
            'bytes_sent': 0,
            'start_time': 0,
            'compression_ratio': 0
        }
        self.stats_lock = threading.Lock()
    
    def get_file_list(self, folder_path):
        """Efficiently scan directory tree and build file list"""
        files = []
        folder_path = Path(folder_path)
        
        print(f"Scanning directory: {folder_path}")
        for root, dirs, filenames in os.walk(folder_path):
            for filename in filenames:
                filepath = Path(root) / filename
                try:
                    size = filepath.stat().st_size
                    rel_path = filepath.relative_to(folder_path)
                    files.append({
                        'path': str(rel_path),
                        'size': size,
                        'full_path': str(filepath)
                    })
                except (OSError, IOError):
                    print(f"Warning: Could not access {filepath}")
                    continue
        
        print(f"Found {len(files)} files")
        return files
    
    def calculate_checksum(self, filepath):
        """Calculate MD5 checksum for file integrity"""
        hash_md5 = hashlib.md5()
        try:
            with open(filepath, "rb") as f:
                for chunk in iter(lambda: f.read(4096), b""):
                    hash_md5.update(chunk)
            return hash_md5.hexdigest()
        except IOError:
            return None

class FileTransferServer:
    def __init__(self, host='0.0.0.0', port=9999, max_workers=8):
        self.host = host
        self.port = port
        self.max_workers = max_workers
        self.transfer = FastFileTransfer(max_workers=max_workers)
        self.file_queue = Queue()
        self.active_connections = 0
        self.connection_lock = threading.Lock()
    
    def start_server(self, folder_path):
        """Start the file transfer server"""
        print(f"Starting server on {self.host}:{self.port}")
        print(f"Serving folder: {folder_path}")
        
        server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1024*1024)  # 1MB send buffer
        server_sock.bind((self.host, self.port))
        server_sock.listen(1)
        
        print("Waiting for client connection...")
        client_sock, addr = server_sock.accept()
        client_sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 1024*1024)
        
        # Use configurable timeout if available
        timeout = getattr(self, '_socket_timeout', 1800)
        client_sock.settimeout(timeout)  # Use configurable timeout
        
        print(f"Client connected from {addr}")
        
        try:
            self.handle_client(client_sock, folder_path)
        except Exception as e:
            print(f"Error handling client: {e}")
        finally:
            print("Closing connections...")
            client_sock.close()
            server_sock.close()
    
    def handle_client(self, sock, folder_path):
        """Handle file transfer to client - Parallel version using workers"""
        # Get file list
        files = self.transfer.get_file_list(folder_path)
        total_size = sum(f['size'] for f in files)
        
        print(f"Total files: {len(files)}")
        print(f"Total size: {total_size / (1024**3):.2f} GB")
        
        # Send file manifest
        manifest = {
            'files': files,
            'total_size': total_size,
            'compress': self.transfer.compress
        }
        manifest_json = json.dumps(manifest).encode()
        
        try:
            self.send_data(sock, struct.pack('!I', len(manifest_json)) + manifest_json)
            print("Manifest sent, starting file transfer...")
        except Exception as e:
            print(f"Failed to send manifest: {e}")
            return
        
        # Start transfer
        self.transfer.stats['start_time'] = time.time()
        successful_files = 0
        failed_files = 0
        skipped_files = 0
        
        # Send files using ThreadPoolExecutor for parallel processing
        print(f"Using {self.max_workers} worker threads for file processing")
        
        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            # Submit all file transfer tasks
            future_to_file = {}
            for file_info in files:
                if os.path.exists(file_info['full_path']):
                    future = executor.submit(self.prepare_file_data, file_info)
                    future_to_file[future] = file_info
                else:
                    print(f"⚠ Skipping missing file: {file_info['path']}")
                    skipped_files += 1
            
            # Process completed tasks and send data
            for i, future in enumerate(as_completed(future_to_file)):
                file_info = future_to_file[future]
                file_path = file_info['path']
                
                try:
                    # Get prepared file data from worker
                    file_data = future.result()
                    
                    if file_data is None:
                        failed_files += 1
                        print(f"⚠ Failed to prepare: {file_path}")
                        continue
                    
                    # Send the prepared data over socket
                    try:
                        self.send_prepared_file(sock, file_data)
                        successful_files += 1
                    except Exception as e:
                        failed_files += 1
                        print(f"⚠ Failed to send {file_path}: {e}")
                    
                    # Progress update
                    if (i + 1) % 50 == 0:
                        elapsed = time.time() - self.transfer.stats['start_time']
                        if elapsed > 0:
                            speed = self.transfer.stats['bytes_sent'] / elapsed / (1024**2)
                            progress = (i + 1) / len(future_to_file) * 100
                            print(f"Progress: {i+1}/{len(future_to_file)} files ({progress:.1f}%) | "
                                  f"Success: {successful_files}, Failed: {failed_files}, Skipped: {skipped_files} | "
                                  f"Speed: {speed:.1f} MB/s")
                        
                except KeyboardInterrupt:
                    print("\n⚠ Transfer interrupted by user")
                    executor.shutdown(wait=False)
                    break
                except Exception as e:
                    failed_files += 1
                    print(f"⚠ Unexpected error with {file_path}: {e}")
                    continue
        
        # Send completion signal
        try:
            completion_msg = json.dumps({
                'status': 'TRANSFER_COMPLETE',
                'successful': successful_files,
                'failed': failed_files,
                'skipped': skipped_files,
                'total': len(files)
            }).encode()
            self.send_data(sock, struct.pack('!I', len(completion_msg)) + completion_msg)
            print("Transfer completion signal sent")
        except Exception as e:
            print(f"Could not send completion signal: {e}")
            
        self.print_stats(successful_files, failed_files, skipped_files, len(files))
    
    def send_data(self, sock, data):
        """Send data with proper error handling and retry logic for large files"""
        total_sent = 0
        max_retries = 3
        retry_delay = 0.5
        
        while total_sent < len(data):
            for attempt in range(max_retries):
                try:
                    # Send data in smaller chunks to avoid timeouts
                    chunk_size = min(64*1024, len(data) - total_sent)  # 64KB chunks
                    sent = sock.send(data[total_sent:total_sent + chunk_size])
                    if sent == 0:
                        raise ConnectionError("Socket connection broken")
                    total_sent += sent
                    break  # Success, exit retry loop
                    
                except socket.timeout:
                    if attempt < max_retries - 1:
                        print(f"Send timeout (attempt {attempt + 1}/{max_retries}), retrying...")
                        time.sleep(retry_delay)
                        retry_delay *= 2  # Exponential backoff
                    else:
                        print("Send timeout - connection may be slow or file too large")
                        raise
                except Exception as e:
                    if attempt < max_retries - 1:
                        print(f"Send error (attempt {attempt + 1}/{max_retries}): {e}, retrying...")
                        time.sleep(retry_delay)
                    else:
                        print(f"Send error: {e}")
                        raise
    
    def prepare_file_data(self, file_info):
        """Prepare file data in worker thread - read and compress if needed"""
        filepath = file_info['full_path']
        rel_path = file_info['path']
        file_size = file_info['size']
        
        try:
            # Check file accessibility
            if not os.path.exists(filepath):
                print(f"⚠ File not found: {rel_path}")
                return None
            
            if not os.access(filepath, os.R_OK):
                print(f"⚠ No read permission: {rel_path}")
                return None
            
            # Verify file size hasn't changed
            current_size = os.path.getsize(filepath)
            if current_size != file_size:
                print(f"⚠ File size changed during scan: {rel_path} ({file_size} -> {current_size})")
                file_size = current_size
                file_info['size'] = current_size
            
            # Calculate checksum
            try:
                checksum = self.transfer.calculate_checksum(filepath)
            except Exception as e:
                print(f"⚠ Could not calculate checksum for {rel_path}: {e}")
                checksum = None
            
            # Prepare file header
            header = {
                'path': rel_path,
                'size': file_size,
                'checksum': checksum
            }
            header_json = json.dumps(header).encode()
            header_data = struct.pack('!I', len(header_json)) + header_json
            
            # Read and prepare file chunks
            chunks = []
            bytes_read = 0
            
            try:
                with open(filepath, 'rb') as f:
                    while bytes_read < file_size:
                        chunk_size = min(self.transfer.chunk_size, file_size - bytes_read)
                        chunk = f.read(chunk_size)
                        
                        if not chunk:
                            if bytes_read < file_size:
                                print(f"⚠ Unexpected EOF in {rel_path} at {bytes_read}/{file_size}")
                            break
                        
                        # Compression logic
                        try:
                            if self.transfer.compress and len(chunk) > 100:
                                compressed = zlib.compress(chunk, level=1)
                                if len(compressed) < len(chunk) * 0.9:
                                    chunk_data = struct.pack('!I?', len(compressed), True) + compressed
                                else:
                                    chunk_data = struct.pack('!I?', len(chunk), False) + chunk
                            else:
                                chunk_data = struct.pack('!I?', len(chunk), False) + chunk
                        except Exception as e:
                            print(f"⚠ Compression failed for chunk in {rel_path}: {e}")
                            chunk_data = struct.pack('!I?', len(chunk), False) + chunk
                        
                        chunks.append(chunk_data)
                        bytes_read += len(chunk)
                        
            except IOError as e:
                print(f"⚠ I/O error reading {rel_path}: {e}")
                return None
            except Exception as e:
                print(f"⚠ Unexpected error preparing {rel_path}: {e}")
                return None
            
            return {
                'header': header_data,
                'chunks': chunks,
                'size': bytes_read,
                'path': rel_path
            }
                
        except Exception as e:
            print(f"⚠ Critical error preparing {rel_path}: {e}")
            return None
    
    def send_prepared_file(self, sock, file_data):
        """Send pre-prepared file data over socket"""
        try:
            # Send header
            self.send_data(sock, file_data['header'])
            
            # Send all chunks
            for chunk_data in file_data['chunks']:
                self.send_data(sock, chunk_data)
            
            # Update statistics
            with self.transfer.stats_lock:
                self.transfer.stats['files_sent'] += 1
                self.transfer.stats['bytes_sent'] += file_data['size']
            
            return True
            
        except Exception as e:
            print(f"⚠ Error sending prepared file {file_data['path']}: {e}")
            return False
    
    def send_file(self, sock, file_info):
        """Send a single file with comprehensive error handling"""
        filepath = file_info['full_path']
        rel_path = file_info['path']
        file_size = file_info['size']
        
        try:
            # Check file accessibility
            if not os.path.exists(filepath):
                print(f"⚠ File not found: {rel_path}")
                return False
            
            if not os.access(filepath, os.R_OK):
                print(f"⚠ No read permission: {rel_path}")
                return False
            
            # Verify file size hasn't changed
            current_size = os.path.getsize(filepath)
            if current_size != file_size:
                print(f"⚠ File size changed during scan: {rel_path} ({file_size} -> {current_size})")
                file_size = current_size  # Use current size
                file_info['size'] = current_size
            
            # Calculate checksum with error handling
            try:
                checksum = self.transfer.calculate_checksum(filepath)
            except Exception as e:
                print(f"⚠ Could not calculate checksum for {rel_path}: {e}")
                checksum = None
            
            # Send file header
            header = {
                'path': rel_path,
                'size': file_size,
                'checksum': checksum
            }
            header_json = json.dumps(header).encode()
            header_data = struct.pack('!I', len(header_json)) + header_json
            
            try:
                self.send_data(sock, header_data)
            except Exception as e:
                print(f"⚠ Failed to send header for {rel_path}: {e}")
                return False
            
            # Send file content with retry logic
            bytes_sent = 0
            retry_count = 0
            max_retries = 3
            
            try:
                with open(filepath, 'rb') as f:
                    while bytes_sent < file_size:
                        try:
                            chunk_size = min(self.transfer.chunk_size, file_size - bytes_sent)
                            chunk = f.read(chunk_size)
                            
                            if not chunk:
                                if bytes_sent < file_size:
                                    print(f"⚠ Unexpected EOF in {rel_path} at {bytes_sent}/{file_size}")
                                break
                            
                            # Compression logic with error handling
                            try:
                                if self.transfer.compress and len(chunk) > 100:
                                    compressed = zlib.compress(chunk, level=1)
                                    if len(compressed) < len(chunk) * 0.9:
                                        chunk_data = struct.pack('!I?', len(compressed), True) + compressed
                                    else:
                                        chunk_data = struct.pack('!I?', len(chunk), False) + chunk
                                else:
                                    chunk_data = struct.pack('!I?', len(chunk), False) + chunk
                            except Exception as e:
                                print(f"⚠ Compression failed for chunk in {rel_path}: {e}")
                                chunk_data = struct.pack('!I?', len(chunk), False) + chunk
                            
                            # Send chunk with retry
                            chunk_sent = False
                            for attempt in range(max_retries):
                                try:
                                    self.send_data(sock, chunk_data)
                                    chunk_sent = True
                                    break
                                except Exception as e:
                                    if attempt < max_retries - 1:
                                        print(f"⚠ Retry {attempt + 1}/{max_retries} for chunk in {rel_path}: {e}")
                                        time.sleep(0.1)  # Brief pause before retry
                                    else:
                                        raise e
                            
                            if not chunk_sent:
                                print(f"⚠ Failed to send chunk after {max_retries} attempts: {rel_path}")
                                return False
                            
                            bytes_sent += len(chunk)
                            
                        except Exception as e:
                            print(f"⚠ Error reading/sending chunk from {rel_path}: {e}")
                            return False
                            
            except IOError as e:
                print(f"⚠ I/O error reading {rel_path}: {e}")
                return False
            except Exception as e:
                print(f"⚠ Unexpected error sending {rel_path}: {e}")
                return False
            
            # Update statistics
            with self.transfer.stats_lock:
                self.transfer.stats['files_sent'] += 1
                self.transfer.stats['bytes_sent'] += bytes_sent
            
            return True
                
        except Exception as e:
            print(f"⚠ Critical error with {rel_path}: {e}")
            return False
    
    def print_stats(self, successful=None, failed=None, skipped=None, total=None):
        """Print transfer statistics"""
        elapsed = time.time() - self.transfer.stats['start_time']
        total_mb = self.transfer.stats['bytes_sent'] / (1024**2)
        speed = total_mb / elapsed if elapsed > 0 else 0
        
        print(f"\n--- Transfer Complete ---")
        if successful is not None:
            print(f"Files processed: {total}")
            print(f"Successfully sent: {successful}")
            print(f"Failed to send: {failed}")
            print(f"Skipped (missing): {skipped}")
            success_rate = (successful / total * 100) if total > 0 else 0
            print(f"Success rate: {success_rate:.1f}%")
        else:
            print(f"Files transferred: {self.transfer.stats['files_sent']}")
        print(f"Data transferred: {total_mb:.1f} MB")
        print(f"Time elapsed: {elapsed:.1f} seconds")
        print(f"Average speed: {speed:.1f} MB/s")

class FileTransferClient:
    def __init__(self, host, port=9999, output_dir="./received", max_workers=8):
        self.host = host
        self.port = port
        self.output_dir = Path(output_dir)
        self.max_workers = max_workers
        self.transfer = FastFileTransfer(max_workers=max_workers)
        
        # Create output directory
        self.output_dir.mkdir(parents=True, exist_ok=True)
    
    def connect_and_receive(self):
        """Connect to server and receive files"""
        print(f"Connecting to {self.host}:{self.port}")
        
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 1024*1024)  # 1MB receive buffer
        sock.settimeout(60)  # 60 second connection timeout
        
        try:
            # Test connection first
            print("Attempting to connect...")
            sock.connect((self.host, self.port))
            print("✓ Connected to server successfully!")
            
            # Use configurable timeout if available
            timeout = getattr(self, '_socket_timeout', 1800)
            sock.settimeout(timeout)  # Use configurable timeout for large files
            
            # Receive manifest
            print("Receiving file manifest...")
            try:
                manifest_size_data = self.recv_exact(sock, 4)
                manifest_size = struct.unpack('!I', manifest_size_data)[0]
                print(f"Manifest size: {manifest_size} bytes")
                
                manifest_data = self.recv_exact(sock, manifest_size)
                manifest = json.loads(manifest_data.decode())
                
                files = manifest['files']
                total_size = manifest['total_size']
                compress = manifest['compress']
                
                print(f"✓ Manifest received: {len(files)} files ({total_size / (1024**3):.2f} GB)")
                
                # Receive files
                self.transfer.stats['start_time'] = time.time()
                self.receive_files(sock, files, compress)
                
            except socket.timeout:
                print("✗ Timeout while receiving data from server")
                print("This usually means:")
                print("  1. Server is not sending data")
                print("  2. Network connection is too slow")
                print("  3. Server crashed or disconnected")
                
        except ConnectionRefusedError:
            print("✗ Connection refused!")
            print("Possible causes:")
            print("  1. Server is not running")
            print("  2. Wrong IP address or port")
            print("  3. Firewall blocking connection")
            print(f"  4. Make sure server is running on {self.host}:{self.port}")
            
        except socket.gaierror:
            print("✗ Cannot resolve hostname!")
            print(f"Check if '{self.host}' is a valid IP address or hostname")
            
        except socket.timeout:
            print("✗ Connection timeout!")
            print("Possible causes:")
            print("  1. Server is not responding")
            print("  2. Network connectivity issues")
            print("  3. Wrong IP address")
            
        except Exception as e:
            print(f"✗ Connection error: {e}")
            
        finally:
            sock.close()
    
    def recv_exact(self, sock, num_bytes):
        """Receive exactly num_bytes from socket with improved timeout handling"""
        data = b""
        max_retries = 3
        retry_delay = 1.0
        
        while len(data) < num_bytes:
            for attempt in range(max_retries):
                try:
                    # Receive data in smaller chunks to avoid timeouts
                    remaining = num_bytes - len(data)
                    chunk_size = min(64*1024, remaining)  # 64KB chunks
                    chunk = sock.recv(chunk_size)
                    
                    if not chunk:
                        raise ConnectionError("Server disconnected")
                    data += chunk
                    break  # Success, exit retry loop
                    
                except socket.timeout:
                    if attempt < max_retries - 1:
                        print(f"Receive timeout (attempt {attempt + 1}/{max_retries}) - got {len(data)}/{num_bytes} bytes, retrying...")
                        time.sleep(retry_delay)
                        retry_delay *= 2  # Exponential backoff
                    else:
                        print(f"Timeout while receiving data (got {len(data)}/{num_bytes} bytes)")
                        print("This may indicate:")
                        print("  1. Network is too slow for large files")
                        print("  2. Server is overloaded")
                        print("  3. File is extremely large")
                        print("  4. Try reducing --workers or using --no-compress")
                        raise
                except Exception as e:
                    if attempt < max_retries - 1:
                        print(f"Receive error (attempt {attempt + 1}/{max_retries}): {e}, retrying...")
                        time.sleep(retry_delay)
                    else:
                        raise
        return data
    
    def receive_files(self, sock, files, compress):
        """Receive all files from server - Parallel version using workers"""
        received_files = 0
        failed_files = 0
        skipped_files = 0
        total_files = len(files)
        
        print("Starting file reception...")
        print(f"Using {self.max_workers} worker threads for file processing")
        
        # Use ThreadPoolExecutor for parallel file writing
        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            futures = []
            
            for i in range(total_files):
                try:
                    # Receive file data from socket
                    file_data = self.receive_file_data(sock, compress)
                    
                    if file_data is None:
                        failed_files += 1
                        continue
                    
                    # Submit file writing task to worker thread
                    future = executor.submit(self.write_received_file, file_data)
                    futures.append((future, file_data['path'], i))
                    
                except KeyboardInterrupt:
                    print("\n⚠ Reception interrupted by user")
                    executor.shutdown(wait=False)
                    break
                except Exception as e:
                    print(f"⚠ Unexpected error receiving file {i+1}: {e}")
                    failed_files += 1
                    continue
            
            # Wait for all file writing tasks to complete
            for future, file_path, index in futures:
                try:
                    result = future.result()
                    if result == "success":
                        received_files += 1
                    elif result == "failed":
                        failed_files += 1
                    elif result == "skipped":
                        skipped_files += 1
                    
                    # Progress update
                    if (index + 1) % 50 == 0:
                        elapsed = time.time() - self.transfer.stats['start_time']
                        speed = self.transfer.stats['bytes_sent'] / elapsed / (1024**2) if elapsed > 0 else 0
                        progress = (index + 1) / total_files * 100
                        print(f"Progress: {index+1}/{total_files} files ({progress:.1f}%) | "
                              f"Success: {received_files}, Failed: {failed_files}, Skipped: {skipped_files} | "
                              f"Speed: {speed:.1f} MB/s")
                        
                except Exception as e:
                    print(f"⚠ Error processing {file_path}: {e}")
                    failed_files += 1
        
        # Try to receive completion signal
        try:
            completion_size = struct.unpack('!I', self.recv_exact(sock, 4))[0]
            completion_data = self.recv_exact(sock, completion_size)
            completion_info = json.loads(completion_data.decode())
            
            if completion_info.get('status') == 'TRANSFER_COMPLETE':
                print("✓ Received transfer completion signal")
                server_stats = completion_info
                print(f"Server reported - Success: {server_stats.get('successful', 'N/A')}, "
                      f"Failed: {server_stats.get('failed', 'N/A')}, "
                      f"Skipped: {server_stats.get('skipped', 'N/A')}")
        except Exception as e:
            print(f"Note: Could not receive completion signal: {e}")
        
        print(f"\n--- Reception Complete ---")
        print(f"Files processed: {total_files}")
        print(f"Successfully received: {received_files}")
        print(f"Failed: {failed_files}")
        print(f"Skipped: {skipped_files}")
        success_rate = (received_files / total_files * 100) if total_files > 0 else 0
        print(f"Success rate: {success_rate:.1f}%")
        
        elapsed = time.time() - self.transfer.stats['start_time']
        total_mb = self.transfer.stats['bytes_sent'] / (1024**2)
        speed = total_mb / elapsed if elapsed > 0 else 0
        print(f"Total data: {total_mb:.1f} MB")
        print(f"Time elapsed: {elapsed:.1f} seconds")
        print(f"Average speed: {speed:.1f} MB/s")
    
    def receive_file_data(self, sock, compress):
        """Receive file data from socket (runs in main thread)"""
        try:
            # Receive file header
            try:
                header_size = struct.unpack('!I', self.recv_exact(sock, 4))[0]
                header_data = self.recv_exact(sock, header_size)
                header = json.loads(header_data.decode())
            except Exception as e:
                print(f"⚠ Failed to receive file header: {e}")
                return None
            
            rel_path = header['path']
            file_size = header['size']
            expected_checksum = header.get('checksum')
            
            # Receive all file chunks
            chunks = []
            bytes_received = 0
            
            try:
                while bytes_received < file_size:
                    # Receive chunk header
                    chunk_header = self.recv_exact(sock, 5)  # 4 bytes size + 1 byte compression flag
                    chunk_size, is_compressed = struct.unpack('!I?', chunk_header)
                    
                    # Receive chunk data
                    chunk_data = self.recv_exact(sock, chunk_size)
                    
                    chunks.append((chunk_data, is_compressed))
                    
                    # Calculate bytes received (decompressed size)
                    if is_compressed:
                        try:
                            decompressed = zlib.decompress(chunk_data)
                            bytes_received += len(decompressed)
                        except:
                            bytes_received += len(chunk_data)  # Fallback
                    else:
                        bytes_received += len(chunk_data)
                        
            except Exception as e:
                print(f"⚠ Error receiving chunks for {rel_path}: {e}")
                return None
            
            return {
                'path': rel_path,
                'size': file_size,
                'checksum': expected_checksum,
                'chunks': chunks
            }
            
        except Exception as e:
            print(f"⚠ Critical error receiving file data: {e}")
            return None
    
    def write_received_file(self, file_data):
        """Write received file data to disk (runs in worker thread)"""
        rel_path = file_data['path']
        file_size = file_data['size']
        expected_checksum = file_data['checksum']
        chunks = file_data['chunks']
        
        try:
            # Create output path
            output_path = self.output_dir / rel_path
            output_path.parent.mkdir(parents=True, exist_ok=True)
            
            # Check write access
            try:
                test_path = output_path.parent / f".write_test_{int(time.time())}_{threading.current_thread().ident}"
                test_path.touch()
                test_path.unlink()
            except Exception as e:
                print(f"⚠ No write permission for {rel_path}: {e}")
                return "failed"
            
            # Write file data
            temp_path = output_path.with_suffix(output_path.suffix + f'.tmp_{threading.current_thread().ident}')
            
            try:
                with open(temp_path, 'wb') as f:
                    for chunk_data, is_compressed in chunks:
                        try:
                            if is_compressed:
                                chunk = zlib.decompress(chunk_data)
                            else:
                                chunk = chunk_data
                        except Exception as e:
                            print(f"⚠ Decompression failed for {rel_path}: {e}")
                            return "failed"
                        
                        f.write(chunk)
                
                # Rename temp file to final name (thread-safe)
                if temp_path.exists():
                    if output_path.exists():
                        output_path.unlink()
                    temp_path.rename(output_path)
                
            except IOError as e:
                print(f"⚠ I/O error writing {rel_path}: {e}")
                if temp_path.exists():
                    try:
                        temp_path.unlink()
                    except:
                        pass
                return "failed"
            except Exception as e:
                print(f"⚠ Unexpected error writing {rel_path}: {e}")
                if temp_path.exists():
                    try:
                        temp_path.unlink()
                    except:
                        pass
                return "failed"
            
            # Verify file size
            try:
                actual_size = output_path.stat().st_size
                if actual_size != file_size:
                    print(f"⚠ Size mismatch for {rel_path}: expected {file_size}, got {actual_size}")
                    return "failed"
            except Exception as e:
                print(f"⚠ Cannot verify size for {rel_path}: {e}")
                return "failed"
            
            # Verify checksum if available
            if expected_checksum:
                try:
                    received_checksum = self.transfer.calculate_checksum(output_path)
                    if received_checksum != expected_checksum:
                        print(f"⚠ Checksum mismatch for {rel_path}")
                        return "failed"
                except Exception as e:
                    print(f"⚠ Cannot verify checksum for {rel_path}: {e}")
            
            with self.transfer.stats_lock:
                self.transfer.stats['bytes_sent'] += file_size
            
            return "success"
            
        except Exception as e:
            print(f"⚠ Critical error writing file {rel_path}: {e}")
            return "failed"

def test_connection(host, port):
    """Test if server is reachable"""
    print(f"Testing connection to {host}:{port}...")
    
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(5)
        result = sock.connect_ex((host, port))
        sock.close()
        
        if result == 0:
            print("✓ Connection successful - server is reachable!")
            return True
        else:
            print("✗ Connection failed - server is not reachable")
            return False
            
    except Exception as e:
        print(f"✗ Connection test failed: {e}")
        return False

def main():
    parser = argparse.ArgumentParser(description='Ultra-fast file transfer system')
    parser.add_argument('mode', choices=['server', 'client', 'test'], help='Run as server, client, or test connection')
    parser.add_argument('--host', default='localhost', help='Server host (default: localhost)')
    parser.add_argument('--port', type=int, default=9999, help='Port number (default: 9999)')
    parser.add_argument('--folder', help='Folder to transfer (server mode)')
    parser.add_argument('--output', default='./received', help='Output directory (client mode)')
    parser.add_argument('--workers', type=int, default=8, help='Number of worker threads')
    parser.add_argument('--no-compress', action='store_true', help='Disable compression')
    parser.add_argument('--chunk-size', type=int, default=1, help='Chunk size in bytes (default: 1MB)')
    parser.add_argument('--timeout', type=int, default=1800, help='Socket timeout in seconds (default: 1800 = 30 min)')
    
    args = parser.parse_args()
    
    if args.mode == 'test':
        test_connection(args.host, args.port)
        
    elif args.mode == 'server':
        if not args.folder:
            print("Error: --folder required in server mode")
            return
        
        if not os.path.exists(args.folder):
            print(f"Error: Folder '{args.folder}' does not exist")
            return
            
        print(f"Server will bind to: {args.host}:{args.port}")
        print(f"Serving folder: {os.path.abspath(args.folder)}")
        print(f"Using {args.workers} worker threads")
        print(f"Chunk size: {args.chunk_size / (1024*1024):.1f} MB")
        print(f"Socket timeout: {args.timeout} seconds")
        
        server = FileTransferServer(args.host, args.port, args.workers)
        server.transfer.compress = not args.no_compress
        server.transfer.chunk_size = args.chunk_size*1024*1024
        
        # Apply timeout to server socket operations
        server._socket_timeout = args.timeout
        server.start_server(args.folder)
    
    elif args.mode == 'client':
        print(f"Client will connect to: {args.host}:{args.port}")
        print(f"Output directory: {os.path.abspath(args.output)}")
        print(f"Using {args.workers} worker threads")
        print(f"Chunk size: {args.chunk_size / (1024*1024):.1f} MB")
        print(f"Socket timeout: {args.timeout} seconds")
        
        client = FileTransferClient(args.host, args.port, args.output, args.workers)
        client.transfer.compress = not args.no_compress
        client.transfer.chunk_size = args.chunk_size*1024*1024
        
        # Apply timeout to client socket operations
        client._socket_timeout = args.timeout
        client.connect_and_receive()

if __name__ == "__main__":
    main()