"""
Job splitting logic for memory-bounded generation.

Handles splitting large generation jobs into smaller chunks based on date ranges
to fit within memory bounds while maintaining monthly parquet file structure.
"""

from datetime import datetime, timedelta, date
from typing import List, Tuple, Dict, Any
from dateutil.relativedelta import relativedelta
import polars as pl

from ..config.models import GenerationJob, Filter


def calculate_hours_in_date_range(start_date: datetime, end_date: datetime) -> int:
    """
    Calculate the number of hours in a date range.
    
    For billing data generation, the end_date should include all hours of that date.
    For example, if end_date is 2025-12-31, we should include hours 00:00 through 23:00
    of 2025-12-31, but not 00:00 of 2026-01-01.
    
    Args:
        start_date: Start date (includes this hour)
        end_date: End date (includes all hours of this date)
        
    Returns:
        Number of hours in the date range
    """
    # If end_date has no time component (00:00:00), add 23 hours to include the full day
    if end_date.hour == 0 and end_date.minute == 0 and end_date.second == 0:
        # Include all hours of the end date (00:00 through 23:00)
        adjusted_end_date = end_date.replace(hour=23, minute=59, second=59)
    else:
        adjusted_end_date = end_date
    
    delta = adjusted_end_date - start_date
    return int(delta.total_seconds() / 3600) + 1  # +1 to include the final hour


def estimate_job_memory_usage(
    job: GenerationJob,
    reference_matches: int,
    bytes_per_record: int = 151
) -> int:
    """
    Estimate memory usage for a generation job.
    
    Args:
        job: Generation job to estimate
        reference_matches: Number of reference data matches after filtering
        bytes_per_record: Average bytes per record in Polars DataFrame
        
    Returns:
        Estimated memory usage in bytes
    """
    hours_in_range = calculate_hours_in_date_range(job.start_date, job.end_date)
    total_records = reference_matches * job.resource_count * hours_in_range
    return total_records * bytes_per_record


def split_job_by_date_range(
    job: GenerationJob,
    target_memory_bytes: int = 1024 * 1024 * 1024,  # 1GB default
    reference_matches: int = 1,
    bytes_per_record: int = 151
) -> List[GenerationJob]:
    """
    Split a generation job into smaller date-based chunks.
    
    Uses monthly splits to align with parquet output structure.
    
    Args:
        job: Original generation job
        target_memory_bytes: Target memory limit per job
        reference_matches: Number of reference data matches after filtering
        bytes_per_record: Average bytes per record in Polars DataFrame
        
    Returns:
        List of smaller generation jobs
    """
    # Check if splitting is needed
    estimated_memory = estimate_job_memory_usage(job, reference_matches, bytes_per_record)
    if estimated_memory <= target_memory_bytes:
        return [job]
    
    # Calculate how many months we can fit in target memory
    total_hours = calculate_hours_in_date_range(job.start_date, job.end_date)
    records_per_hour = reference_matches * job.resource_count
    memory_per_hour = records_per_hour * bytes_per_record
    
    max_hours_per_job = target_memory_bytes // memory_per_hour
    
    # Convert to months (assume 30 days per month for estimation)
    hours_per_month = 30 * 24
    max_months_per_job = max(1, max_hours_per_job // hours_per_month)
    
    # Split into date ranges
    split_jobs = []
    current_date = job.start_date
    job_index = 1
    
    while current_date < job.end_date:
        # Calculate end date for this split (monthly boundaries)
        split_end_date = min(
            current_date + relativedelta(months=max_months_per_job),
            job.end_date
        )
        
        # Create split job, preserving original UUID for cross-month consistency
        split_job = GenerationJob(
            name=f"{job.name}_split_{job_index}",
            reference_data=job.reference_data,
            start_date=current_date,
            end_date=split_end_date,
            total_cost=job.total_cost * (
                calculate_hours_in_date_range(current_date, split_end_date) / total_hours
            ),
            resource_count=job.resource_count,
            filters=job.filters,
            uuid=job.uuid  # Preserve UUID for consistent resource IDs across months
        )
        
        split_jobs.append(split_job)
        current_date = split_end_date
        job_index += 1
    
    return split_jobs


def group_jobs_by_month(jobs: List[GenerationJob]) -> Dict[str, List[GenerationJob]]:
    """
    Group generation jobs by month for parquet file organization.
    
    Args:
        jobs: List of generation jobs
        
    Returns:
        Dictionary mapping month strings (YYYY-MM) to lists of jobs
    """
    month_groups = {}
    
    for job in jobs:
        current_date = job.start_date
        
        while current_date < job.end_date:
            month_key = current_date.strftime("%Y-%m")
            
            if month_key not in month_groups:
                month_groups[month_key] = []
            
            month_groups[month_key].append(job)
            current_date = current_date + relativedelta(months=1)
    
    return month_groups


def validate_job_splits(
    original_job: GenerationJob,
    split_jobs: List[GenerationJob],
    tolerance: float = 0.01
) -> bool:
    """
    Validate that split jobs preserve the original job's properties.
    
    Args:
        original_job: Original generation job
        split_jobs: List of split jobs
        tolerance: Acceptable difference for total cost validation
        
    Returns:
        True if validation passes, False otherwise
    """
    # Check total cost preservation
    original_cost = original_job.total_cost
    split_cost_sum = sum(job.total_cost for job in split_jobs)
    
    if abs(original_cost - split_cost_sum) > original_cost * tolerance:
        return False
    
    # Check date range coverage
    split_start = min(job.start_date for job in split_jobs)
    split_end = max(job.end_date for job in split_jobs)
    
    if split_start != original_job.start_date or split_end != original_job.end_date:
        return False
    
    # Check resource count preservation
    for job in split_jobs:
        if job.resource_count != original_job.resource_count:
            return False
    
    # Check filters preservation
    for job in split_jobs:
        if len(job.filters) != len(original_job.filters):
            return False
    
    return True


class JobSplitter:
    """
    Handles job splitting for memory-bounded generation with intelligent date boundaries.
    """
    
    def __init__(
        self,
        target_memory_bytes: int = 1024 * 1024 * 1024,  # 1GB default
        bytes_per_record: int = 151,
        prefer_month_boundaries: bool = True
    ):
        self.target_memory_bytes = target_memory_bytes
        self.bytes_per_record = bytes_per_record
        self.prefer_month_boundaries = prefer_month_boundaries
    
    def split_jobs_for_memory_bounds(
        self,
        jobs: List[GenerationJob],
        reference_matches_per_job: Dict[str, int]
    ) -> List[GenerationJob]:
        """
        Split multiple jobs to fit within memory bounds using intelligent date boundaries.
        
        Args:
            jobs: List of generation jobs
            reference_matches_per_job: Dictionary mapping job names to reference match counts
            
        Returns:
            List of split jobs that fit within memory bounds
        """
        split_jobs = []
        
        for job in jobs:
            reference_matches = reference_matches_per_job.get(job.name, 1)
            
            # Use intelligent splitting with month boundaries when possible
            if self.prefer_month_boundaries:
                job_splits = self._split_job_with_intelligent_boundaries(
                    job, reference_matches
                )
            else:
                job_splits = split_job_by_date_range(
                    job,
                    self.target_memory_bytes,
                    reference_matches,
                    self.bytes_per_record
                )
            
            # Validate splits
            if not validate_job_splits(job, job_splits):
                raise ValueError(f"Job split validation failed for job: {job.name}")
            
            split_jobs.extend(job_splits)
        
        return split_jobs
    
    def _split_job_with_intelligent_boundaries(
        self,
        job: GenerationJob,
        reference_matches: int
    ) -> List[GenerationJob]:
        """
        Split job using intelligent date boundaries (month boundaries when possible).
        
        This approach prioritizes data locality and upload efficiency by aligning
        splits with natural date boundaries that match our S3 storage structure.
        
        Args:
            job: Job to split
            reference_matches: Number of reference data matches for this job
            
        Returns:
            List of split jobs
        """
        # First check if the job fits within memory bounds without splitting
        estimated_memory = estimate_job_memory_usage(
            job, reference_matches, self.bytes_per_record
        )
        
        if estimated_memory <= self.target_memory_bytes:
            return [job]  # No splitting needed
        
        split_jobs = []
        current_date = job.start_date
        
        while current_date < job.end_date:
            # Try month boundary first
            month_end = (current_date + relativedelta(months=1)).replace(day=1) - timedelta(days=1)
            
            # Don't go beyond job end date
            if month_end > job.end_date:
                month_end = job.end_date
            
            # Estimate memory for this month
            hours_in_period = calculate_hours_in_date_range(current_date, month_end)
            period_memory = (
                reference_matches * job.resource_count * hours_in_period * self.bytes_per_record
            )
            
            if period_memory <= self.target_memory_bytes or current_date == month_end:
                # Month fits within memory bounds or is a single day
                split_job = GenerationJob(
                    name=f"{job.name}_split_{current_date.strftime('%Y%m')}",
                    start_date=current_date,
                    end_date=month_end,
                    total_cost=(
                        job.total_cost * hours_in_period /
                        calculate_hours_in_date_range(job.start_date, job.end_date)
                    ),
                    resource_count=job.resource_count,
                    filters=job.filters,
                    uuid=job.uuid  # Preserve UUID for consistent resource IDs across months
                )
                split_jobs.append(split_job)
                current_date = month_end + timedelta(days=1)
            else:
                # Month is too large, fall back to day-by-day splitting
                day_splits = self._split_by_days(
                    job, current_date, month_end, reference_matches
                )
                split_jobs.extend(day_splits)
                current_date = month_end + timedelta(days=1)
        
        return split_jobs
    
    def _split_by_days(
        self,
        job: GenerationJob,
        start_date: date,
        end_date: date,
        reference_matches: int
    ) -> List[GenerationJob]:
        """
        Split a date range by days when month boundary doesn't fit in memory.
        
        Args:
            job: Original job
            start_date: Start of range to split
            end_date: End of range to split
            reference_matches: Number of reference matches
            
        Returns:
            List of daily split jobs
        """
        split_jobs = []
        current_date = start_date
        split_index = 0
        
        while current_date <= end_date:
            # Calculate how many days we can fit in memory
            max_days = 1
            test_end_date = current_date

            while test_end_date < end_date:
                next_end_date = min(test_end_date + timedelta(days=1), end_date)
                hours_in_period = calculate_hours_in_date_range(current_date, next_end_date)
                period_memory = (
                    reference_matches * job.resource_count * hours_in_period * self.bytes_per_record
                )

                if period_memory <= self.target_memory_bytes:
                    test_end_date = next_end_date
                    max_days = (test_end_date - current_date).days + 1
                else:
                    break

            # Ensure max_days is at least 1 to prevent invalid date ranges
            max_days = max(1, max_days)

            # Create split job for the calculated period
            split_end_date = min(current_date + timedelta(days=max_days - 1), end_date)

            # Defensive check: ensure split_end_date is not before current_date
            if split_end_date < current_date:
                split_end_date = current_date
            hours_in_period = calculate_hours_in_date_range(current_date, split_end_date)

            split_job = GenerationJob(
                name=f"{job.name}_split_{split_index:03d}",
                start_date=current_date,
                end_date=split_end_date,
                total_cost=(
                    job.total_cost * hours_in_period /
                    calculate_hours_in_date_range(job.start_date, job.end_date)
                ),
                resource_count=job.resource_count,
                filters=job.filters,
                uuid=job.uuid  # Preserve UUID for consistent resource IDs across months
            )
            
            split_jobs.append(split_job)
            current_date = split_end_date + timedelta(days=1)
            split_index += 1
        
        return split_jobs
    
    def estimate_total_memory_usage(
        self,
        jobs: List[GenerationJob],
        reference_matches_per_job: Dict[str, int]
    ) -> int:
        """
        Estimate total memory usage for a list of jobs.
        
        Args:
            jobs: List of generation jobs
            reference_matches_per_job: Dictionary mapping job names to reference match counts
            
        Returns:
            Total estimated memory usage in bytes
        """
        total_memory = 0
        
        for job in jobs:
            reference_matches = reference_matches_per_job.get(job.name, 1)
            job_memory = estimate_job_memory_usage(job, reference_matches, self.bytes_per_record)
            total_memory += job_memory
        
        return total_memory