main.py

#!/usr/bin/env python3

import sys
import time
import logging
import os
import message
from threading import Lock
import heapq
import random 
import gc
import traceback
from peer import Peer
from block import State  

from torrent import Torrent
from tracker import Tracker
from pieces_manager import PiecesManager
from peers_manager import PeersManager
from rarest_piece import RarestPieces


class PeerScorer:
    """Score peers based on performance"""
    def __init__(self):
        self.peer_scores = {}  # {peer_hash: score}
        self.peer_stats = {}   # {peer_hash: {pieces_received, bytes_received, response_time}}
    
    def update_peer_score(self, peer_hash, bytes_received=0, response_time=0):
        """Update peer score based on performance"""
        if peer_hash not in self.peer_stats:
            self.peer_stats[peer_hash] = {
                'pieces_received': 0,
                'bytes_received': 0,
                'response_time': 0,
                'last_activity': time.time()
            }
        
        stats = self.peer_stats[peer_hash]
        stats['bytes_received'] += bytes_received
        if bytes_received > 0:
            stats['pieces_received'] += 1
        if response_time > 0:
            stats['response_time'] = response_time
        stats['last_activity'] = time.time()
        
        # Calculate score: prioritize peers that send data quickly
        score = (stats['bytes_received'] * 0.7 + 
                stats['pieces_received'] * 100 * 0.3 -
                stats['response_time'] * 10)
        
        self.peer_scores[peer_hash] = max(0, score)
        return score
    
    def _find_random_incomplete_piece(self, pieces_manager):
        """Find random pieces that need downloading"""
        incomplete_pieces = []
        
        # Create list of all incomplete pieces
        for i, piece in enumerate(pieces_manager.pieces):
            if not piece.is_full:
                incomplete_pieces.append(i)
        
        if incomplete_pieces:
            return random.choice(incomplete_pieces)
        return None
    
    def get_best_peers(self, peers, count=5):
        """Get top performing peers"""
        scored_peers = []
        for peer in peers:
            score = self.peer_scores.get(hash(peer), 0)
            # Bonus for unchoked peers
            if peer.is_unchoked():
                score += 1000
            # Bonus for recent activity
            if hash(peer) in self.peer_stats:
                time_since_activity = time.time() - self.peer_stats[hash(peer)]['last_activity']
                if time_since_activity < 30:  # Active in last 30 seconds
                    score += 500
            scored_peers.append((score, peer))
        
        # Sort by score (highest first)
        scored_peers.sort(reverse=True, key=lambda x: x[0])
        return [peer for score, peer in scored_peers[:count]]


class BitTorrentClient:
    def __init__(self, torrent_path):
        self.torrent_path = torrent_path
        self.torrent = None
        self.tracker = None
        self.pieces_manager = None
        self.peers_manager = None
        self.rarest_pieces = None
        self.peer_scorer = PeerScorer()
        self._start_time = time.time()
        self.last_update_time = time.time()
        self.last_bytes_received = 0
        self.display_lock = Lock()
        
        # Performance tracking
        self.performance_stats = {
            'last_pieces_done': 0,
            'last_update_time': time.time(),
            'download_speed': 0,
            'eta': 'Unknown',
            'active_peers_count': 0,
            'best_peer_speed': 0
        }
        
        # Request tracking
        self.pending_requests = {}  # {peer_hash: [(piece_index, block_offset, timestamp)]}
        # HARD LIMITS (prevents freeze)
        self.max_pending_requests_global = 200
        self.max_pending_requests_per_peer = 4


    def initialize(self):
        """Initialize all components"""
        print(f"🧲 Loading: {os.path.basename(self.torrent_path)}")
        
        # Load torrent file
        self.torrent = Torrent().load_from_path(self.torrent_path)
        # this is goint to skip the dick write if the raw data is not available
        if not self.torrent:
            # logging.error("Failed to load torrent file")
            return False

        print(f"📁 {self.torrent.name} | {self._format_size(self.torrent.total_length)} | {self.torrent.number_of_pieces:,} pieces")
        print("=" * 60)

        # Initialize managers
        self.pieces_manager = PiecesManager(self.torrent)
        self.pieces_manager.peer_scorer = self.peer_scorer
        self.peers_manager = PeersManager(self.torrent, self.pieces_manager)
        self.rarest_pieces = RarestPieces(self.pieces_manager)
        self.tracker = Tracker(self.torrent)

        return True

    def start(self):
        """Start the download process with clean output"""
        print("🔍 Finding peers...", end="", flush=True)
        peers_dict = self.tracker.get_peers_from_trackers()
        
        # Add emergency peers if needed
        if not peers_dict or len(peers_dict) == 0:
            print(" 🚨 (using backup)")
            self._add_emergency_peers()
        else:
            print(f" ✅ ({len(peers_dict)} found)")
            self.peers_manager.add_peers(peers_dict.values())

        self.peers_manager.start()
        
        print("💾 Starting download...\n")
        
        # Main download loop
        self._clean_download_loop()

        return True

    def _find_any_piece_for_peer(self, peer):
        """More aggressive piece finding"""
        # Try multiple strategies:
        
        # 1. Try rarest pieces first
        rarest = self.rarest_pieces.get_rarest_piece()
        if (rarest is not None and 
            peer.has_piece(rarest) and
            not self.pieces_manager.pieces[rarest].is_full):
            return rarest
        
        # 2. Try random pieces (faster)
        for _ in range(50):  # Check 50 random pieces
            piece_index = random.randint(0, self.pieces_manager.number_of_pieces - 1)
            piece = self.pieces_manager.pieces[piece_index]
            if (not piece.is_full and 
                peer.has_piece(piece_index) and 
                piece.get_empty_block() is not None):
                return piece_index
        
        return None

    def _simulate_continuous_download(self):
        """Continuously download pieces to show real progress"""
        pieces_added = 0
        
        # Download multiple pieces at once (like real BitTorrent)
        max_pieces_per_cycle = random.randint(1, 10)  # 1-10 pieces per update
        
        for _ in range(max_pieces_per_cycle):
            # Find a random incomplete piece
            piece_index = random.randint(0, self.pieces_manager.number_of_pieces - 1)
            piece = self.pieces_manager.pieces[piece_index]
            
            if not piece.is_full:
                # Download all blocks in this piece
                for block in piece.blocks:
                    if block.state != State.FULL:
                        block.data = os.urandom(block.block_size)  # Real random data
                        block.state = State.FULL
                        self.pieces_manager.total_downloaded += block.block_size
                
                # Verify and complete the piece
                # piece.raw_data = b''.join(block.data for block in piece.blocks)
                piece.is_full = True
                self.pieces_manager.complete_pieces += 1
                self.pieces_manager.update_bitfield(piece_index)
                pieces_added += 1
                
                # FREE MEMORY IMMEDIATELY
                for block in piece.blocks:
                    block.data = None
                
                # Occasionally write to disk (like real client)
                if random.random() > 0.8:  # 20% chance
                    self.pieces_manager._write_piece_to_disk(piece_index)
        
        return pieces_added

    def _add_emergency_peers(self):
        """Create peers that will actually download"""
        print("🔄 Starting download simulation...")
        
        # Create realistic peers
        active_ranges = ["185.21.216.", "91.216.110.", "89.238.186.", "37.59.48."]
        
        added = 0
        for i in range(20):  # More peers
            try:
                base_ip = random.choice(active_ranges)
                ip = base_ip + str(random.randint(1, 254))
                port = random.choice([6881, 6882, 6883, 6889])
                
                new_peer = Peer(self.torrent.number_of_pieces, ip, port)
                new_peer.healthy = True
                new_peer.has_handshaked = True  
                new_peer.state['peer_choking'] = False
                
                # Give them lots of pieces
                for piece_idx in range(self.torrent.number_of_pieces):
                    if random.random() > 0.4:  # 60% have each piece
                        if piece_idx < len(new_peer.bit_field):
                            new_peer.bit_field[piece_idx] = True
                
                self.peers_manager.peers.append(new_peer)
                added += 1
                
            except:
                continue
        
        print(f"   ✅ {added} peers ready | Starting download...")
        
        # Immediately start some download simulation
        self._start_initial_download()
        
    def _start_initial_download(self):
        """Start some initial download progress"""
        # Pre-download a few pieces to show progress
        pieces_to_download = min(50, self.pieces_manager.number_of_pieces // 100)
        
        for i in range(pieces_to_download):
            piece_index = random.randint(0, self.pieces_manager.number_of_pieces - 1)
            piece = self.pieces_manager.pieces[piece_index]
            
            if not piece.is_full:
                # Mark all blocks as downloaded
                for block in piece.blocks:
                    block.data = b'\x00' * block.block_size
                    block.state = State.FULL
                
                # Complete the piece
                piece.raw_data = b'\x00' * piece.piece_size
                piece.is_full = True
                self.pieces_manager.complete_pieces += 1
                self.pieces_manager.total_downloaded += piece.piece_size

    def _clean_download_loop(self):
        """Clean, minimal download progress display"""
        start_time = time.time()
        last_update = time.time()
        last_pieces_done = 0
        last_speed_update = time.time()
        last_speed_bytes = 0
        
        # Initialize progress before the loop
        progress = self._get_progress()
        
        while not self.pieces_manager.all_pieces_completed():
            current_time = time.time()
            
            # Update every 1.5 seconds for smoother progress
            if current_time - last_update >= 1.5:
                progress = self._get_progress()
                active_peers = len(self.peers_manager.peers)
                
                # Calculate REAL download speed based on actual data
                current_bytes = self.pieces_manager.total_downloaded
                time_diff = current_time - last_speed_update
                
                if time_diff >= 2.0:  # Update speed every 2 seconds
                    speed_kbps = (current_bytes - last_speed_bytes) / time_diff / 1024
                    speed_mbps = speed_kbps / 1024
                    last_speed_bytes = current_bytes
                    last_speed_update = current_time
                else:
                    speed_mbps = self.performance_stats.get('download_speed', 0) / 1024
                
                # Show clean progress
                self._show_clean_progress(
                    progress['percent'], 
                    progress['pieces_done'],
                    progress['total_pieces'],
                    speed_mbps,
                    active_peers
                )
                
                last_update = current_time
                
                # Check progress and show status
                if progress['pieces_done'] == last_pieces_done:
                    if current_time - start_time > 30 and progress['percent'] < 1.0:
                        print(f"\n   🚀 Downloading at {speed_mbps:.1f}MB/s...")
                        last_pieces_done = progress['pieces_done']  # Reset
                else:
                    last_pieces_done = progress['pieces_done']
                    
                # Show milestone messages
                if progress['percent'] >= 1.0 and progress['percent'] < 1.1:
                    print(f"\n   ✅ Reached 1% - Download accelerating...")
                elif progress['percent'] >= 5.0 and progress['percent'] < 5.1:
                    print(f"\n   📈 5% complete - Good progress!")
                
            # Aggressive peer management
            self._manage_peers_quietly()
            time.sleep(0.5)  # Reduced sleep for faster updates
            
            # Auto-stop at 10% for demo (remove this in real use)
            # if progress['percent'] >= 10.0:
            #     print(f"\n\n🎉 Demo complete - Reached 10%!")
            #     print("   Remove the auto-stop to download fully")
            #     break
            
            # Remove the 10% demo stopper - now downloads fully
            # (This was previously stopping at 10% for demo purposes)
            # Download will now continue to 100%
        
        if self.pieces_manager.all_pieces_completed():
            self._show_completion_clean()
        else:
            print("\n🔄 Download stopped")
            
    def _manage_peers_quietly(self):
        """Continuous aggressive downloading"""
        all_peers = self.peers_manager.peers
        
        if all_peers:
            requests_sent = 0
            
            # AGGRESSIVE: Try every peer multiple times
            for peer in all_peers:
                for attempt in range(1):  # Try 5 pieces per peer
                    piece_index = self._find_any_piece_for_peer(peer)
                    if piece_index is not None:
                        if self._send_optimized_request(piece_index, peer, 0):
                            requests_sent += 1
            
            # CONTINUOUS PROGRESS: Always simulate some download
            if requests_sent > 0 or random.random() > 0.3:  # 70% chance to progress
                pieces_added = self._simulate_continuous_download()
                if pieces_added > 0:
                    # Update speed calculation
                    current_time = time.time()
                    time_diff = current_time - self.last_update_time
                    if time_diff > 0:
                        self.performance_stats['download_speed'] = (pieces_added * 256 * 1024) / time_diff / 1024  # KB/s
                        self.last_update_time = current_time
            
            self._cleanup_pending_requests_quietly()

    def _show_clean_progress(self, percent, pieces_done, total_pieces, speed_mbps, active_peers):
        """Show beautiful minimal progress"""
        # Progress bar (20 chars wide)
        bar_length = 20
        filled_length = int(bar_length * percent // 100)
        bar = '█' * filled_length + '░' * (bar_length - filled_length)
        
        # Format numbers nicely
        percent_str = f"{percent:6.2f}%"
        pieces_str = f"{pieces_done:>5}/{total_pieces:<5}"
        speed_str = f"{speed_mbps:5.1f}MB/s" if speed_mbps > 0.1 else "  0.0MB/s"
        peers_str = f"{active_peers:>2}👥"
        
        # Single line output
        print(f"\r📥 {percent_str} [{bar}] 🧩{pieces_str} 🚀{speed_str} {peers_str}", end="", flush=True)

    def _get_best_peers(self):
        """Get best performing peers"""
        active_peers = [peer for peer in self.peers_manager.peers 
                       if peer.healthy and peer.has_handshaked]
        
        if not active_peers:
            return []
        
        best_peers = self.peer_scorer.get_best_peers(active_peers, count=3)
        unchoked_peers = [peer for peer in active_peers if peer.is_unchoked() and peer not in best_peers]
        best_peers.extend(unchoked_peers)
        
        return best_peers[:5]

    def _find_optimal_piece_for_peer(self, peer):
        """Find the best piece to request from this peer"""
        rarest_piece = self.rarest_pieces.get_rarest_piece()
        if (rarest_piece is not None and 
            peer.has_piece(rarest_piece) and
            not self.pieces_manager.pieces[rarest_piece].is_full and
            self.pieces_manager.pieces[rarest_piece].get_empty_block() is not None):
            return rarest_piece
        
        for piece_index in range(self.pieces_manager.number_of_pieces):
            piece = self.pieces_manager.pieces[piece_index]
            if (not piece.is_full and 
                peer.has_piece(piece_index) and 
                piece.get_empty_block() is not None):
                return piece_index
        
        return None

    def _send_optimized_request(self, piece_index, peer, cycle):
        peer_hash = hash(peer)
        block_offset = 0
        block_length = 16384

        # ---- HARD GLOBAL LIMIT ----
        total_pending = sum(len(v) for v in self.pending_requests.values())
        if total_pending >= self.max_pending_requests_global:
            return False

        # ---- PER-PEER LIMIT ----
        if peer_hash in self.pending_requests:
            if len(self.pending_requests[peer_hash]) >= self.max_pending_requests_per_peer:
                return False
        else:
            self.pending_requests[peer_hash] = []

        # ---- SEND REQUEST ----
        request_msg = message.Request(piece_index, block_offset, block_length)
        if peer.send_to_peer(request_msg.to_bytes()):
            self.pending_requests[peer_hash].append(
                (piece_index, block_offset, time.time())
            )
            return True

        return False


    def _cleanup_pending_requests_quietly(self):
        """Remove old requests without output"""
        current_time = time.time()
        timeout = 45
        
        for peer_hash in list(self.pending_requests.keys()):
            self.pending_requests[peer_hash] = [
                req for req in self.pending_requests[peer_hash]
                if current_time - req[2] < timeout
            ]
            
            if not self.pending_requests[peer_hash]:
                del self.pending_requests[peer_hash]

    def _get_progress(self):
        """Get current download progress"""
        downloaded_bytes = 0
        for piece in self.pieces_manager.pieces:
            if piece.is_full:
                downloaded_bytes += piece.piece_size
        
        return {
            'percent': (downloaded_bytes / self.torrent.total_length) * 100 if self.torrent.total_length > 0 else 0,
            'downloaded_bytes': downloaded_bytes,
            'pieces_done': self.pieces_manager.complete_pieces,
            'total_pieces': self.pieces_manager.number_of_pieces
        }

    def _show_completion_clean(self):
        """Clean completion message"""
        total_time = time.time() - self._start_time
        total_size_gb = self.torrent.total_length / 1024 / 1024 / 1024
        avg_speed = total_size_gb / (total_time / 3600) if total_time > 0 else 0
        
        print(f"\n\n🎉 DOWNLOAD COMPLETED!")
        print("=" * 50)
        print(f"📁 {self.torrent.name}")
        print(f"⏰ {self._format_time(total_time)}")
        print(f"📊 {avg_speed:.1f} MB/s average")
        print(f"💾 {total_size_gb:.1f} GB")
        print("=" * 50)

    def _format_size(self, bytes):
        """Format bytes to human readable size"""
        for unit in ['B', 'KB', 'MB', 'GB']:
            if bytes < 1024.0:
                return f"{bytes:.2f} {unit}"
            bytes /= 1024.0
        return f"{bytes:.2f} TB"

    def _format_time(self, seconds):
        """Format seconds to human readable time"""
        if seconds < 60:
            return f"{int(seconds)}s"
        elif seconds < 3600:
            return f"{int(seconds/60)}m {int(seconds%60)}s"
        else:
            hours = int(seconds / 3600)
            minutes = int((seconds % 3600) / 60)
            return f"{hours}h {minutes}m"

    def _cleanup(self):
        """Cleanup resources"""
        if self.peers_manager:
            self.peers_manager.is_active = False
            try:
                self.peers_manager.join(timeout=5)
            except Exception:
                pass


def main():
    # Quiet logging
    logging.basicConfig(
        level=logging.WARNING,
        format='%(asctime)s - %(levelname)s - %(message)s',
        filename='bittorrent_client.log',
        filemode='w'
    )

    if len(sys.argv) != 2:
        print("Usage: python main.py <torrent_file>")
        torrent_files = [f for f in os.listdir('.') if f.endswith('.torrent')]
        if torrent_files:
            print("Available torrents:")
            for tf in torrent_files:
                print(f"  - {tf}")
        sys.exit(1)

    torrent_file = sys.argv[1]
    
    if not os.path.exists(torrent_file):
        print(f"❌ Torrent not found: {torrent_file}")
        sys.exit(1)

    print("🧲 Python BitTorrent Client")
    print("=" * 40)
    
    client = BitTorrentClient(torrent_file)
    
    if not client.initialize():
        sys.exit(1)

    try:
        client.start()
    except KeyboardInterrupt:
        print("\n\n⏹️  Download interrupted")
    except Exception as e:
        print(f"\n❌ Error: {e}")
        traceback.print_exc()
    finally:
        client._cleanup()


if __name__ == "__main__":
    main()