import urllib.request
import os
from pathlib import Path
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

YOLO_CONFIG = {
    'weights': {
        'url': 'https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.weights',
        'filename': 'yolov4.weights'
    },
    'config': {
        'url': 'https://raw.githubusercontent.com/AlexeyAB/darknet/master/cfg/yolov4.cfg',
        'filename': 'yolov4.cfg'
    }
}

def download_file(url: str, filename: str):
    """Download a file and show progress."""
    try:
        logger.info(f"Downloading {filename}...")
        
        def report_progress(count, block_size, total_size):
            percent = int(count * block_size * 100 / total_size)
            print(f"\rProgress: {percent}%", end='')
            
        filepath = Path('app/models') / filename
        filepath.parent.mkdir(parents=True, exist_ok=True)
        
        urllib.request.urlretrieve(
            url,
            filepath,
            reporthook=report_progress
        )
        print()  # New line after progress
        logger.info(f"Successfully downloaded {filename}")
        return True
    except Exception as e:
        logger.error(f"Error downloading {filename}: {e}")
        return False

def setup_yolo():
    """Download and set up YOLO files."""
    success = True
    for file_info in YOLO_CONFIG.values():
        if not download_file(file_info['url'], file_info['filename']):
            success = False
    
    if success:
        logger.info("YOLO setup completed successfully!")
    else:
        logger.error("YOLO setup failed. Please check the errors above.")

if __name__ == "__main__":
    setup_yolo()