"""
Medical Imaging Dataset - GitHub Dataset Downloader
====================================================

Project: Medical Imaging Dataset
Website: https://rskworld.in
Contact: help@rskworld.in, support@rskworld.in
Phone: +91 93305 39277
Founder: Molla Samser
Designer & Tester: Rima Khatun

This script downloads medical imaging datasets directly from GitHub repositories.
"""

import os
import requests
import zipfile
import shutil
from pathlib import Path
from urllib.parse import urlparse


# GitHub repositories with medical imaging data
GITHUB_DATASETS = {
    "covid_chestxray": {
        "name": "COVID-19 Chest X-ray Dataset",
        "repo": "ieee8023/covid-chestxray-dataset",
        "description": "Chest X-rays for COVID-19 detection",
        "license": "MIT",
        "images": "~100+ images",
        "direct_download": True
    },
    "medical_imaging": {
        "name": "Medical Imaging Datasets",
        "repo": "v7labs/medical-imaging-datasets",
        "description": "Collection of medical imaging datasets",
        "license": "Various",
        "images": "Multiple",
        "direct_download": True
    }
}


def download_github_repo(repo_name, output_dir="data/temp"):
    """
    Download a GitHub repository as ZIP.
    
    Args:
        repo_name: Repository name in format "owner/repo"
        output_dir: Directory to save the downloaded files
    """
    url = f"https://github.com/{repo_name}/archive/refs/heads/master.zip"
    
    # Try main branch if master doesn't work
    alt_url = f"https://github.com/{repo_name}/archive/refs/heads/main.zip"
    
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    zip_path = output_path / f"{repo_name.replace('/', '_')}.zip"
    
    print(f"Downloading {repo_name}...")
    print(f"URL: {url}")
    
    try:
        response = requests.get(url, stream=True, timeout=30)
        if response.status_code == 404:
            print("Master branch not found, trying main branch...")
            response = requests.get(alt_url, stream=True, timeout=30)
            url = alt_url
        
        if response.status_code == 200:
            total_size = int(response.headers.get('content-length', 0))
            downloaded = 0
            
            with open(zip_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
                        downloaded += len(chunk)
                        if total_size > 0:
                            percent = (downloaded / total_size) * 100
                            print(f"\rProgress: {percent:.1f}%", end='', flush=True)
            
            print(f"\nDownloaded to: {zip_path}")
            
            # Extract ZIP file
            print("Extracting...")
            extract_dir = output_path / repo_name.split('/')[-1]
            with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                zip_ref.extractall(output_path)
            
            print(f"Extracted to: {extract_dir}")
            return str(extract_dir)
        else:
            print(f"Error: HTTP {response.status_code}")
            return None
            
    except Exception as e:
        print(f"Error downloading: {e}")
        return None


def download_covid_dataset():
    """Download COVID-19 chest X-ray dataset from GitHub."""
    print("=" * 60)
    print("Downloading COVID-19 Chest X-ray Dataset")
    print("=" * 60)
    print()
    
    repo = "ieee8023/covid-chestxray-dataset"
    extract_path = download_github_repo(repo, "data/temp/covid")
    
    if extract_path:
        # Organize images
        source = Path(extract_path)
        target_images = Path("data/xray/images")
        target_labels = Path("data/xray/labels")
        
        target_images.mkdir(parents=True, exist_ok=True)
        target_labels.mkdir(parents=True, exist_ok=True)
        
        # Find and copy images
        image_extensions = ['.png', '.jpg', '.jpeg']
        image_count = 0
        
        for ext in image_extensions:
            for img_file in source.rglob(f'*{ext}'):
                if img_file.is_file() and 'images' in str(img_file):
                    dest = target_images / img_file.name
                    if not dest.exists():
                        shutil.copy2(img_file, dest)
                        image_count += 1
        
        print(f"\nOrganized {image_count} images to data/xray/images/")
        return True
    
    return False


def download_medmnist_images():
    """Download sample images from MedMNIST (if available via direct link)."""
    print("=" * 60)
    print("MedMNIST Dataset Information")
    print("=" * 60)
    print()
    print("MedMNIST is available via:")
    print("1. Kaggle: https://www.kaggle.com/datasets/andrewmvd/medical-mnist")
    print("2. Official: https://medmnist.com/")
    print()
    print("For direct download, please visit the official website or use Kaggle API.")
    print()


def list_public_sources():
    """List additional public sources for medical imaging data."""
    print("=" * 60)
    print("Additional Public Medical Imaging Data Sources")
    print("=" * 60)
    print()
    
    sources = [
        {
            "name": "The Cancer Imaging Archive (TCIA)",
            "url": "https://www.cancerimagingarchive.net/",
            "description": "Open-access database of medical images for cancer research",
            "modalities": "CT, MRI, PET",
            "access": "Free registration required"
        },
        {
            "name": "Radiopaedia",
            "url": "https://radiopaedia.org/",
            "description": "Educational radiology resource with case studies",
            "modalities": "X-ray, CT, MRI, Ultrasound",
            "access": "Free, educational use"
        },
        {
            "name": "OpenNeuro",
            "url": "https://openneuro.org/",
            "description": "Open-science database for brain imaging",
            "modalities": "MRI, fMRI, EEG",
            "access": "Free, open access"
        },
        {
            "name": "Stanford AIMI",
            "url": "https://aimi.stanford.edu/resources/shared-datasets",
            "description": "Stanford AI in Medicine & Imaging shared datasets",
            "modalities": "Multiple",
            "access": "Free for research"
        },
        {
            "name": "MIDRC",
            "url": "https://www.midrc.org/",
            "description": "Medical Imaging Data Resource Center",
            "modalities": "CT, X-ray, MRI",
            "access": "Free registration"
        },
        {
            "name": "MedMNIST",
            "url": "https://medmnist.com/",
            "description": "Large-scale standardized biomedical images",
            "modalities": "Multiple (12 2D + 6 3D datasets)",
            "access": "Free, direct download"
        }
    ]
    
    for source in sources:
        print(f"Source: {source['name']}")
        print(f"  URL: {source['url']}")
        print(f"  Description: {source['description']}")
        print(f"  Modalities: {source['modalities']}")
        print(f"  Access: {source['access']}")
        print()


def main():
    """Main function."""
    print("=" * 60)
    print("Medical Imaging Dataset - GitHub Downloader")
    print("=" * 60)
    print()
    print("This script helps download publicly available medical imaging")
    print("datasets from GitHub and other public sources.")
    print()
    print("Options:")
    print("1. Download COVID-19 Chest X-ray Dataset (GitHub)")
    print("2. List additional public data sources")
    print("3. Show MedMNIST information")
    print()
    
    choice = input("Enter option (1-3) or 'q' to quit: ").strip()
    
    if choice == '1':
        download_covid_dataset()
    elif choice == '2':
        list_public_sources()
    elif choice == '3':
        download_medmnist_images()
    elif choice.lower() == 'q':
        print("Exiting...")
    else:
        print("Invalid choice.")


if __name__ == '__main__':
    main()