"""
Provider copy optimization for time-shift operations.

After time-shifting first org, copies drops to remaining orgs via S3-to-S3 copy.
"""

import boto3
from typing import List
from ..models.time_shift_models import DropCopyResult
from .billing_period_helper import BillingPeriodHelper
from .manifest_manager import ManifestManager


class TimeShiftProviderCopier:
    """
    Copies time-shifted drops across organizations.

    After time-shift transforms first org's data, copies the result
    to remaining orgs using S3-to-S3 operations (no transformation needed).
    """

    def __init__(
        self,
        manifest_manager: ManifestManager,
        billing_helper: BillingPeriodHelper,
        s3_bucket: str,
        s3_root_folder: str = "piedpiper-custom-billing"
    ):
        """
        Initialize provider copier.

        Args:
            manifest_manager: Manifest manager service
            billing_helper: Billing period helper service
            s3_bucket: S3 bucket name
            s3_root_folder: Root folder for billing data
        """
        self.manifest_manager = manifest_manager
        self.billing_helper = billing_helper
        self.s3_bucket = s3_bucket
        self.s3_root_folder = s3_root_folder
        self.s3_client = boto3.client('s3')

    def copy_provider_to_orgs(
        self,
        source_org: str,
        target_orgs: List[str],
        provider: str,
        source_month: str,
        target_month: str
    ) -> List[DropCopyResult]:
        """
        Copy provider's time-shifted data to multiple organizations.

        Args:
            source_org: Organization with time-shifted data
            target_orgs: Organizations to copy to
            provider: Provider name
            source_month: Source month (YYYY-MM)
            target_month: Target month (YYYY-MM)

        Returns:
            List of DropCopyResults for each target org
        """
        import time

        # Convert to billing period
        target_period = self.billing_helper.month_to_billing_period(target_month)

        # Get source drop info
        source_drop_id = self._get_drop_id(source_org, provider, target_period)
        source_prefix = self.billing_helper.build_drop_prefix(
            self.s3_root_folder, source_org, provider, target_period, source_drop_id
        )

        # List files in source drop
        source_files = self._list_files(source_prefix)

        results = []
        for target_org in target_orgs:
            start_time = time.time()

            try:
                # Generate new drop ID for target org
                target_drop_id = self.billing_helper.generate_drop_id()
                target_prefix = self.billing_helper.build_drop_prefix(
                    self.s3_root_folder, target_org, provider, target_period, target_drop_id
                )

                # Copy all files
                self._copy_files(source_files, source_prefix, target_prefix)

                # Create manifest
                self._create_manifest(target_org, provider, target_period, target_drop_id)

                processing_time = time.time() - start_time

                results.append(DropCopyResult(
                    org=target_org,
                    provider=provider,
                    source_month=source_month,
                    target_month=target_month,
                    source_billing_period=target_period,
                    target_billing_period=target_period,
                    source_drop_id=source_drop_id,
                    target_drop_id=target_drop_id,
                    files_copied=len(source_files),
                    total_records=0,  # S3 copy doesn't count records
                    processing_time_sec=processing_time,
                    success=True
                ))

            except Exception as e:
                processing_time = time.time() - start_time
                results.append(DropCopyResult(
                    org=target_org,
                    provider=provider,
                    source_month=source_month,
                    target_month=target_month,
                    processing_time_sec=processing_time,
                    success=False,
                    error_message=str(e)
                ))

        return results

    def _get_drop_id(self, org: str, provider: str, billing_period: str) -> str:
        """Get drop ID from manifest."""
        manifest_key = self.billing_helper.build_manifest_key(
            self.s3_root_folder, org, provider, billing_period
        )
        manifest = self.manifest_manager.read_manifest(manifest_key)

        if 'current_drop_id' not in manifest:
            raise RuntimeError(
                f"Invalid manifest - missing current_drop_id: {manifest_key}"
            )

        return manifest['current_drop_id']

    def _list_files(self, prefix: str) -> List[str]:
        """List all files in S3 prefix."""
        files = []
        paginator = self.s3_client.get_paginator('list_objects_v2')

        for page in paginator.paginate(Bucket=self.s3_bucket, Prefix=prefix):
            if 'Contents' not in page:
                continue

            for obj in page['Contents']:
                key = obj['Key']
                # Only include data files, not directories
                if not key.endswith('/'):
                    files.append(key)

        return files

    def _copy_files(
        self, files: List[str], source_prefix: str, target_prefix: str
    ) -> None:
        """Copy files from source to target via S3-to-S3 copy."""
        for source_key in files:
            # Extract filename
            filename = source_key.split('/')[-1]
            target_key = f"{target_prefix}{filename}"

            # S3-to-S3 copy
            self.s3_client.copy_object(
                Bucket=self.s3_bucket,
                CopySource={'Bucket': self.s3_bucket, 'Key': source_key},
                Key=target_key
            )

    def _create_manifest(
        self, org: str, provider: str, billing_period: str, drop_id: str
    ) -> None:
        """Create manifest for target org."""
        manifest_key = self.billing_helper.build_manifest_key(
            self.s3_root_folder, org, provider, billing_period
        )
        manifest_data = self.manifest_manager.create_manifest(drop_id)
        self.manifest_manager.write_manifest(manifest_key, manifest_data)
