"""Download the Brain Tumor MRI dataset from Kaggle.

Dataset: masoudnickparvar/brain-tumor-mri-dataset
Layout after extraction:
    data/raw/Training/{glioma,meningioma,notumor,pituitary}/*.jpg
    data/raw/Testing/{glioma,meningioma,notumor,pituitary}/*.jpg

Prerequisites:
    1. Create a Kaggle account and an API token at https://www.kaggle.com/settings.
    2. Save the downloaded kaggle.json to ~/.kaggle/kaggle.json and chmod 600.
"""

from __future__ import annotations

import os
import subprocess
import sys
from pathlib import Path

DATASET = "masoudnickparvar/brain-tumor-mri-dataset"
PROJECT_ROOT = Path(__file__).resolve().parent.parent
RAW_DIR = PROJECT_ROOT / "data" / "raw"
EXPECTED_CLASSES = {"glioma", "meningioma", "notumor", "pituitary"}


def ensure_kaggle_cli() -> None:
    try:
        subprocess.run(
            ["kaggle", "--version"], check=True, capture_output=True, text=True
        )
    except (subprocess.CalledProcessError, FileNotFoundError):
        print("Kaggle CLI not found. Installing...")
        subprocess.run([sys.executable, "-m", "pip", "install", "kaggle"], check=True)


def ensure_credentials() -> None:
    cred = Path.home() / ".kaggle" / "kaggle.json"
    if not cred.exists():
        sys.exit(
            "\nMissing ~/.kaggle/kaggle.json.\n"
            "Create one at https://www.kaggle.com/settings -> 'Create New Token'\n"
            "then: mkdir -p ~/.kaggle && mv ~/Downloads/kaggle.json ~/.kaggle/ && chmod 600 ~/.kaggle/kaggle.json\n"
        )
    # Kaggle CLI refuses to run if the file is world-readable.
    cred.chmod(0o600)


def already_downloaded() -> bool:
    training = RAW_DIR / "Training"
    testing = RAW_DIR / "Testing"
    if not training.is_dir() or not testing.is_dir():
        return False
    train_classes = {p.name for p in training.iterdir() if p.is_dir()}
    test_classes = {p.name for p in testing.iterdir() if p.is_dir()}
    return EXPECTED_CLASSES.issubset(train_classes) and EXPECTED_CLASSES.issubset(
        test_classes
    )


def download() -> None:
    RAW_DIR.mkdir(parents=True, exist_ok=True)
    print(f"Downloading {DATASET} -> {RAW_DIR}")
    subprocess.run(
        [
            "kaggle",
            "datasets",
            "download",
            "-d",
            DATASET,
            "-p",
            str(RAW_DIR),
            "--unzip",
        ],
        check=True,
    )


def summarize() -> None:
    for split in ("Training", "Testing"):
        split_dir = RAW_DIR / split
        if not split_dir.is_dir():
            continue
        print(f"\n{split}:")
        for cls_dir in sorted(split_dir.iterdir()):
            if cls_dir.is_dir():
                n = sum(1 for _ in cls_dir.iterdir())
                print(f"  {cls_dir.name:12s} {n:5d} images")


def main() -> None:
    if already_downloaded():
        print(f"Dataset already present at {RAW_DIR}. Skipping download.")
    else:
        ensure_kaggle_cli()
        ensure_credentials()
        download()
    summarize()
    print("\nDone.")


if __name__ == "__main__":
    main()