"""
Line item tracking for CloudZero ingestion cost planning.

Tracks line item counts per provider and month to help estimate CloudZero
ingestion costs. Provides min/max months per provider for cost planning.
"""

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


@dataclass
class LineItemStats:
    """Statistics for line items in a specific month."""
    provider: str
    month: str
    line_item_count: int
    total_cost: float = 0.0
    generated_at: datetime = field(default_factory=datetime.now)


@dataclass
class ProviderStats:
    """Aggregated statistics for a provider."""
    provider: str
    total_line_items: int
    month_count: int
    min_month: str
    max_month: str
    avg_line_items_per_month: float
    months: List[str] = field(default_factory=list)
    annual_cost: float = 0.0
    peak_month_line_items: int = 0
    low_month_line_items: int = 0


class LineItemTracker:
    """
    Tracks line item counts for CloudZero ingestion cost planning.
    """
    
    def __init__(self):
        self.stats: Dict[str, Dict[str, LineItemStats]] = defaultdict(dict)
        self.total_line_items = 0
    
    def track_batch(
        self,
        provider: str,
        month: str,
        line_item_count: int,
        total_cost: float = 0.0
    ) -> None:
        """
        Track line items for a specific provider and month.
        
        Args:
            provider: Cloud provider (AWS, GCP, Azure)
            month: Month in YYYY-MM format
            line_item_count: Number of line items generated
            total_cost: Total cost for this provider/month batch
        """
        if month not in self.stats[provider]:
            self.stats[provider][month] = LineItemStats(
                provider=provider,
                month=month,
                line_item_count=0,
                total_cost=0.0
            )
        
        self.stats[provider][month].line_item_count += line_item_count
        self.stats[provider][month].total_cost += total_cost
        self.total_line_items += line_item_count
    
    def track_dataframe(
        self,
        df: pl.DataFrame,
        provider: str,
        month: str
    ) -> None:
        """
        Track line items from a generated DataFrame.
        
        Args:
            df: Generated billing data DataFrame
            provider: Cloud provider
            month: Month in YYYY-MM format
        """
        if not df.is_empty():
            line_item_count = len(df)
            # Calculate total cost if cost column exists
            total_cost = 0.0
            if "cost" in df.columns:
                total_cost = float(df["cost"].sum())
            
            self.track_batch(provider, month, line_item_count, total_cost)
    
    def get_provider_stats(self, provider: str) -> Optional[ProviderStats]:
        """
        Get aggregated statistics for a provider.
        
        Args:
            provider: Cloud provider
            
        Returns:
            ProviderStats object or None if provider not found
        """
        if provider not in self.stats:
            return None
        
        provider_data = self.stats[provider]
        months = sorted(provider_data.keys())
        
        if not months:
            return None
        
        total_line_items = sum(stat.line_item_count for stat in provider_data.values())
        total_cost = sum(stat.total_cost for stat in provider_data.values())
        
        # Calculate annual cost (extrapolate from available months)
        monthly_avg_cost = total_cost / len(months) if len(months) > 0 else 0
        annual_cost = monthly_avg_cost * 12
        
        # Find peak and low month line item counts
        line_item_counts = [stat.line_item_count for stat in provider_data.values()]
        peak_month_line_items = max(line_item_counts) if line_item_counts else 0
        low_month_line_items = min(line_item_counts) if line_item_counts else 0
        
        return ProviderStats(
            provider=provider,
            total_line_items=total_line_items,
            month_count=len(months),
            min_month=months[0],
            max_month=months[-1],
            avg_line_items_per_month=total_line_items / len(months),
            months=months,
            annual_cost=annual_cost,
            peak_month_line_items=peak_month_line_items,
            low_month_line_items=low_month_line_items
        )
    
    def get_all_provider_stats(self) -> List[ProviderStats]:
        """
        Get statistics for all tracked providers.
        
        Returns:
            List of ProviderStats objects
        """
        stats = []
        for provider in self.stats.keys():
            provider_stats = self.get_provider_stats(provider)
            if provider_stats:
                stats.append(provider_stats)
        
        return sorted(stats, key=lambda x: x.provider)
    
    def get_monthly_breakdown(self) -> Dict[str, Dict[str, int]]:
        """
        Get monthly breakdown of line items by provider.
        
        Returns:
            Dictionary mapping months to provider line item counts
        """
        monthly_breakdown = defaultdict(dict)
        
        for provider, provider_data in self.stats.items():
            for month, stats in provider_data.items():
                monthly_breakdown[month][provider] = stats.line_item_count
        
        return dict(monthly_breakdown)
    
    def get_summary_stats(self) -> Dict[str, any]:
        """
        Get overall summary statistics.
        
        Returns:
            Dictionary with summary statistics
        """
        provider_stats = self.get_all_provider_stats()
        
        if not provider_stats:
            return {
                "total_line_items": 0,
                "total_providers": 0,
                "total_months": 0,
                "date_range": None
            }
        
        all_months = set()
        for stats in provider_stats:
            all_months.update(stats.months)
        
        sorted_months = sorted(all_months)
        
        return {
            "total_line_items": self.total_line_items,
            "total_providers": len(provider_stats),
            "total_months": len(sorted_months),
            "date_range": {
                "min_month": sorted_months[0] if sorted_months else None,
                "max_month": sorted_months[-1] if sorted_months else None
            },
            "providers": [stats.provider for stats in provider_stats]
        }
    
    def format_summary_table(self) -> str:
        """
        Format summary statistics as a readable table.
        
        Returns:
            Formatted table string
        """
        provider_stats = self.get_all_provider_stats()
        
        if not provider_stats:
            return "No line items tracked."
        
        lines = []
        lines.append("Line Item Tracking Summary")
        lines.append("=" * 50)
        
        # Sort providers by annual cost (descending)
        provider_stats.sort(key=lambda x: x.annual_cost, reverse=True)
        
        # Provider breakdown
        lines.append("\nProvider Breakdown:")
        lines.append("┌─────────────┬─────────────┬─────────────┬─────────────┐")
        lines.append("│ Provider    │ Cost/Year   │Peak Month   │ Low Month   │")
        lines.append("│             │             │Line Items   │ Line Items  │")
        lines.append("├─────────────┼─────────────┼─────────────┼─────────────┤")
        
        for stats in provider_stats:
            provider_name = stats.provider[:11]  # Truncate if too long
            
            # Format cost
            if stats.annual_cost >= 1_000_000_000_000:  # Trillions
                cost_str = f"${stats.annual_cost/1_000_000_000_000:.1f}T"
            elif stats.annual_cost >= 1_000_000_000:  # Billions
                cost_str = f"${stats.annual_cost/1_000_000_000:.1f}B"
            elif stats.annual_cost >= 1_000_000:  # Millions
                cost_str = f"${stats.annual_cost/1_000_000:.1f}M"
            elif stats.annual_cost >= 1_000:  # Thousands
                cost_str = f"${stats.annual_cost/1_000:.0f}K"
            else:
                cost_str = f"${stats.annual_cost:.0f}"
            
            # Format line item counts
            if stats.peak_month_line_items >= 1_000_000:
                peak_str = f"{stats.peak_month_line_items/1_000_000:.1f}M items"
            elif stats.peak_month_line_items >= 1_000:
                peak_str = f"{stats.peak_month_line_items/1_000:.0f}K items"
            else:
                peak_str = f"{stats.peak_month_line_items:,} items"
            
            if stats.low_month_line_items >= 1_000_000:
                low_str = f"{stats.low_month_line_items/1_000_000:.1f}M items"
            elif stats.low_month_line_items >= 1_000:
                low_str = f"{stats.low_month_line_items/1_000:.0f}K items"
            else:
                low_str = f"{stats.low_month_line_items:,} items"
            
            lines.append(f"│ {provider_name:<11} │ {cost_str:>11} │ {peak_str:>11} │ {low_str:>11} │")
        
        lines.append("└─────────────┴─────────────┴─────────────┴─────────────┘")
        
        # Overall summary
        summary = self.get_summary_stats()
        lines.append(f"\nTotal: {summary['total_line_items']:>10,} line items")
        lines.append(f"Providers: {summary['total_providers']:>6}")
        lines.append(f"Months: {summary['total_months']:>9}")
        
        if summary['date_range']['min_month']:
            lines.append(f"Date Range: {summary['date_range']['min_month']} to {summary['date_range']['max_month']}")
        
        return "\n".join(lines)
    
    def reset(self) -> None:
        """Reset all tracking statistics."""
        self.stats.clear()
        self.total_line_items = 0
    
    def export_to_dict(self) -> Dict[str, any]:
        """
        Export tracking data to dictionary format.
        
        Returns:
            Dictionary with all tracking data
        """
        export_data = {
            "summary": self.get_summary_stats(),
            "providers": {},
            "monthly_breakdown": self.get_monthly_breakdown()
        }
        
        for provider in self.stats.keys():
            provider_stats = self.get_provider_stats(provider)
            if provider_stats:
                export_data["providers"][provider] = {
                    "total_line_items": provider_stats.total_line_items,
                    "month_count": provider_stats.month_count,
                    "min_month": provider_stats.min_month,
                    "max_month": provider_stats.max_month,
                    "avg_line_items_per_month": provider_stats.avg_line_items_per_month,
                    "months": provider_stats.months
                }
        
        return export_data