"""
CloudZero Billing Format (CBF) mapping functionality.

Provides column mapping and validation for converting datasets to CBF format,
with S3-based default mapping and user-defined CSV mapping support.
"""

import csv
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple

import polars as pl
import yaml

from src.services.s3_data_service import S3DataService


def load_yaml_mapping(content: str) -> Tuple[List[Tuple[str, str]], Set[str]]:
    """Load CBF mapping from YAML content."""
    raw = yaml.safe_load(content)
    
    mapping = []
    required = set()
    
    for entry in raw.get("mappings", []):
        source = entry.get("source", "")
        cbf = entry["cbf"]
        is_required = entry.get("required", False)
        
        if source == "":
            if is_required:
                raise ValueError(f"Required CBF column {cbf} has no source mapping.")
            continue  # skip non-required columns with no source
        
        mapping.append((source, cbf))
        if is_required:
            required.add(cbf)
    
    return mapping, required


def load_csv_mapping(path: Path) -> Tuple[List[Tuple[str, str]], Set[str]]:
    """Load CBF mapping from CSV file."""
    mapping = []
    required = set()
    
    with open(path) as f:
        reader = csv.DictReader(f)
        for row in reader:
            source = row.get("source", "").strip()
            cbf = row.get("cbf", "").strip()
            raw_required = row.get("required")
            is_required = (raw_required or "false").strip().lower() == "true"
            
            if not source:
                if is_required:
                    raise ValueError(f"Required CBF column {cbf} has no source mapping.")
                continue
            
            mapping.append((source, cbf))
            if is_required:
                required.add(cbf)
    
    return mapping, required


def load_mapping_file(map_file: Optional[Path] = None) -> Tuple[List[Tuple[str, str]], Set[str]]:
    """
    Load CBF mapping configuration from default or user-provided file.
    
    Args:
        map_file: Optional path to user-provided mapping file (CSV or YAML)
        
    Returns:
        Tuple of (mapping_list, required_columns_set)
    """
    if map_file is None:
        # Load default mapping from S3
        print("[•] Loading default CBF mapping from S3")
        try:
            s3_service = S3DataService()
            content = s3_service.download_cbf_mapping()
            mapping, required = load_yaml_mapping(content)
            print(f"[✓] Loaded {len(mapping)} column mappings from S3 ({len(required)} required)")
        except Exception as e:
            print(f"[!] Failed to load S3 mapping, using fallback: {e}")
            # Fallback to minimal CBF mapping if S3 is unavailable
            # Using actual CBF column names from the standard CloudZero format
            mapping = [
                ("timestamp", "time/usage_start"),
                ("resource_id", "resource/id"),
                ("cost", "cost/cost"),
                ("cost", "cost/discounted_cost"),
                ("cost", "cost/amortized_cost"),
                ("cost", "cost/discounted_amortized_cost"),
                ("cost", "cost/on_demand_cost"),
                ("usage_amount", "usage/amount"),
                ("pricing_unit", "usage/units"),
            ]
            # Only the three essential CBF columns are required
            required = {"time/usage_start", "resource/id", "cost/cost"}
            print(f"[✓] Using fallback mapping with {len(mapping)} columns ({len(required)} required)")
    else:
        # Load user-provided mapping file
        map_path = Path(map_file)
        if not map_path.exists():
            raise FileNotFoundError(f"Mapping file not found: {map_path}")
        
        if map_path.suffix not in {".csv", ".yaml", ".yml"}:
            raise ValueError(f"Unsupported mapping file extension: {map_path.suffix}")
        
        print(f"[•] Using user-provided mapping: {map_path}")
        if map_path.suffix == ".csv":
            mapping, required = load_csv_mapping(map_path)
        else:
            with open(map_path) as f:
                content = f.read()
            mapping, required = load_yaml_mapping(content)
        
        print(f"[✓] Loaded {len(mapping)} column mappings ({len(required)} required)")
    
    return mapping, required


def apply_mapping_and_validate(
    df: pl.DataFrame, 
    mapping: List[Tuple[str, str]], 
    required_cbf_columns: Set[str]
) -> pl.DataFrame:
    """
    Apply CBF mapping to DataFrame and validate required columns.
    
    Follows the cz-demo-tools logic:
    1. Auto-detect tag columns and add them to mapping
    2. Drop columns not in the mapping
    3. Create multiple CBF columns from the same source column
    4. Validate required columns are present and not null
    
    Args:
        df: Input DataFrame to map
        mapping: List of (source_column, cbf_column) tuples
        required_cbf_columns: Set of required CBF column names
        
    Returns:
        DataFrame with CBF column names
        
    Raises:
        ValueError: If required columns are missing or contain nulls
    """
    # Step 1: Auto-detect tag columns and add them to mapping (similar to cz-demo-tools)
    extended_mapping = list(mapping)  # Copy original mapping
    tag_columns = [col for col in df.columns if col.startswith("resource/tag:")]
    for tag_col in tag_columns:
        # Add identity mapping for tag columns
        extended_mapping.append((tag_col, tag_col))
    
    # Step 2: Remove columns not in the mapping (similar to remove_columns_based_on_mapping)
    relevant_columns = {source_col for source_col, _ in extended_mapping if source_col}  # Skip empty source columns
    columns_to_remove = [col for col in df.columns if col not in relevant_columns]
    
    # Drop unmapped columns
    mapped_df = df.drop(columns_to_remove)
    
    # Step 3: Apply column mapping (similar to rename_columns)
    # Track which source columns map to which CBF columns
    source_to_cbf = {}
    for source_col, cbf_col in extended_mapping:
        if not source_col:  # Skip empty source columns
            continue
        if source_col not in source_to_cbf:
            source_to_cbf[source_col] = []
        source_to_cbf[source_col].append(cbf_col)
    
    # Create CBF columns from source columns
    for source_col, cbf_columns in source_to_cbf.items():
        if source_col in mapped_df.columns:
            for cbf_col in cbf_columns:
                # Skip if source and CBF column are the same (identity mapping)
                if source_col == cbf_col:
                    continue
                # Create new CBF column by copying source column
                if cbf_col not in mapped_df.columns:
                    mapped_df = mapped_df.with_columns(pl.col(source_col).alias(cbf_col))
    
    # Step 4: Remove source columns that have been mapped to CBF columns
    for source_col, cbf_columns in source_to_cbf.items():
        if source_col in mapped_df.columns:
            # Only remove source column if it has been mapped to different CBF columns
            if any(cbf_col != source_col for cbf_col in cbf_columns):
                mapped_df = mapped_df.drop(source_col)
    
    # Step 5: Convert all columns to strings with proper null handling for CloudZero
    mapped_df = _convert_to_cloudz_string_format(mapped_df)
    
    # Step 6: Validate required columns are present
    missing_required = required_cbf_columns - set(mapped_df.columns)
    if missing_required:
        raise ValueError(f"Missing required CBF columns: {missing_required}")
    
    # Step 7: Validate required columns don't contain nulls
    for col in required_cbf_columns:
        if col in mapped_df.columns:
            null_count = mapped_df.select(pl.col(col).is_null().sum()).item()
            if null_count > 0:
                raise ValueError(f"Required CBF column '{col}' contains {null_count} null values")
    
    return mapped_df


def _convert_to_cloudz_string_format(df: pl.DataFrame) -> pl.DataFrame:
    """
    Convert all columns to strings with proper null handling for CloudZero ingest.
    
    CloudZero requires:
    - All fields must be strings
    - Cost/usage amount nulls become "0"
    - All other nulls become empty strings ""
    
    Args:
        df: DataFrame to convert
        
    Returns:
        DataFrame with all columns as strings and proper null handling
    """
    # Define columns that should use "0" for null values (cost and usage amount columns)
    numeric_null_to_zero_columns = {
        'cost/cost',
        'cost/amortized_cost', 
        'cost/discounted_cost',
        'cost/discounted_amortized_cost',
        'cost/on_demand_cost',
        'usage/amount'
    }
    
    # Convert each column to string with appropriate null handling
    string_conversions = []
    
    for col in df.columns:
        if col in numeric_null_to_zero_columns:
            # Cost/usage columns: null -> "0", numbers -> stringified
            conversion = (pl.when(pl.col(col).is_null())
                         .then(pl.lit("0"))
                         .otherwise(pl.col(col).cast(pl.Utf8))
                         .alias(col))
        else:
            # All other columns: null -> "", values -> stringified
            conversion = (pl.when(pl.col(col).is_null())
                         .then(pl.lit(""))
                         .otherwise(pl.col(col).cast(pl.Utf8))
                         .alias(col))
        
        string_conversions.append(conversion)
    
    # Apply all conversions
    result_df = df.with_columns(string_conversions)
    
    return result_df