import os
from datetime import datetime, timedelta

from biolib.biolib_logging import logger_no_user_data
from biolib.compute_node.job_worker.cache_state import LfsCacheState


def prune_lfs_cache(dry_run: bool) -> None:
    logger_no_user_data.info(f'Pruning LFS cache (dry run = {dry_run})...')

    current_time = datetime.utcnow()
    paths_to_delete = set()

    with LfsCacheState() as state:
        for _, partition in state['storage_partitions'].items():
            lfs_storage_path = f"{partition['path']}/lfs"
            if os.path.exists(lfs_storage_path):
                for lfs_uuid in os.listdir(lfs_storage_path):
                    if lfs_uuid not in state['large_file_systems']:
                        paths_to_delete.add(f'{lfs_storage_path}/{lfs_uuid}')

        lfs_uuids_to_keep_in_state = set()
        for lfs_uuid, lfs in state['large_file_systems'].items():
            last_used_at = datetime.fromisoformat(lfs['last_used_at'])
            lfs_time_to_live_in_days = 60 if lfs['state'] == 'ready' else 7

            if last_used_at < current_time - timedelta(days=lfs_time_to_live_in_days):
                partition = state['storage_partitions'][lfs['storage_partition_uuid']]
                paths_to_delete.add(f"{partition['path']}/lfs/{lfs_uuid}")
            else:
                lfs_uuids_to_keep_in_state.add(lfs_uuid)

        state['large_file_systems'] = {
            key: value for key, value in state['large_file_systems'].items() if key in lfs_uuids_to_keep_in_state
        }

    for path in paths_to_delete:
        logger_no_user_data.info(f'Deleting {path}...')
