import os
import zipfile
import gdown

# Google Drive file ID for SPOD_30b_8c.zip
FILE_ID = '1wfoLkfZOxxEtuDyDSWCnAq-N8HZVKdQ1'
DEST_DIR = './data/SPOD_dataset'
ZIP_PATH = os.path.join(DEST_DIR, 'SPOD_30b_8c.zip')
EXTRACTED_PATH = os.path.join(DEST_DIR, 'SPOD_30b_8c')
CONFIG_PATH = './configs/_base_/datasets/hsi_detection.py'

os.makedirs(DEST_DIR, exist_ok=True)

# Download the dataset
if not os.path.exists(ZIP_PATH):
    url = f'[https://drive.google.com/uc?id={FILE_ID}'](https://drive.google.com/uc?id={FILE_ID}')
    print(f'Downloading dataset from {url}...')
    gdown.download(url, ZIP_PATH, quiet=False)
else:
    print('Dataset zip already exists.')

# Extract the dataset
if not os.path.exists(EXTRACTED_PATH):
    print('Extracting dataset...')
    with zipfile.ZipFile(ZIP_PATH, 'r') as zip_ref:
        zip_ref.extractall(DEST_DIR)
else:
    print('Dataset already extracted.')

# Update config file
if os.path.exists(CONFIG_PATH):
    with open(CONFIG_PATH, 'r') as f:
        lines = f.readlines()
    with open(CONFIG_PATH, 'w') as f:
        for line in lines:
            if line.strip().startswith('data_root'):
                f.write(f"data_root = '{EXTRACTED_PATH}/'\n")
            else:
                f.write(line)
    print(f'Updated {CONFIG_PATH} with dataset path.')
else:
    print(f'Config file {CONFIG_PATH} not found!')

print('Done.')