from typing import Optional
"""
DatasetAnalyzer: Centralized dataset size analysis for cz-billing-generator.

Replaces scattered analysis logic with memory-efficient, reusable analysis capabilities.
"""

import logging
import math
from pathlib import Path

import polars as pl

from .models import CostImpactAnalysis, DatasetMetrics, ProviderMetrics


class DatasetAnalyzer:
    """Centralized dataset size analysis across all pipeline contexts."""

    def __init__(
        self,
        tenant_factor: int = 1,
        memory_limit_gb: float = 4.0,
        run_dir: Optional[Path] = None,
    ):
        """
        Initialize DatasetAnalyzer.

        Args:
            tenant_factor: N CloudZero tenants cost multiplication factor
            memory_limit_gb: Memory limit for chunked processing
            run_dir: Pipeline run directory for logging
        """
        self.tenant_factor = tenant_factor
        self.memory_limit_bytes = memory_limit_gb * 1024**3
        self.run_dir = run_dir
        self.stage_metrics = []

        # Keep simple memory management for DatasetAnalyzer
        # Full data management will be handled by DataManager in future

        # Business cost tiers based on real CloudZero usage patterns
        self.COST_TIERS = [
            (200_000, "Demo"),
            (40_000_000, "CloudZero"),
            (500_000_000, "Large"),
            (float("inf"), "Largest"),
        ]

        # Connection estimation constants (from existing system)
        self.ROWS_PER_PAYLOAD = 250_000  # From SummaryLog
        self.PAYLOADS_PER_CONNECTION = 15  # From validator.py
        self.MAX_PAYLOAD_BYTES = 5 * 1024 * 1024  # From planner.py

        # Setup logging (replaces SummaryLog logging)
        self.logger = self._setup_logger(run_dir)

    def _setup_logger(self, run_dir: Optional[Path]) -> logging.Logger:
        """Setup logging to replace SummaryLog functionality."""
        logger = logging.getLogger("cz_billing_generator.analysis")
        logger.setLevel(logging.INFO)

        # Remove existing handlers to avoid duplication
        for handler in logger.handlers[:]:
            logger.removeHandler(handler)

        # Console handler for user-visible analysis output
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)
        console_handler.setFormatter(logging.Formatter("[%(levelname)s] %(message)s"))
        logger.addHandler(console_handler)

        # File handler if run directory provided
        if run_dir:
            log_path = run_dir / "dataset_analysis.log"
            file_handler = logging.FileHandler(log_path)
            file_handler.setLevel(logging.DEBUG)
            file_handler.setFormatter(
                logging.Formatter("%(asctime)s [%(levelname)s] %(message)s")
            )
            logger.addHandler(file_handler)

        return logger

    def classify_cost_tier(self, monthly_line_items: int) -> str:
        """Classify dataset into cost tier based on monthly line items."""
        for threshold, tier in self.COST_TIERS:
            if monthly_line_items <= threshold:
                return tier
        return "Largest"  # Fallback

    def estimate_connections(self, monthly_line_items: int) -> int:
        """Estimate required API connections based on monthly line items."""
        # Based on existing upload planning logic: 250K rows per payload, 15 payloads per connection
        payloads_needed = math.ceil(monthly_line_items / self.ROWS_PER_PAYLOAD)
        connections_needed = math.ceil(payloads_needed / self.PAYLOADS_PER_CONNECTION)
        return max(1, connections_needed)

    def estimate_upload_duration(
        self, monthly_line_items: int, connections: int
    ) -> float:
        """Estimate upload duration in hours."""
        # Conservative estimate: ~1000 line items per second upload rate
        total_seconds = monthly_line_items / 1000
        # Parallel connections reduce time
        parallel_seconds = total_seconds / max(1, connections)
        return parallel_seconds / 3600  # Convert to hours

    def _estimate_dataframe_size(self, df: pl.DataFrame) -> int:
        """Estimate DataFrame memory size in bytes."""
        # Rough estimate: ~100 bytes per row (conservative)
        return df.height * 100

    def _categorize_ingest_speed(self, tier: str) -> str:
        """Categorize ingest speed for user communication."""
        speed_map = {
            "Demo": "Fast (optimized for rapid demo iteration)",
            "CloudZero": "Moderate (realistic feature testing scale)",
            "Large": "Slower (enterprise customer scale)",
            "Largest": "Slowest (maximum customer scale)",
        }
        return speed_map.get(tier, "Unknown")

    def analyze_dataframe(self, df: pl.DataFrame, stage_name: str) -> DatasetMetrics:
        """
        Analyze DataFrame with memory-efficient methods.

        This is the core analysis method that calculates provider-level metrics
        without creating data copies.
        """
        if df.is_empty():
            return DatasetMetrics(
                provider_breakdown={},
                total_line_items=0,
                cost_tier_classification={},
                connection_requirements={},
                memory_estimate_mb=0.0,
                tenant_multiplication_factor=self.tenant_factor,
                stage_name=stage_name,
            )

        # Use lazy evaluation for memory efficiency
        lazy_df = df.lazy()

        # Calculate provider-level metrics using efficient aggregations
        provider_stats = (
            lazy_df.group_by("cloud_provider")
            .agg(
                [
                    pl.len().alias("total_line_items"),
                    pl.col("timestamp")
                    .str.slice(0, 7)
                    .n_unique()
                    .alias("unique_months"),
                ]
            )
            .collect()
        )

        # Calculate monthly line item counts separately for each provider
        monthly_stats = (
            lazy_df.group_by(
                ["cloud_provider", pl.col("timestamp").str.slice(0, 7).alias("month")]
            )
            .agg([pl.len().alias("monthly_count")])
            .group_by("cloud_provider")
            .agg(
                [
                    pl.col("monthly_count").max().alias("max_monthly"),
                    pl.col("monthly_count").min().alias("min_monthly"),
                ]
            )
            .collect()
        )

        provider_breakdown = {}
        cost_tier_classification = {}
        connection_requirements = {}
        total_line_items = 0

        # Join provider stats with monthly stats
        monthly_stats_dict = {
            row["cloud_provider"]: row for row in monthly_stats.iter_rows(named=True)
        }

        for row in provider_stats.iter_rows(named=True):
            provider = row["cloud_provider"]
            total_months = row["unique_months"]
            provider_total = row["total_line_items"]

            # Get monthly stats for this provider
            monthly_info = monthly_stats_dict.get(
                provider, {"max_monthly": provider_total, "min_monthly": provider_total}
            )
            max_monthly = monthly_info["max_monthly"]
            min_monthly = monthly_info["min_monthly"]

            # Estimate unique billing combinations (approximation)
            # This is a rough estimate - in practice would need actual distinct count
            unique_combinations = min(provider_total // total_months, 10000)

            # Classify cost tier
            cost_tier = self.classify_cost_tier(max_monthly)
            cost_tier_classification[provider] = cost_tier

            # Estimate connections
            connections = self.estimate_connections(max_monthly)
            connection_requirements[provider] = connections

            # Estimate upload duration
            upload_duration = self.estimate_upload_duration(max_monthly, connections)

            provider_breakdown[provider] = ProviderMetrics(
                provider=provider,
                max_monthly_line_items=max_monthly,
                min_monthly_line_items=min_monthly,
                total_months=total_months,
                unique_billing_combinations=unique_combinations,
                cost_tier=cost_tier,
                required_connections=connections,
                estimated_upload_duration_hours=upload_duration,
            )

            total_line_items += provider_total

        # Estimate memory usage
        memory_estimate_mb = self._estimate_dataframe_size(df) / (1024 * 1024)

        return DatasetMetrics(
            provider_breakdown=provider_breakdown,
            total_line_items=total_line_items,
            cost_tier_classification=cost_tier_classification,
            connection_requirements=connection_requirements,
            memory_estimate_mb=memory_estimate_mb,
            tenant_multiplication_factor=self.tenant_factor,
            stage_name=stage_name,
        )

    def _analyze_chunked(
        self, df: pl.DataFrame, stage_name: str, chunk_size_mb: float = 100
    ) -> DatasetMetrics:
        """Process large DataFrames in memory-efficient chunks."""
        estimated_size = self._estimate_dataframe_size(df)

        if estimated_size <= self.memory_limit_bytes:
            return self.analyze_dataframe(df, stage_name)

        self.logger.warning(
            f"Large dataset detected ({estimated_size / (1024 * 1024):.1f} MB), using basic chunked analysis"
        )

        # Simple chunking approach until DataManager is implemented
        # This will be replaced by DataManager's optimization
        return self.analyze_dataframe(df, stage_name)

    def _analyze_parquet_stream(
        self, parquet_dir: Path, stage_name: str
    ) -> DatasetMetrics:
        """Memory-efficient analysis of parquet files without full loading."""
        provider_breakdown = {}
        cost_tier_classification = {}
        connection_requirements = {}
        total_line_items = 0
        total_memory_mb = 0.0

        parquet_files = list(parquet_dir.glob("**/*.parquet"))
        if not parquet_files:
            self.logger.warning(f"No parquet files found in {parquet_dir}")
            return DatasetMetrics(
                provider_breakdown={},
                total_line_items=0,
                cost_tier_classification={},
                connection_requirements={},
                memory_estimate_mb=0.0,
                tenant_multiplication_factor=self.tenant_factor,
                stage_name=stage_name,
            )

        for parquet_file in parquet_files:
            # Extract provider from file path - handle different structures:
            # Structure 1: provider/month/data.parquet (e.g., AWS/2024-01/data.parquet)
            # Structure 2: provider/payload_XXX.parquet (e.g., AWS/payload_001.parquet)
            parent_dir = parquet_file.parent.name
            grandparent_dir = parquet_file.parent.parent.name

            # Check if parent directory looks like a date (YYYY-MM format)
            if (
                len(parent_dir) == 7
                and parent_dir[4] == "-"
                and parent_dir[:4].isdigit()
                and parent_dir[5:].isdigit()
            ):
                # Structure 1: provider/month/data.parquet
                provider = grandparent_dir
            else:
                # Structure 2: provider/payload_XXX.parquet
                provider = parent_dir

            # Use lazy scanning for memory efficiency
            lazy_df = pl.scan_parquet(parquet_file)

            try:
                # Calculate metrics for this file
                file_stats = lazy_df.select(
                    [
                        pl.len().alias("line_items"),
                        pl.col("timestamp").str.slice(0, 7).n_unique().alias("months"),
                        pl.col("timestamp")
                        .str.slice(0, 7)
                        .value_counts()
                        .struct.field("count")
                        .max()
                        .alias("max_monthly"),
                    ]
                ).collect()

                if file_stats.height > 0:
                    row = file_stats.row(0, named=True)
                    file_line_items = row["line_items"]
                    file_months = row["months"]
                    file_max_monthly = row["max_monthly"]

                    # Accumulate provider metrics
                    if provider not in provider_breakdown:
                        cost_tier = self.classify_cost_tier(file_max_monthly)
                        connections = self.estimate_connections(file_max_monthly)
                        upload_duration = self.estimate_upload_duration(
                            file_max_monthly, connections
                        )

                        provider_breakdown[provider] = ProviderMetrics(
                            provider=provider,
                            max_monthly_line_items=file_max_monthly,
                            min_monthly_line_items=file_max_monthly,
                            total_months=file_months,
                            unique_billing_combinations=0,  # Would need more analysis
                            cost_tier=cost_tier,
                            required_connections=connections,
                            estimated_upload_duration_hours=upload_duration,
                        )
                        cost_tier_classification[provider] = cost_tier
                        connection_requirements[provider] = connections
                    else:
                        # Update existing metrics
                        existing = provider_breakdown[provider]
                        existing.max_monthly_line_items = max(
                            existing.max_monthly_line_items, file_max_monthly
                        )
                        existing.min_monthly_line_items = min(
                            existing.min_monthly_line_items, file_max_monthly
                        )
                        existing.total_months += file_months

                    total_line_items += file_line_items
                    total_memory_mb += parquet_file.stat().st_size / (1024 * 1024)

            except Exception as e:
                self.logger.warning(f"Failed to analyze {parquet_file}: {e}")
                continue

        return DatasetMetrics(
            provider_breakdown=provider_breakdown,
            total_line_items=total_line_items,
            cost_tier_classification=cost_tier_classification,
            connection_requirements=connection_requirements,
            memory_estimate_mb=total_memory_mb,
            tenant_multiplication_factor=self.tenant_factor,
            stage_name=stage_name,
        )

    def analyze_payload_directory(self, payload_dir: Path) -> DatasetMetrics:
        """Analyze pre-generated payloads for upload planning."""
        return self._analyze_parquet_stream(payload_dir, "pre-upload")

    def classify_cost_impact(self, metrics: DatasetMetrics) -> CostImpactAnalysis:
        """Translate metrics to business cost impact."""
        overall_tier = metrics.overall_cost_tier

        # Generate benefits based on cost tier
        benefits = []
        if overall_tier == "Demo":
            benefits = [
                "Fast iteration for demo development",
                "Low CloudZero ingest costs",
                "Single connection per provider sufficient",
            ]
        elif overall_tier == "CloudZero":
            benefits = [
                "Realistic scale for CloudZero feature testing",
                "Good balance of realism vs. development efficiency",
                "Moderate upload times",
            ]

        # Generate moderate warnings (not alarmist)
        warnings = []
        if overall_tier in ["Large", "Largest"]:
            warnings = [
                "Higher CloudZero ingest costs than typical demo data",
                "Multiple API connections required per provider",
                "Longer upload times - consider scheduling during off-peak hours",
            ]
        elif overall_tier == "CloudZero" and metrics.tenant_multiplication_factor > 1:
            warnings = [
                f"Cost multiplied by {metrics.tenant_multiplication_factor}x for multiple CloudZero tenants",
                "Consider reducing dataset size for cost optimization",
            ]
        elif metrics.tenant_multiplication_factor > 1:
            warnings = [
                f"Cost multiplied by {metrics.tenant_multiplication_factor}x for multiple CloudZero tenants"
            ]

        # Generate optimization recommendations
        recommendations = []
        if overall_tier in ["CloudZero", "Large", "Largest"]:
            recommendations = [
                "Reduce time range (e.g., 1-3 months instead of full year)",
                "Reduce resource count per job",
                "Focus on specific cloud providers instead of all three",
                "Test with Demo scale first before generating larger datasets",
            ]

        return CostImpactAnalysis(
            overall_tier=overall_tier,
            cost_multiplication_factor=metrics.tenant_multiplication_factor,
            ingest_speed_category=self._categorize_ingest_speed(overall_tier),
            optimization_recommendations=recommendations,
            moderate_warnings=warnings,
            benefits=benefits,
        )

    def track_stage_completion(self, stage_name: str, metrics: DatasetMetrics):
        """Track pipeline stage completion with analysis (replaces SummaryLog.log_summary)."""
        self.stage_metrics.append((stage_name, metrics))
        self._log_stage_analysis(stage_name, metrics)

        # Basic storage tracking - will be enhanced by DataManager
        if self.run_dir:
            try:
                run_size_gb = self._calculate_directory_size(self.run_dir) / (1024**3)
                # Log storage info to structured logging only (not console)
                self.logger.debug(
                    "Run storage calculated", stage=stage_name, size_gb=run_size_gb
                )
            except Exception as e:
                self.logger.warning(f"Storage calculation failed for {stage_name}: {e}")

    def _calculate_directory_size(self, directory: Path) -> int:
        """Calculate directory size in bytes."""
        total_size = 0
        try:
            for file_path in directory.rglob("*"):
                if file_path.is_file():
                    total_size += file_path.stat().st_size
        except (OSError, PermissionError):
            pass
        return total_size

    def _log_stage_analysis(self, stage_name: str, metrics: DatasetMetrics):
        """Log stage analysis (enhanced version of SummaryLog output)."""
        self.logger.info(f"\n📊 {stage_name} Stage Analysis:")

        for provider, provider_metrics in metrics.provider_breakdown.items():
            self.logger.info(f"   {provider_metrics.summary()}")

        if metrics.tenant_multiplication_factor > 1:
            self.logger.info(
                f"   Cost Factor: {metrics.tenant_multiplication_factor}x CloudZero tenants"
            )

        # Add cost impact analysis
        cost_impact = self.classify_cost_impact(metrics)
        self.logger.info(f"   Overall Cost Tier: {cost_impact.overall_tier}")
        self.logger.info(f"   Ingest Speed: {cost_impact.ingest_speed_category}")

    def finalize_analysis(self):
        """Generate final comprehensive analysis (replaces SummaryLog final summary)."""
        if not self.stage_metrics:
            self.logger.info("No pipeline stages analyzed")
            return

        self.logger.info("\n" + "=" * 60)
        self.logger.info("FINAL DATASET ANALYSIS SUMMARY")
        self.logger.info("=" * 60)

        # Show progression through stages
        for stage_name, metrics in self.stage_metrics:
            self.logger.info(f"\n{stage_name} Stage:")
            self.logger.info(f"  Total Line Items: {metrics.total_line_items:,}")
            self.logger.info(f"  Memory Usage: {metrics.memory_estimate_mb:.1f} MB")

            for provider, provider_metrics in metrics.provider_breakdown.items():
                self.logger.info(
                    f"  {provider}: {provider_metrics.max_monthly_line_items:,}/month ({provider_metrics.cost_tier})"
                )

        # Final recommendations
        if self.stage_metrics:
            final_metrics = self.stage_metrics[-1][1]
            cost_impact = self.classify_cost_impact(final_metrics)

            if cost_impact.moderate_warnings:
                self.logger.info("\n⚠️  CONSIDERATIONS:")
                for warning in cost_impact.moderate_warnings:
                    self.logger.info(f"   {warning}")

            if cost_impact.optimization_recommendations:
                self.logger.info("\n💡 OPTIMIZATION OPTIONS:")
                for rec in cost_impact.optimization_recommendations:
                    self.logger.info(f"   {rec}")

        self.logger.info("\n" + "=" * 60)

    def get_stage_metrics(self, stage_name: str) -> Optional[DatasetMetrics]:
        """Get metrics for a specific stage."""
        for name, metrics in self.stage_metrics:
            if name == stage_name:
                return metrics
        return None

    def get_final_dataset_metrics(self) -> Optional[DatasetMetrics]:
        """Get final aggregated dataset metrics."""
        if not self.stage_metrics:
            return None

        # Use the last stage metrics as final metrics
        _, final_metrics = self.stage_metrics[-1]
        return final_metrics
