"""
Dataset loading utilities for upload operations.

Provides loading capabilities for both generated datasets from S3 storage
and ad-hoc CSV datasets from user input.
"""

import re
from pathlib import Path
from typing import Tuple, Optional, List, Dict

import polars as pl

from src.services.s3_data_service import S3DataService


class DatasetLoader:
    """
    Dataset loader for upload operations.
    
    Handles loading generated datasets from S3 storage and
    ad-hoc datasets from user-provided CSV files.
    """
    
    def __init__(self):
        """Initialize dataset loader with S3 support."""
        self.s3_loader = S3DataService()
    
    def load_generated_dataset_from_s3(
        self,
        dataset_name: str,
        stage: str = "generate",
        months: Optional[List[str]] = None,
        providers: Optional[List[str]] = None
    ) -> Dict[str, Dict[str, pl.DataFrame]]:
        """
        Load generated dataset from S3 with month/provider filtering.
        
        Args:
            dataset_name: Name of the dataset to load
            stage: Stage to load from (default: "generate")
            months: List of months to load (format: "YYYY-MM")
            providers: List of providers to load (e.g., ["AWS", "GCP"])
            
        Returns:
            Dict[provider, Dict[month, DataFrame]] - Nested structure for easy processing
        """
        return self.s3_loader.load_dataset_for_upload(
            dataset_name=dataset_name,
            stage=stage,
            months=months,
            providers=providers
        )
    
    def get_s3_dataset_summary(self, dataset_name: str, stage: str = "generate") -> Dict[str, Dict[str, int]]:
        """
        Get summary of S3 dataset contents without loading data.
        
        Args:
            dataset_name: Name of the dataset
            stage: Stage to summarize (default: "generate")
            
        Returns:
            Dict[provider, Dict[month, line_count]] - Summary of available data
        """
        return self.s3_loader.get_upload_dataset_summary(dataset_name, stage)
    
    def get_s3_dataset_summary_fast(self, dataset_name: str, stage: str = "generate") -> Dict[str, List[str]]:
        """
        Get fast summary of S3 dataset contents without loading data or counting rows.
        
        Only checks file existence for validation purposes.
        
        Args:
            dataset_name: Name of the dataset
            stage: Stage to summarize (default: "generate")
            
        Returns:
            Dict[provider, List[month]] - Summary of available files without row counts
        """
        return self.s3_loader.get_upload_dataset_summary_fast(dataset_name, stage)
    
    def validate_s3_dataset_access(self, dataset_name: str, stage: str = "generate") -> bool:
        """
        Validate that we can access the S3 dataset without loading it.
        
        Args:
            dataset_name: Name of the dataset
            stage: Stage to validate (default: "generate")
            
        Returns:
            True if dataset is accessible, False otherwise
        """
        return self.s3_loader.validate_upload_dataset_access(dataset_name, stage)
    
    def load_generated_dataset(
        self, 
        org: str, 
        run_id: Optional[str] = None, 
        stage: str = "generate"
    ) -> Tuple[pl.DataFrame, Path]:
        """
        Load generated dataset from pipeline run (legacy local file support).
        
        Extracted from loader.load_pipeline_data_with_config() to provide
        consistent loading of generated datasets for S3 upload.
        """
        # Use relative path from workspace root
        base = Path("orgs") / org / "billing" / "runs"
        if not base.exists():
            raise FileNotFoundError(f"No billing runs directory found at {base}")
        
        # Resolve run directory
        if run_id == "latest" or run_id is None:
            # Get most recent run
            run_dirs = sorted(base.glob("*"), key=lambda p: p.name, reverse=True)
            if not run_dirs:
                raise FileNotFoundError(f"No run directories found in {base}")
            run_dir = run_dirs[0]
        else:
            # Use specific run ID
            run_dir = base / run_id
            if not run_dir.exists():
                raise FileNotFoundError(f"Run directory not found: {run_dir}")
        
        # Load from specified stage
        stage_dir = run_dir / stage
        if not stage_dir.exists():
            raise FileNotFoundError(f"Stage directory not found: {stage_dir}")
        
        # Look for consolidated billing data first
        billing_path = run_dir / "billing.parquet"
        if billing_path.exists():
            print(f"[•] Loading consolidated billing data from {billing_path}")
            df = pl.read_parquet(billing_path)
            return df, billing_path.parent
        
        # Look for stage-specific billing data
        stage_billing_path = stage_dir / "billing.parquet"
        if stage_billing_path.exists():
            print(f"[•] Loading stage billing data from {stage_billing_path}")
            df = pl.read_parquet(stage_billing_path)
            return df, stage_billing_path.parent
        
        raise FileNotFoundError(
            f"No billing.parquet found in {run_dir} or {stage_dir}"
        )
    
    def load_adhoc_dataset(self, csv_path: Path) -> pl.DataFrame:
        """
        Load ad-hoc dataset from user-provided CSV file.
        
        Extracted from loader.load_files_from_folder() to provide
        consistent loading of user CSV files for API upload.
        """
        if not csv_path.exists():
            raise FileNotFoundError(f"CSV file not found: {csv_path}")
        
        print(f"[•] Loading user-provided CSV: {csv_path}")
        return pl.read_csv(csv_path)
    
    def load_user_mapping(self, mapping_path: Path) -> List[Tuple[str, str]]:
        """
        Load user-defined column mapping for ad-hoc datasets.
        
        Expected format: CSV with 'source_column' and 'cbf_column' headers.
        """
        if not mapping_path.exists():
            raise FileNotFoundError(f"Mapping file not found: {mapping_path}")
        
        print(f"[•] Loading user mapping from {mapping_path}")
        mapping_df = pl.read_csv(mapping_path)
        
        # Validate mapping file structure
        required_columns = {"source_column", "cbf_column"}
        if not required_columns.issubset(mapping_df.columns):
            raise ValueError(
                f"Mapping file must contain columns: {required_columns}. "
                f"Found: {set(mapping_df.columns)}"
            )
        
        # Convert to list of tuples
        mapping_tuples = [
            (row["source_column"], row["cbf_column"])
            for row in mapping_df.to_dicts()
        ]
        
        return mapping_tuples


def resolve_env_var(value: str) -> str:
    """
    Resolve environment variable references in configuration values.
    
    Extracted from loader.resolve_env_var() to support environment
    variable resolution in upload configurations.
    """
    import os
    
    match = re.fullmatch(r"\${(\w+)}", value)
    if not match:
        return value  # Not a variable reference
    
    var_name = match.group(1)
    resolved = os.getenv(var_name)
    if not resolved:
        raise OSError(f"Environment variable {var_name} is not set")
    
    return resolved