#!/usr/bin/env python3
"""
Выкачивание всех данных с FTP сервера
"""
import ftplib
import socket
import os
from pathlib import Path
from datetime import datetime

def download_ftp_data(host, port, username, password, output_dir):
    """Выкачать все данные с FTP сервера"""
    try:
        # Подключаемся к FTP серверу
        ftp = ftplib.FTP()
        ftp.connect(host, port, timeout=30)
        ftp.login(username, password)
        
        # Получаем приветственное сообщение
        welcome = ftp.getwelcome()
        print(f"Connected: {welcome}")
        
        # Создаем директорию для пользователя
        user_dir = output_dir / username
        user_dir.mkdir(exist_ok=True)
        
        # Счетчики
        files_downloaded = 0
        dirs_created = 0
        errors = []
        
        def download_recursive(remote_path, local_path):
            nonlocal files_downloaded, dirs_created, errors
            
            try:
                # Получаем листинг текущей директории
                current_listing = []
                ftp.cwd(remote_path)
                ftp.retrlines('LIST', current_listing.append)
                
                for line in current_listing:
                    parts = line.split()
                    if len(parts) < 9:
                        continue
                    
                    item_name = ' '.join(parts[8:])
                    if item_name in ['.', '..']:
                        continue
                    
                    remote_item_path = f"{remote_path}/{item_name}" if remote_path != '/' else f"/{item_name}"
                    local_item_path = local_path / item_name
                    
                    if parts[0].startswith('d'):  # Директория
                        print(f"  Creating directory: {remote_item_path}")
                        local_item_path.mkdir(exist_ok=True)
                        dirs_created += 1
                        
                        # Рекурсивно скачиваем содержимое директории
                        download_recursive(remote_item_path, local_item_path)
                        
                    else:  # Файл
                        print(f"  Downloading file: {remote_item_path}")
                        try:
                            with open(local_item_path, 'wb') as f:
                                ftp.retrbinary(f'RETR {remote_item_path}', f.write)
                            files_downloaded += 1
                        except Exception as e:
                            error_msg = f"Failed to download {remote_item_path}: {e}"
                            errors.append(error_msg)
                            print(f"    ERROR: {error_msg}")
                
            except Exception as e:
                error_msg = f"Failed to process {remote_path}: {e}"
                errors.append(error_msg)
                print(f"  ERROR: {error_msg}")
        
        # Начинаем скачивание с корневой директории
        print(f"Starting download from root directory...")
        download_recursive('/', user_dir)
        
        # Сохраняем статистику
        stats_file = user_dir / 'download_stats.txt'
        with open(stats_file, 'w', encoding='utf-8') as f:
            f.write(f"=== FTP Download Statistics ===\n")
            f.write(f"Date: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write(f"Host: {host}\n")
            f.write(f"Username: {username}\n")
            f.write(f"Files downloaded: {files_downloaded}\n")
            f.write(f"Directories created: {dirs_created}\n")
            f.write(f"Errors: {len(errors)}\n")
            if errors:
                f.write(f"\n=== Errors ===\n")
                for error in errors:
                    f.write(f"  - {error}\n")
        
        ftp.quit()
        
        return {
            'username': username,
            'files_downloaded': files_downloaded,
            'dirs_created': dirs_created,
            'errors': len(errors),
            'user_dir': user_dir
        }
        
    except Exception as e:
        print(f"Error downloading data for {username}: {e}")
        return None

def main():
    print("=== Выкачивание всех данных с FTP ===\n")
    
    # FTP credentials
    ftp_creds = [
        {
            'db': '2210PIGOTTOILUAP01',
            'username': 'marketmax',
            'password': 'IMP$1220mm',
            'host': 'ftp.imperialtrading.com',
            'type': 'ftp_vendor'
        },
        {
            'db': '2610PIGOTTOILUAP01',
            'username': 'marketmaxinv',
            'password': 'IMP$0606ma',
            'host': 'ftp.imperialtrading.com',
            'type': 'ftp_vendor'
        }
    ]
    
    output_dir = Path('/root/ir-assessment/redteam/irelydata/ftp_downloads')
    output_dir.mkdir(exist_ok=True)
    
    total_files = 0
    total_dirs = 0
    total_errors = 0
    
    for cred in ftp_creds:
        print(f"\n=== Downloading data for {cred['username']} ===")
        
        result = download_ftp_data(
            cred['host'], 
            21, 
            cred['username'], 
            cred['password'],
            output_dir
        )
        
        if result:
            total_files += result['files_downloaded']
            total_dirs += result['dirs_created']
            total_errors += result['errors']
            
            print(f"  ✓ Downloaded {result['files_downloaded']} files, {result['dirs_created']} directories")
            if result['errors'] > 0:
                print(f"  ⚠ {result['errors']} errors occurred")
        else:
            print(f"  ✗ Failed to download data")
    
    # Сохраняем общую статистику
    summary_file = output_dir / 'download_summary.txt'
    with open(summary_file, 'w', encoding='utf-8') as f:
        f.write(f"=== FTP Download Summary ===\n")
        f.write(f"Date: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"Total users processed: {len(ftp_creds)}\n")
        f.write(f"Total files downloaded: {total_files}\n")
        f.write(f"Total directories created: {total_dirs}\n")
        f.write(f"Total errors: {total_errors}\n")
        f.write(f"Output directory: {output_dir}\n")
    
    print(f"\n=== Download Complete ===")
    print(f"Total files downloaded: {total_files}")
    print(f"Total directories created: {total_dirs}")
    print(f"Total errors: {total_errors}")
    print(f"Output directory: {output_dir}")

if __name__ == '__main__':
    main()