main.py

from torrent import Torrent
from peer import Peer
from TrackersManager import TrackersManager
from PeerManager import PeerManager
from BlockandPiece import BLOCK_SIZE
import threading
import time
from datetime import datetime
import os
import yappi
import traceback
from logger import Logger
from collections import deque

NEG_INF = float('-inf')

def log_error(msg, exc=None, flags = None, name = ''):
    if flags is None:
        flags = []
    flags.insert(0, 'ERROR')
    flags.insert(0, f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    log_entry = f'[{timestamp}] {msg}\n'
    if exc is not None:
        if isinstance(exc, ConnectionResetError) or ('10054' in str(exc)):
            flags.append('Network')
            log_entry += f'Пир разорвал соединение (обычно для BitTorrent): {msg} — {exc}'
        else:
            log_entry = f'{msg}: {exc}'
    else:
        log_entry = f'{msg}'
    log_entry += '\n'
    res = ''
    for flag in flags:
        res += f"[{flag}]"
    res += ' '    
    res += log_entry    
    if name:
        path = os.path.join('logs', 'main')
    else:
        name = 'main.log'
        path = 'logs'
        
    with open(os.path.join(f'{path}', f"{name}"), 'a', encoding='utf-8') as f:
        f.write(res)

def log_info(msg, flags = None, name = ''):
    if flags is None:
        flags = []
    flags.insert(0, f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    log_entry = f'{msg}\n'
    res = ''
    for flag in flags:
        res += f"[{flag}]"
    res += ' '    
    res += log_entry    
    if name:
        path = os.path.join('logs', 'main')
    else:
        name = 'main.log'
        path = 'logs'
        
    with open(os.path.join(f'{path}', f"{name}"), 'a', encoding='utf-8') as f:
        f.write(res)

class Bittorrent:
    def __init__(self) -> None:
        self.max_concurrent_blocks = 150000
        self.semaphore = threading.Semaphore(self.max_concurrent_blocks)
        self._progress_thread = None
        self._monitor_thread = None
        self.running = True
        self.rarest_pieces_updated = False

    def _clear_logs_directory(self):
        """Очищает содержимое папки logs при запуске"""
        import shutil
        logs_dir = "logs"
        if os.path.exists(logs_dir):
            try:
                for filename in os.listdir(logs_dir):
                    file_path = os.path.join(logs_dir, filename)
                    if os.path.isfile(file_path):
                        os.remove(file_path)
                        Logger.info(f"Удален лог файл: {filename}")
                    elif os.path.isdir(file_path):
                        shutil.rmtree(file_path)
                        Logger.info(f"Удалена папка с логами: {filename}")
                Logger.info("Папка logs очищена")
            except Exception as e:
                Logger.error(f"Ошибка при очистке папки logs: {e}")
        else:
            Logger.info("Папка logs не найдена, создание не требуется")
        os.makedirs('logs', exist_ok=True)
        os.makedirs(os.path.join('logs', "peermanager"), exist_ok=True)
    
    def start_downloading(self, torrent_path: str, download_path: str):
        self._clear_logs_directory()
        torrent_obj = Torrent(torrent_path, download_path) 
        self.tracker = TrackersManager(torrent_obj)
        self.peer_manager = PeerManager(self.tracker, self.set_updated)
        self.tracker.downloaded_on_init = self.peer_manager.downloaded_bytes
        Logger.info(f"Total torrent length: {self.tracker.torrent_obj.total_length}")

        try:
            self._initialize()
            try:
                self._download_loop()
            except Exception as e:
                log_error(f"self._download_loop(): {e} \nTraceback:\n{traceback.format_exc()}")
        finally:
            self._finalize()

    def _initialize(self):
        self.tracker.start_periodic_updates()

    def _download_loop(self):
        
        while not self.peer_manager.seeding:
            try:
                print("main_loop")
                self._download_rarest_first()
                self._update_stats()
            except Exception as e:
                log_error(f"Error in download loop: {e}")
                time.sleep(1)
        print("seeding now")
        Logger.info("Torrent Complete")
        Logger.info(f"Total length: {self.tracker.torrent_obj.total_length}")
        while True:
            self._update_stats()
            time.sleep(5)        

    def set_updated(self):
        self.rarest_pieces_updated = True

    def _download_rarest_first(self):
        try:
            rarest_pieces: deque[deque[tuple[int, list[Peer]]]] = self.peer_manager.get_rarest_piece_rarest_pieces_copy()
            self.rarest_pieces_updated = False
            
            
            sent_blocks = 0
            sent = False
            if rarest_pieces:
                sent = True
                while rarest_pieces and not self.rarest_pieces_updated:
                    sent_this_iter = False
                    for i, dq in enumerate(rarest_pieces):
                        if sent_this_iter or self.rarest_pieces_updated:
                            break
                        while dq and not self.rarest_pieces_updated:
                            piece_i, peers = dq[0]
                            best_peer = max(peers, key=lambda p: p.peer_score(), default=None)
                            if best_peer and best_peer.peer_score() > NEG_INF:
                                sent_blocks_cur = self.peer_manager.prefetch_next_blocks(piece_i, best_peer)
                                sent_blocks += sent_blocks_cur
                                dq.popleft()
                                if dq and i != 0 and sent_blocks_cur:
                                    sent_this_iter = True
                                    break
                            else:
                                break
                            
                    
                    for _ in range(len(rarest_pieces)):
                        dq = rarest_pieces.popleft()
                        if dq:
                            rarest_pieces.append(dq)
                    
            if not sent or sent_blocks == 0:
                time.sleep(0.1)
        except Exception as e:
            log_error(f"Error in _download_rarest_first: {e}")

    def _update_stats(self):
        try:
            downloaded = self.peer_manager.downloaded_bytes
            uploaded = self.peer_manager.uploaded_bytes
            left = max(0, self.tracker.torrent_obj.total_length - downloaded)

            self.tracker.update_stats(downloaded, uploaded, left)
        except Exception as e:
            log_error(f"Error updating stats: {e}")

    def _finalize(self):
        self.running = False
        self.tracker.stop_periodic_updates()
        self.peer_manager.exitPeerThreads()
            
def main():
    # if len(sys.argv) != 3:
        # print("Usage: python main.py <torrent_file> <download_path>")
        # sys.exit(1)
    # torrent = sys.argv[1]
    # path = sys.argv[2]

    # torrent = r'.\torrents\music.torrent'
    # path = r"./downloads"
    
    # torrent = os.path.join('torrents', 'The_Jackbox_Party_Pack_3_MANY_PEERS_680MB.torrent')
    # torrent = os.path.join('torrents', 'REPO_300.torrent')
    # torrent = os.path.join('torrents', 'music.torrent')
    # torrent = os.path.join('torrents', 'manyLeeches5.torrent')
    torrent_path = os.path.join('torrents', 'Andr.torrent')
    # torrent = os.path.join('torrents', 'Madison.torrent')
    # torrent = os.path.join('torrents', 'ninja.torrent')
    # torrent = os.path.join('torrents', 'FoxLake.torrent')
    # torrent = os.path.join('torrents', '245_rut.torrent')
    # torrent = os.path.join('torrents', 'Photoshop_4gb.torrent')
    # torrent = os.path.join('torrents', 'novichok.torrent')
    # torrent = os.path.join('torrents', 'People_Playground.torrent')
    # torrent = os.path.join('torrents', '1PieceManyManyFiles.torrent')
    # torrent = os.path.join('torrents', 'Photoshop_2.58gb_rutr.torrent')
    download_path = os.path.join('.', 'downloads')
    b = Bittorrent()
    b.start_downloading(torrent_path, download_path)

if __name__ == "__main__":
    yappi.set_clock_type("cpu")
    yappi.start()
    try:
        main()
    finally:
        yappi.stop()
        stats = yappi.get_func_stats()
        stats.save("profile/profile.callgrind", type="CALLGRIND")
        stats.save("profile/profile.pstat", type="pstat")


        # yappi.get_func_stats().save("profile/yappi.prof", type="pstat")