"""
Metadata Orchestrator for complete month-by-month processing workflow.

Coordinates the metadata allocation stage:
- Month discovery and validation
- Month-by-month processing with MetadataProcessor
- S3 data pipeline management
- Success tracking and error handling
- Memory-efficient operations
"""

from typing import List, Dict, Any, Optional
import polars as pl
from datetime import datetime

from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
from rich.console import Console

from src.log_context import get_logger
from .metadata_data_service import MetadataDataService
from .month_discovery_service import MonthDiscoveryService
from ..provider_normalizer import ProviderNormalizer
from ..processor import MetadataProcessor
from ..config import MetadataConfig


class MetadataOrchestrator:
    """
    Orchestrates complete metadata allocation workflow.

    Handles:
    - Month discovery and candidate selection
    - Month-by-month processing coordination
    - S3 data pipeline management
    - Success tracking and resume capability
    - Memory-efficient processing
    """

    def __init__(self, config: MetadataConfig):
        """Initialize orchestrator with configuration."""
        self.logger = get_logger(__name__)
        self.config = config
        self.console = Console()

        # Initialize services
        self.data_service = MetadataDataService()
        self.discovery_service = MonthDiscoveryService(self.data_service)
        self.processor = MetadataProcessor(config)

        # Processing stats
        self.stats = {
            "total_candidates": 0,
            "processed_successfully": 0,
            "processing_errors": 0,
            "validation_errors": 0,
            "providers_skipped": 0,
            "total_rows_processed": 0,
            "providers_processed": set(),
            "months_processed": set()
        }

    def run_metadata_allocation(self, provider: str = None, month: str = None,
                              force: bool = False) -> Dict[str, Any]:
        """
        Run complete metadata allocation workflow.

        Args:
            provider: Optional provider filter
            month: Optional month filter
            force: If True, reprocess completed months

        Returns:
            Processing summary with results and statistics
        """
        try:
            self.logger.info("Starting metadata allocation workflow",
                           dataset=self.config.dataset_name, provider=provider, month=month)

            # Step 1: Discover processing candidates
            candidates = self.discovery_service.discover_processing_candidates(
                self.config.dataset_name, provider, month, force
            )

            self.stats["total_candidates"] = len(candidates)

            if not candidates:
                self.logger.warning("No processing candidates found")
                return self._create_summary("No candidates found")

            # Step 2: Process each provider-month combination with progress bar
            with Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                BarColumn(),
                TaskProgressColumn(),
                console=self.console
            ) as progress:
                task = progress.add_task(
                    "[cyan]Processing months...",
                    total=len(candidates)
                )

                for i, (prov, mnth) in enumerate(candidates, 1):
                    progress.update(
                        task,
                        description=f"[cyan]Processing {prov} {mnth} ({i}/{len(candidates)})..."
                    )

                    try:
                        self.logger.debug(f"Starting processing for {prov} {mnth}")
                        result = self._process_provider_month(prov, mnth)
                        self.logger.debug(f"Processing result for {prov} {mnth}: {result}")

                        if result == "skipped":
                            self.stats["providers_skipped"] += 1
                            progress.update(task, advance=1)
                            self.console.print(f"⏭️  [blue]{prov} {mnth} skipped - not configured[/blue]")
                        elif result:
                            self.stats["processed_successfully"] += 1
                            self.stats["providers_processed"].add(prov)
                            self.stats["months_processed"].add(mnth)
                            progress.update(task, advance=1)
                            self.console.print(f"✅ [green]{prov} {mnth} completed successfully[/green]")
                        else:
                            self.stats["processing_errors"] += 1
                            progress.update(task, advance=1)
                            self.console.print(f"❌ [red]{prov} {mnth} failed (returned False)[/red]")

                    except Exception as e:
                        self.logger.error(f"Exception processing {prov} {mnth}: {e}", exc_info=True)
                        self.stats["processing_errors"] += 1
                        progress.update(task, advance=1)
                        # Show full error message, not truncated
                        error_msg = str(e).replace('\n', ' ')  # Replace newlines with spaces for console display
                        self.console.print(f"❌ [red]{prov} {mnth} failed: {error_msg}[/red]")

            # Step 3: Write stage success marker if any processing succeeded
            if self.stats["processed_successfully"] > 0:
                stage_summary = self._create_summary("Completed")
                success_written = self.data_service.write_stage_success_marker(
                    self.config.dataset_name, stage_summary
                )
                if not success_written:
                    self.logger.warning("Failed to write stage success marker")

            # Step 4: Generate final summary
            summary = self._create_summary("Completed")
            self.logger.info("Metadata allocation completed",
                           processed=self.stats["processed_successfully"],
                           errors=self.stats["processing_errors"])

            return summary

        except Exception as e:
            self.logger.error(f"Metadata allocation workflow failed: {e}")
            raise RuntimeError(f"Metadata allocation failed: {e}") from e

    def _normalize_provider_name(self, provider: str) -> str:
        """Normalize provider name to match config expectations."""
        return ProviderNormalizer.normalize(provider)

    def _process_provider_month(self, provider: str, month: str):
        """
        Process a single provider-month combination.

        Args:
            provider: Provider name (original S3 name, e.g. 'AWS')
            month: Month string (YYYY-MM)

        Returns:
            "skipped" if provider not configured, True if processing succeeded, False otherwise
        """
        try:
            # Normalize provider name for config lookup
            config_provider = self._normalize_provider_name(provider)

            # Early check: Skip if provider not configured
            if not self.config.has_provider(config_provider):
                self.logger.info(f"Skipping {provider} {month} - not configured")
                return "skipped"

            # Step 1: Validate month data
            validation_report = self.discovery_service.validate_month_data(
                self.config.dataset_name, provider, month
            )

            if not validation_report["valid"]:
                self.logger.error(f"Validation failed for {provider} {month}: {validation_report['errors']}")
                self.stats["validation_errors"] += 1
                return False

            # Step 2: Download month data (use original S3 provider name)
            df = self.data_service.download_month_data(
                self.config.dataset_name, provider, month
            )

            if df is None:
                self.logger.error(f"No data available for {provider} {month}")
                return False

            row_count = len(df)

            # Show chunking info if dataset is large
            chunk_threshold = 5_000_000
            if row_count > chunk_threshold:
                num_chunks = (row_count + chunk_threshold - 1) // chunk_threshold
                self.console.print(f"  📦 Large dataset ({row_count:,} rows) - processing in {num_chunks} chunks")

            # Step 3: Process with MetadataProcessor (use normalized provider name for config lookup)
            self.logger.debug(f"Before processing: about to process {provider} {month}")
            processed_df = self.processor.process_provider_month(config_provider, month, df)
            self.logger.debug(f"After processing: got result with {processed_df.height if processed_df is not None else 0} rows")

            if processed_df is None:
                self.logger.error(f"Processing failed for {provider} {month}")
                return False

            # Free original DataFrame memory before upload
            self.logger.debug(f"Before df cleanup: freeing original DataFrame memory")
            del df
            import gc
            gc.collect()
            self.logger.debug(f"After df cleanup: original DataFrame freed")

            # Step 4: Upload processed data (use original S3 provider name)
            self.logger.debug(f"Before upload: starting upload for {provider} {month}")
            success = self.data_service.upload_month_data(
                self.config.dataset_name, provider, month, processed_df
            )
            self.logger.debug(f"After upload: upload {'succeeded' if success else 'failed'} for {provider} {month}")

            if success:
                self.stats["total_rows_processed"] += len(processed_df)
                self.logger.debug(f"Returning success for {provider} {month}")
                return True
            else:
                self.logger.error(f"Upload failed for {provider} {month}")
                return False

        except Exception as e:
            self.logger.error(f"Error processing {provider} {month}: {e}")
            # Log the full traceback for debugging
            import traceback
            self.logger.debug(f"Full traceback for {provider} {month}: {traceback.format_exc()}")
            return False

    def get_processing_summary(self, provider: str = None, month: str = None,
                             force: bool = False) -> Dict[str, Any]:
        """
        Get comprehensive processing summary without running processing.

        Args:
            provider: Optional provider filter
            month: Optional month filter
            force: Include completed months

        Returns:
            Processing summary dictionary
        """
        try:
            summary = self.discovery_service.get_processing_summary(
                self.config.dataset_name, provider, month, force
            )

            # Add processor configuration info
            summary["processor_config"] = {
                "configured_providers": [p.provider for p in self.config.providers],
                "total_provider_configs": len(self.config.providers)
            }

            # Add detailed processor summaries
            processor_summaries = {}
            for provider_config in self.config.providers:
                processor_summary = self.processor.get_processing_summary(provider_config.provider)
                processor_summaries[provider_config.provider] = processor_summary

            summary["processor_details"] = processor_summaries

            return summary

        except Exception as e:
            self.logger.error(f"Error generating processing summary: {e}")
            return {"error": str(e)}

    def validate_prerequisites(self) -> List[str]:
        """
        Validate prerequisites for metadata processing.

        Returns:
            List of validation errors (empty if all valid)
        """
        errors = []

        # 1. Check AWS credentials (needed for S3 data access)
        try:
            test_list = self.data_service.list_available_months(self.config.dataset_name)
            self.logger.debug("AWS credentials validated successfully")
        except Exception as e:
            errors.append(f"AWS credentials error: {e}")
            return errors  # No point continuing if we can't access S3

        # 2. Validate basic configuration requirements
        if not self.config.providers:
            errors.append("No providers configured")
            return errors

        # 3. Validate each configured provider has keys OR correlated groups
        for provider_config in self.config.providers:
            if not provider_config.keys and not provider_config.correlated_groups:
                errors.append(f"Provider '{provider_config.provider}' has no keys or correlated groups")

        # 4. Validate configured providers have accessible data
        try:
            summary = self.discovery_service.get_processing_summary(
                self.config.dataset_name, None, None, False
            )

            available_providers = set(summary.get("providers", []))
            configured_providers = set(p.provider for p in self.config.providers)

            # Normalize provider names for comparison (case-insensitive)
            normalized_available = set(self._normalize_provider_name(p) for p in available_providers)

            # Check if any configured providers have data
            providers_with_data = configured_providers.intersection(normalized_available)

            if not providers_with_data:
                configured_list = ', '.join(configured_providers)
                available_list = ', '.join(available_providers)
                errors.append(f"No data found for configured providers. "
                             f"Configured: [{configured_list}], "
                             f"Available: [{available_list}]")
        except Exception as e:
            errors.append(f"Data availability check error: {e}")

        return errors

    def _create_summary(self, status: str) -> Dict[str, Any]:
        """Create processing summary with current statistics."""
        return {
            "status": status,
            "dataset_name": self.config.dataset_name,
            "processing_stats": {
                "total_candidates": self.stats["total_candidates"],
                "processed_successfully": self.stats["processed_successfully"],
                "processing_errors": self.stats["processing_errors"],
                "validation_errors": self.stats["validation_errors"],
                "providers_skipped": self.stats["providers_skipped"],
                "total_rows_processed": self.stats["total_rows_processed"]
            },
            "coverage": {
                "providers_processed": sorted(list(self.stats["providers_processed"])),
                "months_processed": sorted(list(self.stats["months_processed"])),
                "unique_provider_months": self.stats["processed_successfully"]
            },
            "timestamp": datetime.utcnow().isoformat()
        }

    def cleanup_cache(self):
        """Clean up local cache files."""
        try:
            self.data_service.cleanup_cache(self.config.dataset_name)
            self.logger.info("Cache cleanup completed")
        except Exception as e:
            self.logger.warning(f"Cache cleanup failed: {e}")