"""Service for computing detailed dataset metadata analysis."""

import polars as pl
from datetime import datetime
from typing import Dict, List, Optional, Tuple
from collections import defaultdict

from ..s3_data_service import S3DataService
from .models import (
    ProviderCostBreakdown,
    ProviderRowBreakdown,
    SummaryInsights,
    DetailedMetadata,
    BasicMetadata
)


class MetadataComputationService:
    """Service for computing detailed dataset metadata from parquet files."""

    def __init__(self, s3_service: S3DataService):
        self.s3_service = s3_service

    def compute_detailed_metadata(self, basic_metadata: BasicMetadata) -> DetailedMetadata:
        """
        Compute detailed metadata for a dataset by analyzing parquet files.

        This is the slow operation that scans actual data files to compute
        provider breakdowns, costs, and insights.
        """
        dataset_id = basic_metadata.dataset_id

        # Find the most recent completed stage with data
        latest_stage = self._find_latest_stage_with_data(dataset_id, basic_metadata.stages_completed)

        if not latest_stage:
            # Return basic detailed metadata with empty breakdowns
            return self._create_empty_detailed_metadata(basic_metadata)

        # Compute provider breakdowns from the latest stage
        cost_breakdown = self._compute_provider_cost_breakdown(dataset_id, latest_stage)
        row_breakdown = self._compute_provider_row_breakdown(dataset_id, latest_stage)

        # Generate summary insights
        summary_insights = self._compute_summary_insights(cost_breakdown, row_breakdown)

        return DetailedMetadata(
            basic=basic_metadata,
            metadata_computed_at=datetime.now(),
            provider_cost_breakdown=cost_breakdown,
            provider_row_breakdown=row_breakdown,
            summary_insights=summary_insights
        )

    def _find_latest_stage_with_data(self, dataset_id: str, completed_stages: List[str]) -> Optional[str]:
        """Find the latest stage that has meaningful data for analysis."""
        # Preferred order: latest stages have more enriched data
        stage_priority = ["discount", "tax", "usage", "metadata", "k8s", "generate"]

        for stage in stage_priority:
            if stage in completed_stages:
                stage_path = f"czbillgen/datasets/{dataset_id}/{stage}/"
                try:
                    files = self.s3_service.list_files(stage_path, max_files=5)
                    # Look for parquet files (the actual data)
                    parquet_files = [f for f in files if f.endswith('.parquet')]
                    if parquet_files:
                        return stage
                except Exception:
                    continue

        return None

    def _compute_provider_cost_breakdown(self, dataset_id: str, stage: str) -> ProviderCostBreakdown:
        """Compute cost breakdown by provider and month."""
        provider_costs = defaultdict(lambda: defaultdict(float))

        try:
            # Get all provider-month combinations for this stage
            stage_path = f"czbillgen/datasets/{dataset_id}/{stage}/"
            folders = self.s3_service.list_folders(stage_path)

            for folder in folders:
                # Extract provider from folder name
                folder_name = folder.rstrip('/').split('/')[-1]
                if not folder_name:
                    continue

                # Folder format is typically: provider/ or provider-month/
                provider = folder_name.split('-')[0].upper()

                # Get parquet files in this provider folder
                provider_files = self.s3_service.list_files(folder)
                parquet_files = [f for f in provider_files if f.endswith('.parquet')]

                for parquet_file in parquet_files[:10]:  # Limit to avoid too much processing
                    try:
                        # Load and analyze the parquet file
                        df = self.s3_service.scan_parquet_from_s3(parquet_file).collect()

                        if df.height == 0:
                            continue

                        # Extract month and cost information
                        cost_data = self._extract_cost_data_from_df(df, provider)

                        for month, cost in cost_data.items():
                            provider_costs[provider][month] += cost

                    except Exception:
                        # Continue if individual file fails
                        continue

            # Calculate yearly totals
            final_costs = {}
            for provider, monthly_costs in provider_costs.items():
                final_costs[provider] = dict(monthly_costs)
                final_costs[provider]["yearly_total"] = sum(monthly_costs.values())

        except Exception:
            final_costs = {}

        return ProviderCostBreakdown(provider_costs=final_costs)

    def _compute_provider_row_breakdown(self, dataset_id: str, stage: str) -> ProviderRowBreakdown:
        """Compute row count breakdown by provider and month."""
        provider_rows = defaultdict(lambda: defaultdict(int))

        try:
            # List provider folders in the stage folder
            stage_path = f"czbillgen/datasets/{dataset_id}/{stage}/"
            provider_folders = self.s3_service.list_folders(stage_path)

            for folder in provider_folders:
                folder_name = folder.rstrip('/').split('/')[-1]
                if not folder_name:
                    continue

                # Extract provider name (handles formats like "aws-12345" -> "AWS")
                provider = folder_name.split('-')[0].upper()

                # Get parquet files in this provider folder
                provider_files = self.s3_service.list_files(folder)
                parquet_files = [f for f in provider_files if f.endswith('.parquet')]

                # Each parquet file represents one month of data
                # Simply count files and get row counts - no complex month extraction needed
                for i, parquet_file in enumerate(parquet_files, 1):
                    try:
                        # Get row count for this month's file
                        row_count = self.s3_service.get_parquet_row_count(parquet_file)
                        # Use generic month key - the count is what matters for "Total Months"
                        provider_rows[provider][f"month_{i}"] = row_count
                    except Exception:
                        # Still count the month even if we can't read row count
                        provider_rows[provider][f"month_{i}"] = 0

        except Exception:
            provider_rows = {}

        return ProviderRowBreakdown(provider_rows=dict(provider_rows))

    def _extract_cost_data_from_df(self, df: pl.DataFrame, provider: str) -> Dict[str, float]:
        """Extract cost data by month from a DataFrame."""
        cost_by_month = {}

        try:
            # Common cost column names across different stages
            cost_columns = ["cost", "unblended_cost", "blended_cost", "net_unblended_cost", "line_item_unblended_cost"]
            cost_col = None

            for col in cost_columns:
                if col in df.columns:
                    cost_col = col
                    break

            if not cost_col:
                return cost_by_month

            # Common date column names
            date_columns = ["usage_start_date", "date", "usage_date", "billing_period_start_date"]
            date_col = None

            for col in date_columns:
                if col in df.columns:
                    date_col = col
                    break

            if not date_col:
                # If no date column, return total cost with unknown month
                total_cost = df.select(pl.col(cost_col).sum()).item()
                cost_by_month["unknown"] = float(total_cost) if total_cost else 0.0
                return cost_by_month

            # Group by month and sum costs
            monthly_costs = (df
                .with_columns([
                    pl.col(date_col).dt.strftime("%Y-%m").alias("month")
                ])
                .group_by("month")
                .agg([
                    pl.col(cost_col).sum().alias("total_cost")
                ])
            )

            for row in monthly_costs.iter_rows(named=True):
                month = row["month"]
                total_cost = row["total_cost"]
                if month and total_cost:
                    cost_by_month[month] = float(total_cost)

        except Exception:
            # Fallback: try to get total cost without month breakdown
            try:
                if cost_col and cost_col in df.columns:
                    total_cost = df.select(pl.col(cost_col).sum()).item()
                    cost_by_month["unknown"] = float(total_cost) if total_cost else 0.0
            except Exception:
                pass

        return cost_by_month


    def _compute_summary_insights(self, cost_breakdown: ProviderCostBreakdown, row_breakdown: ProviderRowBreakdown) -> SummaryInsights:
        """Generate summary insights from cost and row data."""

        # Top providers by total cost
        provider_totals = []
        for provider, costs in cost_breakdown.provider_costs.items():
            total = costs.get("yearly_total", 0)
            if total > 0:
                provider_totals.append((provider, total))

        top_providers_by_cost = [p[0] for p in sorted(provider_totals, key=lambda x: x[1], reverse=True)]

        # Top providers by maximum monthly rows
        provider_max_rows = []
        for provider, rows in row_breakdown.provider_rows.items():
            max_monthly = max(rows.values()) if rows else 0
            if max_monthly > 0:
                provider_max_rows.append((provider, max_monthly))

        top_providers_by_max_monthly_rows = [p[0] for p in sorted(provider_max_rows, key=lambda x: x[1], reverse=True)]

        # Date ranges by provider
        date_ranges_by_provider = {}
        for provider, costs in cost_breakdown.provider_costs.items():
            months = [k for k in costs.keys() if k != "yearly_total" and k != "unknown"]
            if months:
                months.sort()
                if len(months) == 1:
                    date_ranges_by_provider[provider] = f"{months[0]}"
                else:
                    date_ranges_by_provider[provider] = f"{months[0]} to {months[-1]}"

        return SummaryInsights(
            top_providers_by_cost=top_providers_by_cost,
            top_providers_by_max_monthly_rows=top_providers_by_max_monthly_rows,
            date_ranges_by_provider=date_ranges_by_provider
        )

    def _create_empty_detailed_metadata(self, basic_metadata: BasicMetadata) -> DetailedMetadata:
        """Create detailed metadata with empty breakdowns when no data is available."""
        return DetailedMetadata(
            basic=basic_metadata,
            metadata_computed_at=datetime.now(),
            provider_cost_breakdown=ProviderCostBreakdown(provider_costs={}),
            provider_row_breakdown=ProviderRowBreakdown(provider_rows={}),
            summary_insights=SummaryInsights(
                top_providers_by_cost=[],
                top_providers_by_max_monthly_rows=[],
                date_ranges_by_provider={}
            )
        )