"""
Workload resource distribution validation script for K8s allocation.

Validates that all workloads in processed K8s datasets meet the minimum
resource requirements. Processes month-by-month for memory efficiency.
"""

import io
import sys
from collections import defaultdict
from pathlib import Path
from typing import Dict, List, Set, Tuple, Any, Optional

import polars as pl
import typer
from rich.console import Console
from rich.table import Table
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeElapsedColumn

from ..config import K8sConfig
from ...services.s3_data_service import S3DataService
from ...services.environment_service import EnvironmentService
from src.log_context import get_logger


class WorkloadResourceValidator:
    """
    Validates that K8s workloads meet minimum resource requirements.
    
    Features:
    - Month-by-month processing for memory efficiency
    - Detailed validation reporting
    - Cross-month consistency checking
    - Resource distribution statistics
    """
    
    def __init__(self, config: K8sConfig, console: Console):
        self.config = config
        self.console = console
        self.s3_service = S3DataService()
        self.logger = get_logger(__name__)
        
        # Validation state
        self.validation_results = {}
        self.total_workloads_processed = 0
        self.total_violations = 0
        
    def validate_dataset(self, month_filter: Optional[str] = None) -> Dict[str, Any]:
        """
        Validate workload resource requirements for the entire dataset.
        
        Args:
            month_filter: Optional specific month to validate (YYYY-MM format)
            
        Returns:
            Dictionary containing validation summary and results
        """
        self.console.print(f"🔍 [bold blue]Validating K8s workload resource distribution[/bold blue]")
        self.console.print(f"📊 Dataset: {self.config.dataset_name}")
        self.console.print(f"📋 Minimum resources per workload: {self.config.min_resources_per_workload}")
        
        # Discover available months
        months = self._discover_k8s_months()
        if not months:
            self.console.print("❌ [red]No K8s data found for dataset[/red]")
            return {"status": "no_data", "months": []}
        
        # Filter to specific month if provided
        if month_filter:
            if month_filter in months:
                months = [month_filter]
                self.console.print(f"🗓️  Validating specific month: {month_filter}")
            else:
                self.console.print(f"❌ [red]Month {month_filter} not found in dataset[/red]")
                return {"status": "month_not_found", "available_months": months}
        else:
            self.console.print(f"🗓️  Validating {len(months)} months: {', '.join(sorted(months))}")
        
        # Process each month
        validation_summary = {
            "status": "success",
            "dataset_name": self.config.dataset_name,
            "months_processed": 0,
            "months_failed": 0,
            "total_workloads": 0,
            "total_violations": 0,
            "month_results": {},
            "cross_month_consistency": True
        }
        
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            console=self.console
        ) as progress:
            task = progress.add_task("Validating months...", total=len(months))
            
            for month in sorted(months):
                try:
                    month_result = self._validate_month(month)
                    validation_summary["month_results"][month] = month_result
                    validation_summary["months_processed"] += 1
                    validation_summary["total_workloads"] += month_result["total_workloads"]
                    validation_summary["total_violations"] += month_result["violations"]
                    
                    if month_result["violations"] > 0:
                        validation_summary["status"] = "violations_found"
                    
                except Exception as e:
                    self.logger.error(f"Failed to validate month {month}", error=str(e))
                    validation_summary["months_failed"] += 1
                    validation_summary["month_results"][month] = {
                        "status": "error",
                        "error": str(e)
                    }
                
                progress.advance(task)
        
        # Check cross-month consistency
        if len(validation_summary["month_results"]) > 1:
            validation_summary["cross_month_consistency"] = self._check_cross_month_consistency(
                validation_summary["month_results"]
            )
        
        # Display summary
        self._display_validation_summary(validation_summary)
        
        return validation_summary
    
    def _validate_month(self, month: str) -> Dict[str, Any]:
        """
        Validate workload resource requirements for a specific month.
        
        Args:
            month: Month to validate (YYYY-MM format)
            
        Returns:
            Dictionary containing month validation results
        """
        self.logger.debug(f"Validating K8s data for month {month}")
        
        # Load K8s data for the month
        k8s_data = self._load_k8s_month_data(month)
        if k8s_data is None or k8s_data.height == 0:
            return {
                "status": "no_data",
                "month": month,
                "total_workloads": 0,
                "violations": 0
            }
        
        # Analyze workload-resource mapping
        workload_resources = self._analyze_workload_resources(k8s_data)
        
        # Check minimum resource requirements
        violations = []
        for workload, resource_count in workload_resources.items():
            if resource_count < self.config.min_resources_per_workload:
                violations.append({
                    "workload": workload,
                    "resource_count": resource_count,
                    "minimum_required": self.config.min_resources_per_workload
                })
        
        # Calculate statistics
        resource_counts = list(workload_resources.values())
        month_result = {
            "status": "pass" if len(violations) == 0 else "violations",
            "month": month,
            "total_workloads": len(workload_resources),
            "violations": len(violations),
            "violation_details": violations[:10],  # Limit to first 10 for reporting
            "statistics": {
                "min_resources": min(resource_counts) if resource_counts else 0,
                "max_resources": max(resource_counts) if resource_counts else 0,
                "avg_resources": sum(resource_counts) / len(resource_counts) if resource_counts else 0,
                "workloads_below_minimum": len(violations)
            }
        }
        
        self.logger.info(
            f"Month {month} validation completed",
            total_workloads=month_result["total_workloads"],
            violations=month_result["violations"],
            min_resources_per_workload=self.config.min_resources_per_workload
        )
        
        return month_result
    
    def _analyze_workload_resources(self, k8s_data: pl.DataFrame) -> Dict[str, int]:
        """
        Analyze workload-to-resource mapping from K8s data.
        
        Args:
            k8s_data: K8s allocation data for a month
            
        Returns:
            Dictionary mapping workload -> resource count
        """
        # Filter to only rows with K8s allocation (non-null workloads)
        k8s_allocated = k8s_data.filter(pl.col("k8s_workload").is_not_null())
        
        if k8s_allocated.height == 0:
            return {}
        
        # Get resource ID column name using ColumnMapper
        from ..services.column_mapper import column_mapper
        resource_id_col = column_mapper.get_column_name(k8s_allocated, "ResourceID")
        if not resource_id_col:
            raise ValueError("ResourceID column not found in DataFrame")
        
        # Group by workload and count unique resources
        workload_resource_counts = (
            k8s_allocated
            .group_by("k8s_workload")
            .agg(pl.col(resource_id_col).n_unique().alias("resource_count"))
        )
        
        # Convert to dictionary
        return dict(zip(
            workload_resource_counts["k8s_workload"].to_list(),
            workload_resource_counts["resource_count"].to_list()
        ))
    
    def _check_cross_month_consistency(self, month_results: Dict[str, Any]) -> bool:
        """
        Check if workload topology is consistent across months.
        
        Args:
            month_results: Validation results for all months
            
        Returns:
            True if topology is consistent across months, False otherwise
        """
        workload_sets = []
        
        for month, result in month_results.items():
            if result.get("status") == "error" or result.get("total_workloads", 0) == 0:
                continue
                
            # For detailed consistency check, we'd need to reload data
            # For now, just check if workload counts are similar
            workload_sets.append(result.get("total_workloads", 0))
        
        if len(workload_sets) < 2:
            return True
        
        # Check if workload counts vary significantly (more than 5%)
        min_count = min(workload_sets)
        max_count = max(workload_sets)
        
        if min_count == 0:
            return False
        
        variance = (max_count - min_count) / min_count
        return variance <= 0.05  # Allow 5% variance
    
    def _discover_k8s_months(self) -> List[str]:
        """
        Discover available months with K8s data.
        
        Returns:
            List of available months in YYYY-MM format
        """
        try:
            prefix = f"czbillgen/datasets/{self.config.dataset_name}/k8s/"
            months = set()
            
            # List all parquet files in K8s data folder (structure: k8s/provider/month.parquet)
            paginator = self.s3_service.s3_client.get_paginator('list_objects_v2')
            page_iterator = paginator.paginate(
                Bucket=self.s3_service.bucket,
                Prefix=prefix
            )
            
            for page in page_iterator:
                if 'Contents' in page:
                    for obj in page['Contents']:
                        key = obj['Key']
                        # Extract month from paths like: czbillgen/datasets/dataset/k8s/AWS/2025-02.parquet
                        if key.endswith('.parquet'):
                            parts = key.split('/')
                            if len(parts) >= 5:  # dataset/k8s/provider/month.parquet
                                filename = parts[-1]  # "2025-02.parquet"
                                month = filename.replace('.parquet', '')  # "2025-02"
                                if len(month) == 7 and month[4] == '-':
                                    months.add(month)
            
            return sorted(list(months))
            
        except Exception as e:
            self.logger.error("Failed to discover K8s months", error=str(e))
            return []
    
    def _load_k8s_month_data(self, month: str) -> Optional[pl.DataFrame]:
        """
        Load K8s data for a specific month across all providers.
        
        Args:
            month: Month to load (YYYY-MM format)
            
        Returns:
            K8s data as Polars DataFrame, or None if not found
        """
        try:
            # Load data from all providers for the month
            all_month_data = []
            providers = ["AWS", "GCP", "Azure"]
            
            for provider in providers:
                file_key = f"czbillgen/datasets/{self.config.dataset_name}/k8s/{provider}/{month}.parquet"
                
                try:
                    # Check if file exists
                    self.s3_service.s3_client.head_object(
                        Bucket=self.s3_service.bucket,
                        Key=file_key
                    )
                    
                    # Read the file using S3 get_object
                    response = self.s3_service.s3_client.get_object(
                        Bucket=self.s3_service.bucket,
                        Key=file_key
                    )
                    provider_data = pl.read_parquet(io.BytesIO(response['Body'].read()))
                    if provider_data is not None and provider_data.height > 0:
                        all_month_data.append(provider_data)
                        
                except Exception:
                    # Provider file doesn't exist or can't be read - skip
                    continue
            
            if not all_month_data:
                return None
            
            # Combine all provider data for the month
            combined_data = pl.concat(all_month_data, how="vertical_relaxed")
            return combined_data
            
        except Exception as e:
            self.logger.error(f"Failed to load K8s data for month {month}", error=str(e))
            return None
    
    def _display_validation_summary(self, summary: Dict[str, Any]) -> None:
        """Display validation summary in a formatted table."""
        
        # Overall status
        status_color = "green" if summary["status"] == "success" else "red"
        self.console.print(f"\n📋 [bold {status_color}]Validation Summary[/bold {status_color}]")
        
        # Summary stats table
        stats_table = Table(title="Overall Statistics")
        stats_table.add_column("Metric", style="cyan")
        stats_table.add_column("Value", style="white")
        
        stats_table.add_row("Dataset", summary["dataset_name"])
        stats_table.add_row("Months Processed", str(summary["months_processed"]))
        stats_table.add_row("Months Failed", str(summary["months_failed"]))
        stats_table.add_row("Total Workloads", str(summary["total_workloads"]))
        stats_table.add_row("Total Violations", str(summary["total_violations"]))
        stats_table.add_row("Cross-Month Consistent", "✅ Yes" if summary["cross_month_consistency"] else "❌ No")
        
        self.console.print(stats_table)
        
        # Month-by-month results
        if summary["month_results"]:
            month_table = Table(title="Month-by-Month Results")
            month_table.add_column("Month", style="cyan")
            month_table.add_column("Status", style="white")
            month_table.add_column("Workloads", style="white")
            month_table.add_column("Violations", style="white")
            month_table.add_column("Min Resources", style="white")
            month_table.add_column("Avg Resources", style="white")
            
            for month, result in sorted(summary["month_results"].items()):
                if result.get("status") == "error":
                    month_table.add_row(
                        month,
                        "❌ Error",
                        "-",
                        "-",
                        "-",
                        "-"
                    )
                else:
                    status_icon = "✅ Pass" if result["violations"] == 0 else "❌ Fail"
                    stats = result.get("statistics", {})
                    month_table.add_row(
                        month,
                        status_icon,
                        str(result["total_workloads"]),
                        str(result["violations"]),
                        str(stats.get("min_resources", 0)),
                        f"{stats.get('avg_resources', 0):.1f}"
                    )
            
            self.console.print(month_table)
        
        # Show some violation examples if any
        if summary["total_violations"] > 0:
            self.console.print("\n❌ [bold red]Sample Violations:[/bold red]")
            violation_count = 0
            for month, result in summary["month_results"].items():
                if result.get("violation_details"):
                    self.console.print(f"  Month {month}:")
                    for violation in result["violation_details"][:3]:  # Show max 3 per month
                        self.console.print(
                            f"    • Workload '{violation['workload']}': "
                            f"{violation['resource_count']} resources "
                            f"(need {violation['minimum_required']})"
                        )
                        violation_count += 1
                        if violation_count >= 10:  # Limit total violations shown
                            break
                    if violation_count >= 10:
                        break


def main():
    """CLI entry point for workload resource validation."""
    app = typer.Typer(help="Validate K8s workload resource distribution")
    
    @app.command()
    def validate(
        config: str = typer.Option(..., help="Path to K8s configuration YAML file"),
        month: str = typer.Option(None, help="Specific month to validate (YYYY-MM)"),
        fail_on_violations: bool = typer.Option(
            True, 
            help="Exit with error code if violations found"
        ),
    ):
        """Validate workload resource distribution for K8s datasets."""
        console = Console()
        
        try:
            # Validate AWS credentials
            env_service = EnvironmentService()
            env_service.validate_aws_credentials_and_access()
            
            # Load configuration
            k8s_config = K8sConfig.from_yaml(config)
            k8s_config.validate()
            
            # Run validation
            validator = WorkloadResourceValidator(k8s_config, console)
            results = validator.validate_dataset(month)
            
            # Exit with appropriate code
            if fail_on_violations and results.get("total_violations", 0) > 0:
                console.print("\n❌ [bold red]Validation failed: workload resource violations found[/bold red]")
                raise typer.Exit(1)
            elif results.get("status") == "success":
                console.print("\n✅ [bold green]Validation passed: all workloads meet resource requirements[/bold green]")
            
        except FileNotFoundError as e:
            console.print(f"❌ [bold red]Configuration file not found:[/bold red] {str(e)}")
            raise typer.Exit(1)
        except ValueError as e:
            console.print(f"❌ [bold red]Configuration error:[/bold red] {str(e)}")
            raise typer.Exit(1)
        except Exception as e:
            console.print(f"❌ [bold red]Validation error:[/bold red] {str(e)}")
            raise typer.Exit(1)
    
    app()


if __name__ == "__main__":
    main()