"""
Month Discovery Service for metadata allocation stage.

Provides month discovery and data validation capabilities:
- Discovers available provider-month combinations
- Validates DataFrame schemas for processing
- Handles missing data scenarios
- Provides processing summaries
"""

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

from src.log_context import get_logger
from .metadata_data_service import MetadataDataService


class MonthDiscoveryService:
    """
    Service for discovering and validating month data for metadata processing.

    Handles:
    - Provider-month combination discovery
    - DataFrame schema validation
    - Missing data scenario handling
    - Processing readiness assessment
    """

    def __init__(self, data_service: MetadataDataService):
        """Initialize with metadata data service."""
        self.logger = get_logger(__name__)
        self.data_service = data_service
        self.schema_cache = {}  # provider -> schema_validation_result

    def discover_processing_candidates(self, dataset_name: str, provider: str = None,
                                     month: str = None, force: bool = False) -> List[Tuple[str, str]]:
        """
        Discover provider-month combinations ready for processing.

        Args:
            dataset_name: Dataset name
            provider: Optional provider filter
            month: Optional month filter
            force: If True, include already completed months

        Returns:
            List of (provider, month) tuples ready for processing
        """
        try:
            # Get all available combinations
            all_combinations = self.data_service.list_available_months(dataset_name, provider)

            # Apply month filter if provided
            if month:
                all_combinations = [(p, m) for p, m in all_combinations if m == month]

            candidates = []

            for prov, mnth in all_combinations:
                # Check if already completed (unless force is True)
                if not force and self.data_service.check_completion_status(dataset_name, prov, mnth):
                    self.logger.debug(f"Skipping completed {prov} {mnth} (use --force to reprocess)")
                    continue

                candidates.append((prov, mnth))

            self.logger.info(f"Found {len(candidates)} provider-month combinations to process")
            return candidates

        except Exception as e:
            self.logger.error(f"Error discovering processing candidates: {e}")
            return []

    def validate_month_data(self, dataset_name: str, provider: str, month: str) -> Dict[str, Any]:
        """
        Validate month data for processing readiness.

        Args:
            dataset_name: Dataset name
            provider: Provider name
            month: Month string

        Returns:
            Validation report dictionary
        """
        validation_report = {
            "provider": provider,
            "month": month,
            "valid": False,
            "errors": [],
            "warnings": [],
            "row_count": 0,
            "data_source": None,
            "schema_info": {}
        }

        try:
            # Try to download the data
            df = self.data_service.download_month_data(dataset_name, provider, month)

            if df is None:
                validation_report["errors"].append("No data found in k8s or generate stages")
                return validation_report

            # Basic data validation
            validation_report["row_count"] = len(df)
            validation_report["schema_info"] = {
                "columns": df.columns,
                "column_count": len(df.columns)
            }

            # Check for required columns
            required_columns = ["cloud_provider", "service", "resource_id"]
            missing_columns = [col for col in required_columns if col not in df.columns]

            if missing_columns:
                validation_report["errors"].append(f"Missing required columns: {missing_columns}")
                return validation_report

            # Check for empty data
            if len(df) == 0:
                validation_report["warnings"].append("DataFrame is empty")
            else:
                # Additional validations for non-empty data
                self._validate_data_quality(df, validation_report)

            # Determine data source
            validation_report["data_source"] = self._determine_data_source(df)

            validation_report["valid"] = len(validation_report["errors"]) == 0
            return validation_report

        except Exception as e:
            validation_report["errors"].append(f"Validation error: {str(e)}")
            return validation_report

    def _validate_data_quality(self, df: pl.DataFrame, validation_report: Dict[str, Any]):
        """Validate data quality for processing."""
        # Check for null values in required columns
        required_columns = ["cloud_provider", "service", "resource_id"]

        for col in required_columns:
            if col in df.columns:
                null_count = df.select(pl.col(col).is_null().sum()).item()
                if null_count > 0:
                    validation_report["warnings"].append(f"Column '{col}' has {null_count} null values")

        # Check cloud_provider consistency
        if "cloud_provider" in df.columns:
            providers = df.select("cloud_provider").unique().to_series().to_list()
            if len(providers) > 1:
                validation_report["warnings"].append(f"Multiple cloud providers in data: {providers}")

        # Check for K8s columns
        k8s_columns = ["k8s_cluster", "k8s_namespace", "k8s_labels"]
        has_k8s_data = any(col in df.columns for col in k8s_columns)

        if has_k8s_data:
            validation_report["schema_info"]["has_k8s_data"] = True
            k8s_rows = len(df.filter(pl.col("k8s_cluster").is_not_null())) if "k8s_cluster" in df.columns else 0
            validation_report["schema_info"]["k8s_row_count"] = k8s_rows
        else:
            validation_report["schema_info"]["has_k8s_data"] = False

    def _determine_data_source(self, df: pl.DataFrame) -> str:
        """Determine if data came from k8s or generate stage."""
        # K8s stage data typically has k8s columns
        k8s_columns = ["k8s_cluster", "k8s_namespace", "k8s_labels"]
        has_k8s_columns = any(col in df.columns for col in k8s_columns)

        if has_k8s_columns:
            # Check if there's actual K8s data
            if "k8s_cluster" in df.columns:
                k8s_rows = len(df.filter(pl.col("k8s_cluster").is_not_null()))
                if k8s_rows > 0:
                    return "k8s"

        return "generate"

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

        Args:
            dataset_name: Dataset name
            provider: Optional provider filter
            month: Optional month filter
            force: Include completed months

        Returns:
            Processing summary dictionary
        """
        summary = {
            "dataset_name": dataset_name,
            "total_combinations": 0,
            "processing_candidates": 0,
            "completed_combinations": 0,
            "validation_results": [],
            "providers": [],
            "months": [],
            "data_sources": {"k8s": 0, "generate": 0},
            "total_rows": 0
        }

        try:
            # Get all available combinations
            all_combinations = self.data_service.list_available_months(dataset_name, provider)
            summary["total_combinations"] = len(all_combinations)

            # Apply month filter
            if month:
                all_combinations = [(p, m) for p, m in all_combinations if m == month]

            # Track providers and months
            providers = set()
            months = set()
            processing_candidates = []
            completed_count = 0
            providers_schema_validated = set()  # Track which providers we've schema-validated

            for prov, mnth in all_combinations:
                providers.add(prov)
                months.add(mnth)

                # Check completion status
                is_completed = self.data_service.check_completion_status(dataset_name, prov, mnth)

                if is_completed:
                    completed_count += 1
                    if not force:
                        continue

                processing_candidates.append((prov, mnth))

                # Do lightweight availability check for all months
                lightweight_check = self._lightweight_availability_check(dataset_name, prov, mnth)

                # Only do schema validation for first month of each provider
                if lightweight_check["available"] and prov not in providers_schema_validated:
                    # Perform schema validation for this provider (cached)
                    schema_validation = self.validate_provider_schema(dataset_name, prov, mnth)
                    summary["validation_results"].append(schema_validation)
                    providers_schema_validated.add(prov)

                    if schema_validation["valid"]:
                        # Track data sources efficiently
                        data_source = lightweight_check.get("data_source", "k8s")
                        if data_source in summary["data_sources"]:
                            summary["data_sources"][data_source] += 1
                else:
                    # For subsequent months of same provider, just do availability check
                    summary["validation_results"].append(lightweight_check)

                    if lightweight_check["available"]:
                        # Track data sources
                        data_source = lightweight_check.get("data_source", "k8s")
                        if data_source in summary["data_sources"]:
                            summary["data_sources"][data_source] += 1

            summary["providers"] = sorted(list(providers))
            summary["months"] = sorted(list(months))
            summary["processing_candidates"] = len(processing_candidates)
            summary["completed_combinations"] = completed_count

            return summary

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

    def _lightweight_availability_check(self, dataset_name: str, provider: str, month: str) -> Dict[str, Any]:
        """
        Lightweight check for data availability without downloading data.

        Args:
            dataset_name: Dataset name
            provider: Provider name
            month: Month string

        Returns:
            Basic availability report without expensive operations
        """
        report = {
            "provider": provider,
            "month": month,
            "available": False,
            "data_source": None
        }

        try:
            # Check k8s stage first (primary source)
            k8s_available = self.data_service._check_data_availability(dataset_name, provider, month, "k8s")
            if k8s_available:
                report["available"] = True
                report["data_source"] = "k8s"
                return report

            # Fallback to generate stage
            generate_available = self.data_service._check_data_availability(dataset_name, provider, month, "generate")
            if generate_available:
                report["available"] = True
                report["data_source"] = "generate"
                return report

        except Exception as e:
            self.logger.debug(f"Error checking availability for {provider}/{month}: {e}")

        return report

    def validate_provider_schema(self, dataset_name: str, provider: str, sample_month: str) -> Dict[str, Any]:
        """
        Validate schema for a provider using one sample month (cached per provider).

        Args:
            dataset_name: Dataset name
            provider: Provider name
            sample_month: Sample month to check schema

        Returns:
            Schema validation report with caching
        """
        # Check cache first
        cache_key = f"{dataset_name}:{provider}"
        if cache_key in self.schema_cache:
            cached_result = self.schema_cache[cache_key].copy()
            # Update month info for current request
            cached_result["month"] = sample_month
            return cached_result

        # Perform full schema validation for first time
        validation_report = {
            "provider": provider,
            "month": sample_month,
            "valid": False,
            "errors": [],
            "warnings": [],
            "schema_info": {}
        }

        try:
            # Download sample data for schema validation
            df = self.data_service.download_month_data(dataset_name, provider, sample_month)

            if df is None:
                validation_report["errors"].append("No data found in k8s or generate stages")
                self.schema_cache[cache_key] = validation_report
                return validation_report

            # Schema validation only - no row counting
            validation_report["schema_info"] = {
                "columns": df.columns,
                "column_count": len(df.columns)
            }

            # Check for required columns
            required_columns = ["cloud_provider", "service", "resource_id"]
            missing_columns = [col for col in required_columns if col not in df.columns]

            if missing_columns:
                validation_report["errors"].append(f"Missing required columns: {missing_columns}")
            elif len(df) == 0:
                validation_report["warnings"].append("DataFrame is empty")
            else:
                validation_report["valid"] = True

            # Cache the result for this provider
            self.schema_cache[cache_key] = validation_report.copy()

        except Exception as e:
            validation_report["errors"].append(f"Error validating schema: {str(e)}")
            self.schema_cache[cache_key] = validation_report

        return validation_report