"""
Batch processor - handles folders of CSV files.
"""

from pathlib import Path
from typing import List
from dataclasses import dataclass
import pandas as pd
from concurrent.futures import ThreadPoolExecutor, as_completed

from app.core.logging import get_logger
from app.autoclean.issue_detector import IssueDetector, DataIssue
from app.autoclean.auto_fixer import AutoFixer, FixResult
from app.storage.file_storage import FileStorage

logger = get_logger(__name__)


@dataclass
class FileResult:
    """Result of processing a single file."""
    
    filename: str
    success: bool
    issues_found: List[DataIssue]
    fixes_applied: List[FixResult]
    original_rows: int
    final_rows: int
    error: str | None = None


@dataclass
class BatchResult:
    """Result of batch processing."""
    
    total_files: int
    successful_files: int
    failed_files: int
    file_results: List[FileResult]
    merged_df: pd.DataFrame | None
    total_issues_found: int
    total_fixes_applied: int


class BatchProcessor:
    """Process folders of messy CSV files."""
    
    def __init__(self, conservative: bool = True, max_workers: int = 4):
        """
        Initialize batch processor.
        
        Args:
            conservative: Use conservative fixing strategy
            max_workers: Number of parallel workers
        """
        self.detector = IssueDetector()
        self.fixer = AutoFixer(conservative=conservative)
        self.storage = FileStorage()
        self.max_workers = max_workers
        self.logger = logger
    
    def process_folder(
        self,
        input_folder: Path,
        output_folder: Path,
        merge: bool = True,
        flag_only: bool = False,
    ) -> BatchResult:
        """
        Process all CSV files in a folder.
        
        Args:
            input_folder: Folder containing messy CSVs
            output_folder: Folder to save cleaned files
            merge: Whether to merge all files into one
            flag_only: If True, only detect issues without fixing (conservative mode)
            
        Returns:
            BatchResult with processing details
        """
        input_folder = Path(input_folder)
        output_folder = Path(output_folder)
        output_folder.mkdir(parents=True, exist_ok=True)
        
        # Find all CSV files
        csv_files = list(input_folder.glob("*.csv"))
        self.logger.info(f"Found {len(csv_files)} CSV files")
        
        if not csv_files:
            return BatchResult(
                total_files=0,
                successful_files=0,
                failed_files=0,
                file_results=[],
                merged_df=None,
                total_issues_found=0,
                total_fixes_applied=0,
            )
        
        # Process files in parallel
        file_results = []
        cleaned_dfs = []
        
        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            future_to_file = {
                executor.submit(self._process_single_file, file_path, flag_only): file_path
                for file_path in csv_files
            }
            
            for future in as_completed(future_to_file):
                file_path = future_to_file[future]
                try:
                    result, cleaned_df = future.result()
                    file_results.append(result)
                    
                    if result.success and cleaned_df is not None:
                        cleaned_dfs.append(cleaned_df)
                        
                        # Save individual cleaned file
                        output_path = output_folder / f"cleaned_{file_path.name}"
                        cleaned_df.to_csv(output_path, index=False)
                        
                except Exception as e:
                    self.logger.error(f"Failed to process {file_path}: {e}")
                    file_results.append(FileResult(
                        filename=file_path.name,
                        success=False,
                        issues_found=[],
                        fixes_applied=[],
                        original_rows=0,
                        final_rows=0,
                        error=str(e),
                    ))
        
        # Merge if requested
        merged_df = None
        if merge and cleaned_dfs:
            merged_df = self._merge_dataframes(cleaned_dfs)
            
            if merged_df is not None:
                merged_path = output_folder / "merged_clean.csv"
                merged_df.to_csv(merged_path, index=False)
                self.logger.info(f"Saved merged file: {merged_path}")
        
        # Calculate summary
        successful = sum(1 for r in file_results if r.success)
        total_issues = sum(len(r.issues_found) for r in file_results)
        total_fixes = sum(len(r.fixes_applied) for r in file_results)
        
        return BatchResult(
            total_files=len(csv_files),
            successful_files=successful,
            failed_files=len(csv_files) - successful,
            file_results=file_results,
            merged_df=merged_df,
            total_issues_found=total_issues,
            total_fixes_applied=total_fixes,
        )
    
    def _process_single_file(self, file_path: Path, flag_only: bool = False) -> tuple[FileResult, pd.DataFrame | None]:
        """Process a single CSV file."""
        try:
            # Read file
            df = pd.read_csv(file_path, encoding='utf-8', on_bad_lines='skip')
            original_rows = len(df)
            
            # Detect issues
            issues = self.detector.detect_all(df)
            
            # Apply fixes (unless flag_only mode)
            fixes_applied = []
            if not flag_only:
                for issue in issues:
                    df, fix_result = self.fixer.fix_issue(df, issue)
                    if fix_result.success:
                        fixes_applied.append(fix_result)
            
            final_rows = len(df)
            
            result = FileResult(
                filename=file_path.name,
                success=True,
                issues_found=issues,
                fixes_applied=fixes_applied,
                original_rows=original_rows,
                final_rows=final_rows,
            )
            
            # In flag_only mode, return original df (no modifications)
            return_df = df if not flag_only else pd.read_csv(file_path, encoding='utf-8', on_bad_lines='skip')
            return result, return_df
            
        except Exception as e:
            self.logger.error(f"Error processing {file_path}: {e}", exc_info=True)
            return FileResult(
                filename=file_path.name,
                success=False,
                issues_found=[],
                fixes_applied=[],
                original_rows=0,
                final_rows=0,
                error=str(e),
            ), None
    
    def _merge_dataframes(self, dfs: List[pd.DataFrame]) -> pd.DataFrame | None:
        """
        Merge multiple DataFrames intelligently.
        
        Strategy:
        1. If all have same columns → simple concat
        2. If different columns → merge by column names (outer join)
        """
        if not dfs:
            return None
        
        if len(dfs) == 1:
            return dfs[0]
        
        try:
            # Check if all have same columns
            first_cols = set(dfs[0].columns)
            all_same = all(set(df.columns) == first_cols for df in dfs)
            
            if all_same:
                # Simple concatenation
                merged = pd.concat(dfs, ignore_index=True)
                self.logger.info("Merged with simple concatenation")
            else:
                # Merge by aligning columns (outer join)
                all_columns = set()
                for df in dfs:
                    all_columns.update(df.columns)
                
                # Reindex all DataFrames to have same columns
                aligned_dfs = []
                for df in dfs:
                    aligned = df.reindex(columns=all_columns)
                    aligned_dfs.append(aligned)
                
                merged = pd.concat(aligned_dfs, ignore_index=True)
                self.logger.info(f"Merged with column alignment ({len(all_columns)} total columns)")
            
            return merged
            
        except Exception as e:
            self.logger.error(f"Merge failed: {e}", exc_info=True)
            return None
