from prefect import task
from prefect.logging import get_run_logger
from pathlib import Path
import os
import re
from tools.copernicus_data_request import CopernicusDataExtractor

import yaml

with open('config/params.yml') as file:
    config = yaml.safe_load(file)

start_date = config['start_date']
end_date = config['end_date']
request_script = config['request_script']


def get_existing_timestamps(consortium: str) -> set:
    """
    Extract timestamps from existing satellite data files.

    Args:
        consortium: Consortium name

    Returns:
        set: Set of existing timestamps
    """
    output_dir = Path(os.getcwd()) / "data" / "01_raw" / \
                 config['consortia_data_folders'][consortium] / "satellite_data"

    if not output_dir.exists():
        return set()

    existing_files = list(output_dir.glob('*.tiff'))
    existing_timestamps = set()

    for file in existing_files:
        try:
            # get timestamp from filename
            # --> format: consortium_vi_values_SENTINEL2_L2A_YYYY-MM-DDTHH_MM_SS.SSSSZ.tiff
            filename = file.stem
            # split and get the timestamp part (last part before .tiff)
            parts = filename.split('_')

            # find the date part (starts with year YYYY-)
            for i, part in enumerate(parts):
                if re.match(r'\d{4}-\d{2}-\d{2}T', part):
                    # reconstruct timestamp by joining remaining parts
                    timestamp_parts = parts[i:]
                    timestamp = '_'.join(timestamp_parts)
                    # repl underscores back to colons in time portion
                    # 2024-01-01T10_30_00.123456Z -> 2024-01-01T10:30:00.123456Z
                    timestamp = re.sub(r'T(\d{2})_(\d{2})_(\d{2})', r'T\1:\2:\3', timestamp)
                    existing_timestamps.add(timestamp)
                    break
        except Exception as e:
            # skip files that don't match expected format
            continue

    return existing_timestamps


@task(
    task_run_name="get_satellite_data_{consortium}",
    retries=2,
    retry_delay_seconds=60
)
def get_satellite_data(
        consortium: str,
        overwrite: bool = False
) -> dict:
    """
    Download satellite data for a specific consortium.

    Args:
        consortium: Name of the consortium
        overwrite: If True, re-download existing files

    Returns:
        dict: Summary with downloaded/skipped counts
    """
    logger = get_run_logger()
    logger.info(f'Starting satellite data collection for {consortium}')

    # init extractor
    try:
        extractor = CopernicusDataExtractor(
            consortium=consortium,
            evalscript=request_script
        )
    except Exception as e:
        logger.error(f'Failed to initialize extractor for {consortium}: {e}')
        raise

    # obtain available timestamps in catalogue
    timespan = [start_date, end_date]
    available_timestamps = extractor._get_timestamps(timespan=timespan)

    if not available_timestamps:
        logger.warning(f'No satellite data available for {consortium} in timespan {start_date} to {end_date}')
        return {
            'consortium': consortium,
            'status': 'no_data',
            'downloaded': 0,
            'skipped': 0,
            'failed': 0,
            'total_available': 0
        }

    logger.info(f'Found {len(available_timestamps)} available acquisitions')

    # check existing files if not overwriting
    timestamps_to_download = available_timestamps
    skipped_count = 0

    if not overwrite:
        existing_timestamps = get_existing_timestamps(consortium)

        if existing_timestamps:
            logger.info(f'Found {len(existing_timestamps)} existing files')

            # filter out existing timestamps
            # 16 is to trunctate to minutes. noticed seconds might not match, but not relevant at this scale
            timestamps_to_download = [
                ts for ts in available_timestamps
                if not any(ts[:16] in ex for ex in existing_timestamps)
            ]
            skipped_count = len(available_timestamps) - len(timestamps_to_download)

            if skipped_count > 0:
                logger.info(f'Skipping {skipped_count} already downloaded acquisitions')

    if not timestamps_to_download:
        logger.info(f'All data already downloaded for {consortium}')
        return {
            'consortium': consortium,
            'status': 'complete',
            'downloaded': 0,
            'skipped': skipped_count,
            'failed': 0,
            'total_available': len(available_timestamps)
        }

    # download only non existing data
    logger.info(f'Downloading {len(timestamps_to_download)} new acquisitions for {consortium}')

    # override datetimes to download only needed ones
    extractor.datetimes = timestamps_to_download

    try:
        extractor.data_request()

        # # newly downloaded files
        downloaded = len(timestamps_to_download)
        failed = 0

        result = {
            'consortium': consortium,
            'status': 'success',
            'downloaded': downloaded,
            'skipped': skipped_count,
            'failed': failed,
            'total_available': len(available_timestamps)
        }

        logger.info(f'Completed {consortium}: {downloaded} downloaded, {skipped_count} skipped')

    except Exception as e:
        logger.error(f'Download failed for {consortium}: {e}')
        result = {
            'consortium': consortium,
            'status': 'failed',
            'downloaded': 0,
            'skipped': skipped_count,
            'failed': len(timestamps_to_download),
            'total_available': len(available_timestamps),
            'error': str(e)
        }

    return result