"""
Source month selector for time-shifting operations.

Identifies and validates 31-day source months for date transformation.
Only 31-day months can be transformed to any target month without data loss.
"""

import boto3
from typing import List, Optional
from calendar import monthrange
from datetime import datetime
from dateutil.relativedelta import relativedelta


class SourceMonthSelector:
    """
    Identifies valid 31-day source months.

    Ensures source month has 31 days and exists for all org/provider
    combinations in S3 before allowing time-shift operation.
    """

    def __init__(self, s3_bucket: str, s3_root_folder: str = "piedpiper-custom-billing"):
        """
        Initialize selector with S3 configuration.

        Args:
            s3_bucket: S3 bucket name
            s3_root_folder: Root folder for billing data
        """
        self.s3_bucket = s3_bucket
        self.s3_root_folder = s3_root_folder
        self.s3_client = boto3.client('s3')

    def select_source_month(
        self,
        orgs: List[str],
        providers: List[str],
        user_specified: Optional[str] = None
    ) -> str:
        """
        Select or validate 31-day source month.

        If user_specified:
            - Validate it has 31 days
            - Validate it exists for all org/provider combinations
        Else:
            - Auto-select first 31-day month found
            - Validate it exists for all org/provider combinations

        Args:
            orgs: List of organization names
            providers: List of provider names
            user_specified: Optional user-specified source month (YYYY-MM)

        Returns:
            Validated source month in YYYY-MM format

        Raises:
            ValueError: If no valid 31-day month found or specified month invalid
            RuntimeError: If source month missing for any org/provider combination
        """
        if user_specified:
            # Validate user-specified month
            self._validate_31_day_month(user_specified)
            source_month = user_specified
        else:
            # Auto-select first 31-day month
            source_month = self._find_first_31_day_month(orgs, providers)

        # Validate exists for all combinations
        self._validate_exists_for_all(source_month, orgs, providers)

        return source_month

    def _validate_31_day_month(self, month: str) -> None:
        """
        Validate month has 31 days.

        Args:
            month: Month in YYYY-MM format

        Raises:
            ValueError: If month doesn't have 31 days
        """
        try:
            year_str, month_str = month.split("-")
            year = int(year_str)
            month_num = int(month_str)
        except (ValueError, AttributeError):
            raise ValueError(
                f"Invalid month format: {month}. Expected YYYY-MM"
            )

        days = monthrange(year, month_num)[1]
        if days != 31:
            raise ValueError(
                f"Source month {month} has {days} days, not 31. "
                "Only 31-day months can be used as source for time-shifting."
            )

    def _find_first_31_day_month(
        self,
        orgs: List[str],
        providers: List[str]
    ) -> str:
        """
        Find first available 31-day month in S3.

        Args:
            orgs: Organizations to check
            providers: Providers to check

        Returns:
            First 31-day month found (YYYY-MM)

        Raises:
            RuntimeError: If no 31-day month found
        """
        # Check first org/provider combination
        org = orgs[0]
        provider = providers[0]
        prefix = f"{self.s3_root_folder}/{org}/{provider}/"

        # List available months
        response = self.s3_client.list_objects_v2(
            Bucket=self.s3_bucket,
            Prefix=prefix,
            Delimiter='/'
        )

        if 'CommonPrefixes' not in response:
            raise RuntimeError(
                f"No data found for {org}/{provider} in S3 "
                f"bucket {self.s3_bucket}"
            )

        # Find first 31-day month
        for prefix_dict in response['CommonPrefixes']:
            folder = prefix_dict['Prefix'].rstrip('/').split('/')[-1]
            # Folder format: drop_YYYY-MM-DD or YYYY-MM
            if 'drop_' in folder:
                month = '-'.join(folder.split('_')[1].split('-')[:2])
            else:
                month = folder

            try:
                self._validate_31_day_month(month)
                return month
            except ValueError:
                # Not a 31-day month, continue searching
                continue

        raise RuntimeError(
            f"No 31-day source month found for {org}/{provider}"
        )

    def _validate_exists_for_all(
        self,
        month: str,
        orgs: List[str],
        providers: List[str]
    ) -> None:
        """
        Validate source month exists for first org/provider combinations.

        Time-shift uses provider copy optimization: transforms first org's data,
        then S3-copies to remaining orgs. Therefore, only the first org needs
        source month data.

        Checks that manifest.json exists for the billing period corresponding
        to the source month for first org × all providers.

        Args:
            month: Source month to validate (YYYY-MM)
            orgs: All organizations (only first is validated)
            providers: All providers

        Raises:
            RuntimeError: If month missing for any first org/provider combination
        """
        missing = []
        billing_period = self._month_to_billing_period(month)

        # Only validate first org (used for transformation)
        # Remaining orgs receive S3 copies of transformed data
        first_org = orgs[0]

        for provider in providers:
            # Check for manifest.json in billing period
            manifest_key = f"{self.s3_root_folder}/{first_org}/{provider}/{billing_period}/manifest.json"

            try:
                self.s3_client.head_object(
                    Bucket=self.s3_bucket,
                    Key=manifest_key
                )
            except self.s3_client.exceptions.ClientError as e:
                if e.response['Error']['Code'] == '404':
                    missing.append(f"{first_org}/{provider}/{billing_period}")
                else:
                    # Other errors (permissions, etc.) should raise
                    raise

        if missing:
            raise RuntimeError(
                f"Source month {month} not found for first org ({first_org}): {', '.join(missing)}. "
                f"First org must have source month data for all providers (remaining orgs receive copies). "
                f"Missing manifests: {', '.join([f's3://{self.s3_bucket}/{self.s3_root_folder}/{m}/manifest.json' for m in missing])}"
            )

    def _month_to_billing_period(self, month: str) -> str:
        """
        Convert month to billing period format.

        Args:
            month: Month in YYYY-MM format

        Returns:
            Billing period in YYYYMMDD-YYYYMMDD format
        """
        date = datetime.strptime(month, "%Y-%m")
        start_date = date.replace(day=1)
        end_date = start_date + relativedelta(months=1)
        return f"{start_date.strftime('%Y%m%d')}-{end_date.strftime('%Y%m%d')}"
