"""Pre-write validation for usage allocation data quality.

This module implements the streamlined validation requirements:
- Row count preservation
- Required columns added  
- Usage calculations validity
- Pricing unit consistency
"""

from dataclasses import dataclass
from typing import List, Optional
import polars as pl
from rich.console import Console

from src.log_context.core import get_logger

console = Console()
logger = get_logger(__name__)


@dataclass
class ValidationCheck:
    """Result of a single validation check."""
    passed: bool
    error: Optional[str] = None
    warning: Optional[str] = None


@dataclass
class ValidationResult:
    """Result of validation process."""
    passed: bool
    should_skip: bool = False  # New field to indicate month should be skipped
    error: Optional[str] = None
    failed_checks: List[ValidationCheck] = None
    

def validate_row_count_preserved(original: pl.DataFrame, enhanced: pl.DataFrame) -> ValidationCheck:
    """Ensure no rows lost during processing.
    
    Args:
        original: Original billing data
        enhanced: Enhanced data with usage columns
        
    Returns:
        ValidationCheck result
    """
    if len(original) != len(enhanced):
        return ValidationCheck(
            passed=False,
            error=f"Row count mismatch: {len(original)} -> {len(enhanced)}"
        )
    return ValidationCheck(passed=True)


def validate_required_columns_added(enhanced: pl.DataFrame) -> ValidationCheck:
    """Ensure usage columns were added.
    
    Args:
        enhanced: Enhanced data that should have usage columns
        
    Returns:
        ValidationCheck result
    """
    required_cols = ["pricing_unit", "usage_amount"]
    missing = [col for col in required_cols if col not in enhanced.columns]
    if missing:
        return ValidationCheck(
            passed=False,
            error=f"Missing required columns: {missing}"
        )
    return ValidationCheck(passed=True)


def validate_usage_calculations(enhanced: pl.DataFrame) -> ValidationCheck:
    """Validate usage calculation requirements.
    
    Requirements:
    1. At least one record has usage_amount > "0.0"
    2. All usage_amount values are valid numeric strings
    3. No negative usage amounts (calculation error indicator)
    
    Args:
        enhanced: Enhanced data with usage calculations
        
    Returns:
        ValidationCheck result
    """
    
    # Check 1: At least one record has usage_amount > "0.0"
    non_zero_usage = enhanced.filter(pl.col("usage_amount") != "0.0")
    if len(non_zero_usage) == 0:
        return ValidationCheck(
            passed=False,
            error="No records have usage_amount > 0 - all calculations failed"
        )
    
    # Check 2: All usage_amount values are valid numeric strings
    try:
        # Convert all usage_amount to float to validate format
        enhanced.select(pl.col("usage_amount").cast(pl.Float64))
    except Exception as e:
        return ValidationCheck(
            passed=False,
            error=f"Invalid usage_amount format (should be numeric strings): {e}"
        )
    
    # Check 3: No negative usage amounts (calculation error indicator)
    try:
        negative_usage = enhanced.filter(
            (pl.col("usage_amount") != "0.0") & 
            (pl.col("usage_amount").cast(pl.Float64) < 0)
        )
        if len(negative_usage) > 0:
            return ValidationCheck(
                passed=False,
                error=f"Found {len(negative_usage)} negative usage amounts (calculation error)"
            )
    except Exception as e:
        return ValidationCheck(
            passed=False,
            error=f"Error checking for negative usage amounts: {e}"
        )
    
    return ValidationCheck(passed=True)


def validate_pricing_unit_consistency(enhanced: pl.DataFrame) -> ValidationCheck:
    """Validate pricing_unit consistency with usage_amount.
    
    Key requirement: If usage_amount > "0.0", then pricing_unit should not be ""
    
    Args:
        enhanced: Enhanced data with usage and pricing columns
        
    Returns:
        ValidationCheck result
    """
    
    try:
        # Key requirement: If usage_amount > "0.0", then pricing_unit should not be ""
        invalid_records = enhanced.filter(
            (pl.col("usage_amount") != "0.0") & 
            (pl.col("usage_amount").cast(pl.Float64) > 0) &
            (pl.col("pricing_unit") == "")
        )
        
        if len(invalid_records) > 0:
            return ValidationCheck(
                passed=False,
                error=f"Found {len(invalid_records)} records with usage > 0 but empty pricing_unit"
            )
    except Exception as e:
        return ValidationCheck(
            passed=False,
            error=f"Error validating pricing unit consistency: {e}"
        )
    
    return ValidationCheck(passed=True)


def validate_before_write(
    original_data: pl.DataFrame,
    enhanced_data: pl.DataFrame, 
    provider: str,
    month: str
) -> ValidationResult:
    """Comprehensive validation before S3 write.
    
    Executes all validation checks and determines if processing should continue,
    skip, or fail.
    
    Args:
        original_data: Original billing data
        enhanced_data: Enhanced data with usage columns
        provider: Cloud provider
        month: Month being processed
        
    Returns:
        ValidationResult with overall success status and skip flag
    """
    
    # Check for data integrity first
    integrity_checks = [
        validate_row_count_preserved(original_data, enhanced_data),
        validate_required_columns_added(enhanced_data),
    ]
    
    # Check for critical integrity failures
    critical_failures = [check for check in integrity_checks if not check.passed]
    if critical_failures:
        error_messages = [check.error for check in critical_failures if check.error]
        combined_error = "; ".join(error_messages)
        
        logger.error(
            "Critical validation failed",
            provider=provider,
            month=month,
            failed_checks=len(critical_failures),
            errors=error_messages
        )
        
        return ValidationResult(
            passed=False,
            error=f"Critical validation failed: {combined_error}",
            failed_checks=critical_failures
        )
    
    # Check for usage calculation results
    usage_check = validate_usage_calculations(enhanced_data)
    
    # If no usage records found, skip instead of failing
    if not usage_check.passed and "No records have usage_amount > 0" in usage_check.error:
        logger.info(
            "No usage matches found - skipping month",
            provider=provider,
            month=month
        )
        
        return ValidationResult(
            passed=True,
            should_skip=True,
            error="No usage matches found for this month"
        )
    
    # Execute remaining checks
    remaining_checks = [usage_check, validate_pricing_unit_consistency(enhanced_data)]
    
    failed_checks = [check for check in remaining_checks if not check.passed]
    if failed_checks:
        error_messages = [check.error for check in failed_checks if check.error]
        combined_error = "; ".join(error_messages)
        
        logger.error(
            "Usage validation failed",
            provider=provider,
            month=month,
            failed_checks=len(failed_checks),
            errors=error_messages
        )
        
        return ValidationResult(
            passed=False,
            error=f"Usage validation failed: {combined_error}",
            failed_checks=failed_checks
        )
    
    logger.info(
        "Pre-write validation passed",
        provider=provider,
        month=month,
        checks_passed=len(integrity_checks) + len(remaining_checks)
    )
    
    return ValidationResult(passed=True)