mors.data

def DownloadEvolutionTracks(fname=''):
    """
    Download evolution track data

    Inputs :
        - fname (optional) :    folder name, "Spada" or "Baraffe"
                                if not provided download both

    Baraffe is fetched through fwl-io into the versioned data layout; Spada
    stays on the legacy Zenodo/OSF path until fwl-io gains archive extraction.
    """

    # If no folder name specified download both Spada and Baraffe
    if not fname:
        folder_list = ('Spada', 'Baraffe')
    elif fname in ('Spada', 'Baraffe'):
        folder_list = [fname]
    else:
        raise ValueError(f'Unrecognised folder name: {fname}')

    if 'Baraffe' in folder_list:
        _fetch_baraffe()
    if 'Spada' in folder_list:
        _download_spada()

    return

def GetFWLData() -> Path:
    """
    Get path to FWL data directory on the disk
    """
    return Path(FWL_DATA_DIR).absolute()

def _baraffe_dataset():
    """Return the Baraffe dataset declared in the shipped manifest."""
    from fwl_io import load_manifest

    datasets = {ds.key: ds for ds in load_manifest(manifest_path())}
    return datasets[_BARAFFE_KEY]

def _baraffe_fetcher():
    """Build an fwl-io fetcher for the Baraffe tracks from the manifest pin."""
    from fwl_io import create_fetcher

    ds = _baraffe_dataset()
    return create_fetcher(
        subdir=ds.subdir,
        zenodo=ds.zenodo,
        registry=ds.registry(),
        data_root=GetFWLData(),
    )

def _download_spada():
    """Download and unpack the Spada grid via Zenodo, falling back to OSF."""
    # Create stellar evolution tracks data repository if not existing
    data_dir = GetFWLData() / 'stellar_evolution_tracks'
    data_dir.mkdir(parents=True, exist_ok=True)

    folder = 'Spada'
    folder_dir = data_dir / folder
    if folder_dir.exists():
        return

    # Link with OSF project repository (fallback mirror)
    osf = OSF()
    project = osf.project(project_id)
    storage = project.storage('osfstorage')

    max_tries = 2  # Maximum download attempts, could be a function argument
    log.info(f'Downloading stellar evolution tracks to {data_dir}')
    for i in range(max_tries):
        log.info(f'Attempt {i + 1} of {max_tries}')
        success = False

        try:
            download_zenodo_folder(folder=folder, data_dir=data_dir)
            success = True
        except (subprocess.CalledProcessError, OSError) as e:
            # zenodo_get exits non-zero on failure (CalledProcessError via
            # check=True) or is missing (FileNotFoundError); neither is a
            # RuntimeError, so both must be caught for the fallback to run.
            log.error(f'Zenodo download failed: {e}')
            # A non-zero exit can leave partial files, so clear the whole tree;
            # rmdir would raise on a non-empty directory and mask the failure.
            shutil.rmtree(folder_dir, ignore_errors=True)

        if not success:
            try:
                download_OSF_folder(storage=storage, folders=folder, data_dir=data_dir)
                success = True
            except (RuntimeError, OSError) as e:
                log.error(f'OSF download failed: {e}')

        if success:
            break

        if i < max_tries - 1:
            log.info('Retrying download...')
            sleep(5)  # Wait 5 seconds before retrying
        else:
            log.error('Max retries reached. Download failed.')

    # Unzip Spada evolution tracks (only when the download populated the folder)
    if folder_dir.exists():
        wrk_dir = os.getcwd()
        os.chdir(os.path.join(data_dir, 'Spada'))
        subprocess.call(['tar', 'xvfz', 'fs255_grid.tar.gz'])
        subprocess.call(['rm', '-f', 'fs255_grid.tar.gz'])
        os.chdir(wrk_dir)

    return

def _fetch_baraffe():
    """Fetch the Baraffe tracks through fwl-io (idempotent, hash-verified)."""
    _baraffe_fetcher().fetch_all()

def baraffe_data_dir() -> Path:
    """Return the versioned directory that holds the Baraffe track files.

    The ``r<record-id>`` version segment is resolved by fwl-io from the pinned
    manifest, so the layout stays a single source of truth. Resolving the path
    creates the ``FWL_DATA`` root if absent (an fwl-io side effect) but does not
    download the tracks; call ``DownloadEvolutionTracks("Baraffe")`` (or
    ``mors download baraffe``) to populate it.
    """
    fetcher = _baraffe_fetcher()
    # A pre-versioning fwl-io (before the r<record-id> layout) resolves the bare
    # subdir and has no version_dir, which would silently mislocate the tracks.
    if getattr(fetcher, 'version_dir', None) is None:
        raise RuntimeError(
            f'fwl-io resolved an unversioned Baraffe directory {fetcher.target_dir}; '
            'upgrade to fwl-io>=26.7.20 for the versioned data layout.'
        )
    return fetcher.target_dir

def download_OSF_folder(*, storage, folders: list[str], data_dir: Path):
    """
    Download a specific folder in the OSF repository

    Inputs :
        - storage : OSF storage name
        - folders : folder names to download
        - data_dir : local repository where data are saved
    """
    for file in storage.files:
        for folder in folders:
            if not file.path[1:].startswith(folder):
                continue
            parts = file.path.split('/')[1:]
            target = Path(data_dir, *parts)
            target.parent.mkdir(parents=True, exist_ok=True)
            log.info(f'Downloading {file.path}...')
            with open(target, 'wb') as f:
                file.write_to(f)
            break

def download_zenodo_folder(folder: str, data_dir: Path):
    """
    Download a specific Zenodo record into specified folder

    Inputs :
        - folder : str
            Folder name to download
        - folder_dir : Path
            local repository where data are saved
    """

    folder_dir = data_dir / folder
    folder_dir.mkdir(parents=True)
    zenodo_id = get_zenodo_record(folder)
    cmd = ['zenodo_get', zenodo_id, '-o', folder_dir]
    out = os.path.join(GetFWLData(), 'zenodo.log')
    log.debug('    logging to %s' % out)
    with open(out, 'w') as hdl:
        subprocess.run(cmd, check=True, stdout=hdl, stderr=hdl)

def get_zenodo_record(folder: str) -> str | None:
    """
    Get Zenodo record ID for a given folder.

    Inputs :
        - folder : str
            Folder name to get the Zenodo record ID for

    Returns :
        - str | None : Zenodo record ID or None if not found
    """
    # Baraffe is fetched through fwl-io and is intentionally absent here.
    zenodo_map = {
        'Spada': '15729101',
    }
    return zenodo_map.get(folder, None)

def manifest_path() -> Path:
    """Entry-point target: the MORS dataset manifest read by fwl-io."""
    return Path(__file__).parent / 'data' / 'mors_manifest.toml'