"""
Medical Imaging Dataset - Automatic 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 automatically downloads publicly available medical imaging datasets.
"""

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


def download_file(url, output_path, description=""):
    """Download a file from URL."""
    print(f"Downloading {description}...")
    print(f"URL: {url}")
    
    try:
        response = requests.get(url, stream=True, timeout=60)
        if response.status_code == 200:
            total_size = int(response.headers.get('content-length', 0))
            downloaded = 0
            
            output_path.parent.mkdir(parents=True, exist_ok=True)
            
            with open(output_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: {output_path}")
            return True
        else:
            print(f"Error: HTTP {response.status_code}")
            return False
    except Exception as e:
        print(f"Error: {e}")
        return False


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()
    
    # Try downloading the repository as ZIP
    url = "https://github.com/ieee8023/covid-chestxray-dataset/archive/refs/heads/master.zip"
    alt_url = "https://github.com/ieee8023/covid-chestxray-dataset/archive/refs/heads/main.zip"
    
    zip_path = Path("data/temp/covid_dataset.zip")
    
    # Try master branch first
    if not download_file(url, zip_path, "COVID-19 Dataset"):
        # Try main branch
        print("Trying main branch...")
        if not download_file(alt_url, zip_path, "COVID-19 Dataset"):
            print("Failed to download COVID-19 dataset")
            return False
    
    # Extract
    print("Extracting...")
    extract_dir = Path("data/temp/covid")
    extract_dir.mkdir(parents=True, exist_ok=True)
    
    try:
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(extract_dir)
        print(f"Extracted to: {extract_dir}")
    except Exception as e:
        print(f"Error extracting: {e}")
        return False
    
    # Organize images
    print("Organizing images...")
    source = extract_dir
    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 all image files
    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():
                # Skip if already exists
                dest = target_images / img_file.name
                if not dest.exists():
                    try:
                        shutil.copy2(img_file, dest)
                        image_count += 1
                        
                        # Create label file
                        label_file = target_labels / f"{img_file.stem}.json"
                        if not label_file.exists():
                            label_data = {
                                "diagnosis": "COVID-19" if "covid" in img_file.name.lower() else "Normal",
                                "confidence": 0.90,
                                "date": "2024-01-15",
                                "modality": "X-ray",
                                "body_part": "Chest",
                                "source": "COVID-19 Chest X-ray Dataset (GitHub)",
                                "license": "MIT",
                                "annotations": {
                                    "findings": ["From public COVID-19 dataset"],
                                    "recommendations": ["Educational use only"]
                                },
                                "metadata": {
                                    "patient_id": "ANONYMIZED",
                                    "age": "N/A",
                                    "gender": "N/A",
                                    "technique": "Digital radiography"
                                }
                            }
                            with open(label_file, 'w') as f:
                                json.dump(label_data, f, indent=4)
                    except Exception as e:
                        print(f"Error copying {img_file}: {e}")
    
    print(f"Organized {image_count} images to data/xray/images/")
    return image_count > 0


def download_sample_images_from_medmnist():
    """Download sample images information (MedMNIST requires direct access)."""
    print("=" * 60)
    print("MedMNIST Dataset Information")
    print("=" * 60)
    print()
    print("MedMNIST is available at: https://medmnist.com/")
    print("You can download it directly from the website.")
    print("Or use Kaggle: kaggle datasets download -d andrewmvd/medical-mnist")
    print()


def create_download_summary():
    """Create a summary of downloaded datasets."""
    summary = {
        "downloaded_datasets": [],
        "total_images": 0,
        "xray_images": 0,
        "ct_images": 0,
        "mri_images": 0
    }
    
    # Count images
    xray_dir = Path("data/xray/images")
    ct_dir = Path("data/ct_scan/images")
    mri_dir = Path("data/mri/images")
    
    if xray_dir.exists():
        xray_count = len(list(xray_dir.glob("*.png")) + list(xray_dir.glob("*.jpg")) + list(xray_dir.glob("*.jpeg")))
        summary["xray_images"] = xray_count
        summary["total_images"] += xray_count
        if xray_count > 0:
            summary["downloaded_datasets"].append("X-ray images")
    
    if ct_dir.exists():
        ct_count = len(list(ct_dir.glob("*.png")) + list(ct_dir.glob("*.jpg")) + list(ct_dir.glob("*.jpeg")) + list(ct_dir.glob("*.dcm")))
        summary["ct_images"] = ct_count
        summary["total_images"] += ct_count
        if ct_count > 0:
            summary["downloaded_datasets"].append("CT scan images")
    
    if mri_dir.exists():
        mri_count = len(list(mri_dir.glob("*.png")) + list(mri_dir.glob("*.jpg")) + list(mri_dir.glob("*.jpeg")) + list(mri_dir.glob("*.dcm")))
        summary["mri_images"] = mri_count
        summary["total_images"] += mri_count
        if mri_count > 0:
            summary["downloaded_datasets"].append("MRI images")
    
    # Save summary
    with open("download_summary.json", 'w') as f:
        json.dump(summary, f, indent=4)
    
    return summary


def main():
    """Main function to download all available datasets."""
    print("=" * 60)
    print("Medical Imaging Dataset - Automatic Downloader")
    print("=" * 60)
    print()
    print("This script will download publicly available medical imaging datasets.")
    print()
    
    downloaded = False
    
    # Download COVID-19 dataset (easiest, no account needed)
    print("\n[1/1] Downloading COVID-19 Chest X-ray Dataset...")
    if download_covid_dataset():
        downloaded = True
        print("[OK] COVID-19 dataset downloaded successfully!")
    else:
        print("[WARNING] Could not download COVID-19 dataset")
    
    print()
    print("=" * 60)
    print("Download Summary")
    print("=" * 60)
    
    summary = create_download_summary()
    
    print(f"Total images: {summary['total_images']}")
    print(f"  - X-ray images: {summary['xray_images']}")
    print(f"  - CT scan images: {summary['ct_images']}")
    print(f"  - MRI images: {summary['mri_images']}")
    print()
    
    if summary['total_images'] > 0:
        print("Downloaded datasets:")
        for dataset in summary['downloaded_datasets']:
            print(f"  - {dataset}")
    else:
        print("No images downloaded yet.")
        print()
        print("To download more datasets:")
        print("1. Read DOWNLOAD_INSTRUCTIONS.md for detailed instructions")
        print("2. Use Kaggle API for larger datasets (requires free account)")
        print("3. Visit https://medmnist.com/ for MedMNIST dataset")
    
    print()
    print("=" * 60)
    print("Note: For more datasets, see DOWNLOAD_INSTRUCTIONS.md")
    print("=" * 60)


if __name__ == '__main__':
    main()