"""Stage discovery for usage allocation input sources."""

from typing import Optional, List
from src.services.s3_data_service import S3DataService
from src.log_context.core import get_logger


class StageDiscovery:
    """
    Discovers available pipeline stages for a dataset.

    Single Responsibility: Find which stages have data for input.
    Does not load data - just discovery.

    Stage Priority (fallback order): metadata → k8s → generate
    """

    STAGE_PRIORITY: List[str] = ["metadata", "k8s", "generate"]

    def __init__(self, s3_service: S3DataService):
        """
        Initialize stage discovery.

        Args:
            s3_service: S3 data service (dependency injection)
        """
        self.s3_service = s3_service
        self.logger = get_logger(__name__)

    def find_first_available_stage(self, dataset_name: str) -> Optional[str]:
        """
        Find the first available stage with data in priority order.

        Tries stages in order: metadata → k8s → generate

        Args:
            dataset_name: Dataset identifier

        Returns:
            First available stage name, or None if no stages have data
        """
        for stage in self.STAGE_PRIORITY:
            if self.check_stage_exists(dataset_name, stage):
                self.logger.info(
                    f"Found available stage: {stage}",
                    dataset_name=dataset_name,
                    stage=stage
                )
                return stage

        self.logger.warning(
            "No stages with data found",
            dataset_name=dataset_name,
            checked_stages=self.STAGE_PRIORITY
        )
        return None

    def check_stage_exists(self, dataset_name: str, stage: str) -> bool:
        """
        Check if a specific stage has data.

        Args:
            dataset_name: Dataset identifier
            stage: Stage name to check

        Returns:
            True if stage exists and has data
        """
        try:
            available_stages = self.s3_service.list_stages(dataset_name)
            return stage in available_stages
        except Exception as e:
            self.logger.debug(
                f"Stage check failed: {stage}",
                dataset_name=dataset_name,
                error=str(e)
            )
            return False
