bluetooth_image_protocol.py

#!/usr/bin/env python3
"""
Bluetooth Image Transfer Protocol
=================================

Protocol for transferring images over Bluetooth mesh network using Base64 encoding and chunking.

Message Format:
- Text: [ID:sender:recipient] message
- Image Start: [ID:sender:recipient:IMG_START] filename.ext|base64_chunk_count|original_size
- Image Chunk: [ID:sender:recipient:IMG_CHUNK] chunk_index|base64_data
- Image End: [ID:sender:recipient:IMG_END] filename.ext|checksum

Usage:
    from bluetooth_image_protocol import BluetoothImageTransfer
    
    transfer = BluetoothImageTransfer()
    
    # Send image
    chunks = transfer.prepare_image_for_transfer("path/to/image.jpg")
    for chunk in chunks:
        socket.send(chunk.encode())
    
    # Receive image
    transfer.process_received_message(received_message)
    if transfer.is_image_complete(image_id):
        image_data = transfer.get_completed_image(image_id)
"""

import base64
import os
import hashlib
import time
from typing import List, Dict, Optional, Tuple
import uuid

class BluetoothImageTransfer:
    def __init__(self, max_chunk_size: int = 800):
        """
        Initialize the image transfer handler
        
        Args:
            max_chunk_size: Maximum size of each chunk in bytes (leaving room for headers)
        """
        self.max_chunk_size = max_chunk_size
        self.incoming_images: Dict[str, Dict] = {}  # image_id -> {chunks, metadata}
        self.outgoing_images: Dict[str, Dict] = {}  # image_id -> {chunks, sent_count}
        
    def prepare_image_for_transfer(self, image_path: str, sender: str, recipient: str) -> List[str]:
        """
        Prepare an image file for Bluetooth transfer by converting to Base64 and chunking
        
        Args:
            image_path: Path to the image file
            sender: Sender's name
            recipient: Recipient's name (or 'broadcast')
            
        Returns:
            List of message strings ready to be sent over Bluetooth
        """
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"Image file not found: {image_path}")
            
        # Read and encode the image
        with open(image_path, 'rb') as f:
            image_data = f.read()
            
        # Get file info
        filename = os.path.basename(image_path)
        file_size = len(image_data)
        base64_data = base64.b64encode(image_data).decode('utf-8')
        
        # Calculate checksum
        checksum = hashlib.md5(image_data).hexdigest()[:8]
        
        # Create unique message ID
        msg_id = f"{sender[:3].upper()}{int(time.time() % 1000):03d}"
        
        print(f"🖼️  PREPARING IMAGE FOR BLUETOOTH TRANSFER")
        print(f"   📁 File: {filename}")
        print(f"   📊 Size: {file_size:,} bytes")
        print(f"   🔢 Base64 size: {len(base64_data):,} characters")
        print(f"   📦 Total chunks: {total_chunks}")
        print(f"   🔍 Checksum: {checksum}")
        print(f"   🆔 Message ID: {msg_id}")
        
        # Split into chunks
        chunks = []
        chunk_size = self.max_chunk_size
        total_chunks = (len(base64_data) + chunk_size - 1) // chunk_size
        
        messages = []
        
        # Start message
        start_msg = f"[{msg_id}:{sender}:{recipient}:IMG_START] {filename}|{total_chunks}|{file_size}|{checksum}"
        messages.append(start_msg)
        print(f"   📤 Created START message")
        
        # Chunk messages
        print(f"   📦 Creating {total_chunks} chunk messages...")
        for i in range(total_chunks):
            start_idx = i * chunk_size
            end_idx = min(start_idx + chunk_size, len(base64_data))
            chunk_data = base64_data[start_idx:end_idx]
            
            chunk_msg = f"[{msg_id}:{sender}:{recipient}:IMG_CHUNK] {i}|{chunk_data}"
            messages.append(chunk_msg)
            
            if (i + 1) % 10 == 0:  # Log every 10 chunks
                print(f"      📦 Created chunks 1-{i+1}/{total_chunks}")
                
        # End message
        end_msg = f"[{msg_id}:{sender}:{recipient}:IMG_END] {filename}|{checksum}"
        messages.append(end_msg)
        print(f"   📤 Created END message")
        print(f"   ✅ Total messages prepared: {len(messages)}")
        
        # Store for tracking
        self.outgoing_images[msg_id] = {
            'filename': filename,
            'total_chunks': total_chunks,
            'sent_count': 0,
            'checksum': checksum
        }
        
        return messages
        
    def process_received_message(self, message: str) -> Optional[Dict]:
        """
        Process a received message and handle image reconstruction
        
        Args:
            message: The received message string
            
        Returns:
            Dictionary with message info if it's an image message, None otherwise
        """
        message = message.strip()
        
        if not message.startswith('[') or ']' not in message:
            return None
            
        try:
            header_end = message.index(']')
            header = message[1:header_end]
            content = message[header_end + 1:].strip()
            
            parts = header.split(':')
            if len(parts) < 4:
                return None
                
            msg_id, sender, recipient, msg_type = parts[:4]
            
            if msg_type == 'IMG_START':
                return self._handle_image_start(msg_id, sender, recipient, content)
            elif msg_type == 'IMG_CHUNK':
                return self._handle_image_chunk(msg_id, sender, recipient, content)
            elif msg_type == 'IMG_END':
                return self._handle_image_end(msg_id, sender, recipient, content)
                
        except (ValueError, IndexError) as e:
            print(f"⚠️ Error parsing image message: {e}")
            return None
            
        return None
        
    def _handle_image_start(self, msg_id: str, sender: str, recipient: str, content: str) -> Dict:
        """Handle IMG_START message"""
        try:
            parts = content.split('|')
            if len(parts) < 4:
                raise ValueError("Invalid IMG_START format")
                
            filename, total_chunks, file_size, checksum = parts
            total_chunks = int(total_chunks)
            file_size = int(file_size)
            
            start_time = time.time()
            self.incoming_images[msg_id] = {
                'sender': sender,
                'recipient': recipient,
                'filename': filename,
                'total_chunks': total_chunks,
                'file_size': file_size,
                'checksum': checksum,
                'chunks': {},
                'received_chunks': 0,
                'start_time': start_time,
                'milestone_25': False,
                'milestone_50': False,
                'milestone_75': False,
            }
            
            print(f"\n📸 Starting new image transfer")
            print(f"   📄 File: {filename}")
            print(f"   📊 Size: {file_size / 1024:.1f} KB")
            print(f"   📦 Total chunks: {total_chunks}")
            print(f"   📨 From: {sender} → {recipient}")
            print(f"   🆔 Transfer ID: {msg_id}")
            print(f"   ⏱️ Started at: {time.strftime('%H:%M:%S', time.localtime(start_time))}")
            
            return {
                'type': 'image_start',
                'msg_id': msg_id,
                'sender': sender,
                'recipient': recipient,
                'filename': filename,
                'total_chunks': total_chunks,
                'file_size': file_size
            }
            
        except (ValueError, IndexError) as e:
            print(f"❌ Error handling IMG_START: {e}")
            return None
            
    def _handle_image_chunk(self, msg_id: str, sender: str, recipient: str, content: str) -> Dict:
        """Handle IMG_CHUNK message"""
        if msg_id not in self.incoming_images:
            print(f"⚠️ Received chunk for unknown image: {msg_id}")
            return None
            
        try:
            pipe_idx = content.index('|')
            chunk_index = int(content[:pipe_idx])
            chunk_data = content[pipe_idx + 1:]
            
            image_info = self.incoming_images[msg_id]
            image_info['chunks'][chunk_index] = chunk_data
            
            # Calculate progress statistics
            image_info['received_chunks'] = len(image_info['chunks'])
            progress = (image_info['received_chunks'] / image_info['total_chunks']) * 100
            
            # Log progress milestones
            if progress >= 25 and not image_info.get('milestone_25'):
                print(f"📊 Image transfer 25% complete ({image_info['filename']}) - {msg_id}")
                image_info['milestone_25'] = True
            elif progress >= 50 and not image_info.get('milestone_50'):
                print(f"📊 Image transfer 50% complete ({image_info['filename']}) - {msg_id}")
                image_info['milestone_50'] = True
            elif progress >= 75 and not image_info.get('milestone_75'):
                print(f"📊 Image transfer 75% complete ({image_info['filename']}) - {msg_id}")
                image_info['milestone_75'] = True
                
            # Log every 10% for detailed progress
            if chunk_index % (image_info['total_chunks'] // 10) == 0:
                print(f"   📦 Received chunk {chunk_index}/{image_info['total_chunks']} ({progress:.1f}%)")
            image_info['received_chunks'] += 1
            
            progress = (image_info['received_chunks'] / image_info['total_chunks']) * 100
            
            # More detailed progress logging
            if chunk_index == 0:
                print(f"📦 First chunk received for {image_info['filename']}")
            elif progress >= 25 and not hasattr(image_info, '_25_logged'):
                print(f"📊 25% progress: {image_info['received_chunks']}/{image_info['total_chunks']} chunks")
                image_info['_25_logged'] = True
            elif progress >= 50 and not hasattr(image_info, '_50_logged'):
                print(f"📊 50% progress: {image_info['received_chunks']}/{image_info['total_chunks']} chunks")
                image_info['_50_logged'] = True
            elif progress >= 75 and not hasattr(image_info, '_75_logged'):
                print(f"📊 75% progress: {image_info['received_chunks']}/{image_info['total_chunks']} chunks")
                image_info['_75_logged'] = True
            elif (chunk_index + 1) % 10 == 0:  # Log every 10 chunks
                print(f"📦 Received chunk {chunk_index + 1}/{image_info['total_chunks']} ({progress:.1f}%)")
            else:
                print(f"📦 Received chunk {chunk_index + 1}/{image_info['total_chunks']} ({progress:.1f}%)")
            
            return {
                'type': 'image_chunk',
                'msg_id': msg_id,
                'sender': sender,
                'chunk_index': chunk_index,
                'progress': progress,
                'received_chunks': image_info['received_chunks'],
                'total_chunks': image_info['total_chunks']
            }
            
        except (ValueError, IndexError) as e:
            print(f"❌ Error handling IMG_CHUNK: {e}")
            return None
            
    def _handle_image_end(self, msg_id: str, sender: str, recipient: str, content: str) -> Dict:
        """Handle IMG_END message"""
        if msg_id not in self.incoming_images:
            print(f"⚠️ Received end for unknown image: {msg_id}")
            return None
            
        try:
            parts = content.split('|')
            if len(parts) < 2:
                raise ValueError("Invalid IMG_END format")
                
            filename, expected_checksum = parts
            image_info = self.incoming_images[msg_id]
            
            print(f"\n🔍 Finalizing image transfer for {filename}")
            print(f"   📦 Total chunks received: {image_info['received_chunks']}/{image_info['total_chunks']}")
            
            # Calculate transfer statistics
            transfer_time = time.time() - image_info.get('start_time', 0)
            transfer_size = image_info.get('file_size', 0)
            transfer_speed = transfer_size / transfer_time if transfer_time > 0 else 0
            
            print(f"   ⏱️ Transfer time: {transfer_time:.1f} seconds")
            print(f"   📊 Transfer speed: {transfer_speed / 1024:.1f} KB/s")
            
            # Check if we have all chunks
            if image_info['received_chunks'] != image_info['total_chunks']:
                print(f"❌ Image incomplete: {image_info['received_chunks']}/{image_info['total_chunks']} chunks")
                return {
                    'type': 'image_error',
                    'msg_id': msg_id,
                    'error': 'Incomplete image transfer'
                }
                
            # Reconstruct the image
            base64_data = ''
            for i in range(image_info['total_chunks']):
                if i not in image_info['chunks']:
                    print(f"❌ Missing chunk {i}")
                    return {
                        'type': 'image_error',
                        'msg_id': msg_id,
                        'error': f'Missing chunk {i}'
                    }
                base64_data += image_info['chunks'][i]
                
            # Decode and verify
            try:
                image_data = base64.b64decode(base64_data)
                actual_checksum = hashlib.md5(image_data).hexdigest()[:8]
                
                if actual_checksum != expected_checksum:
                    print(f"❌ Checksum mismatch: expected {expected_checksum}, got {actual_checksum}")
                    return {
                        'type': 'image_error',
                        'msg_id': msg_id,
                        'error': 'Checksum verification failed'
                    }
                    
                # Store the completed image
                image_info['image_data'] = image_data
                image_info['completed_at'] = time.time()
                transfer_time = image_info['completed_at'] - image_info['started_at']
                
                print(f"✅ Image transfer complete: {filename} ({len(image_data)} bytes in {transfer_time:.1f}s)")
                print(f"   📊 Transfer rate: {len(image_data) / transfer_time:.0f} bytes/sec")
                print(f"   🔍 Checksum verified: {actual_checksum}")
                
                return {
                    'type': 'image_complete',
                    'msg_id': msg_id,
                    'sender': sender,
                    'recipient': recipient,
                    'filename': filename,
                    'image_data': image_data,
                    'file_size': len(image_data),
                    'transfer_time': transfer_time
                }
                
            except Exception as e:
                print(f"❌ Error decoding image: {e}")
                return {
                    'type': 'image_error',
                    'msg_id': msg_id,
                    'error': f'Decoding failed: {str(e)}'
                }
                
        except (ValueError, IndexError) as e:
            print(f"❌ Error handling IMG_END: {e}")
            return None
            
    def is_image_complete(self, msg_id: str) -> bool:
        """Check if an image transfer is complete"""
        if msg_id not in self.incoming_images:
            return False
        return 'image_data' in self.incoming_images[msg_id]
        
    def get_completed_image(self, msg_id: str) -> Optional[Tuple[str, bytes]]:
        """Get completed image data"""
        if not self.is_image_complete(msg_id):
            return None
            
        image_info = self.incoming_images[msg_id]
        return image_info['filename'], image_info['image_data']
        
    def cleanup_completed_transfer(self, msg_id: str):
        """Clean up completed transfer data"""
        if msg_id in self.incoming_images:
            del self.incoming_images[msg_id]
        if msg_id in self.outgoing_images:
            del self.outgoing_images[msg_id]
            
    def get_transfer_status(self) -> Dict:
        """Get status of all ongoing transfers"""
        return {
            'incoming': {
                msg_id: {
                    'filename': info['filename'],
                    'progress': (info['received_chunks'] / info['total_chunks']) * 100,
                    'sender': info['sender']
                } for msg_id, info in self.incoming_images.items()
            },
            'outgoing': {
                msg_id: {
                    'filename': info['filename'],
                    'sent_chunks': info['sent_count'],
                    'total_chunks': info['total_chunks']
                } for msg_id, info in self.outgoing_images.items()
            }
        }