"""Parallel chunk processor for large file analysis."""

import math
import time
import logging
from pathlib import Path
from typing import List, Dict, Any, Optional
from concurrent.futures import ProcessPoolExecutor, as_completed
from tqdm import tqdm
import json

from .file_reader import FileReader
from .cache_manager import CacheManager

logger = logging.getLogger(__name__)


def process_single_chunk(
    file_path: str,
    start: int,
    end: int,
    analyzer_configs: List[Dict],
    s3_category_filter: Optional[List[str]] = None,
) -> Dict[str, Any]:
    """Process a single chunk with all analyzers.

    This function runs in a separate process.

    Args:
        file_path: Path to the JSONL file
        start: Starting line number
        end: Ending line number
        analyzer_configs: List of analyzer configurations
        s3_category_filter: Optional list of keywords to filter records by s3_filepath

    Returns:
        Dictionary with results from each analyzer
    """
    chunk_start_time = time.time()

    # Import analyzers here to avoid pickling issues
    from ..analyzers.id_analyzer import IDAnalyzer
    from ..analyzers.tags_analyzer import TagsAnalyzer
    from ..analyzers.text_analyzer import TextAnalyzer
    from ..analyzers.lang_analyzer import LanguageAnalyzer
    from ..analyzers.weight_analyzer import WeightAnalyzer
    from ..analyzers.stems_analyzer import StemsAnalyzer
    from ..analyzers.paths_analyzer import PathsAnalyzer
    from ..analyzers.lists_analyzer import ListsAnalyzer
    from ..analyzers.sample_collector import SampleCollector

    analyzer_classes = {
        "id": IDAnalyzer,
        "tags": TagsAnalyzer,
        "text": TextAnalyzer,
        "language": LanguageAnalyzer,  # Fixed: was 'lang', should be 'language'
        "weight": WeightAnalyzer,
        "stems": StemsAnalyzer,
        "paths": PathsAnalyzer,
        "lists": ListsAnalyzer,
        "samples": SampleCollector,
    }

    # Initialize analyzers
    analyzers = []
    for config in analyzer_configs:
        analyzer_class = analyzer_classes.get(config["name"])
        if analyzer_class:
            analyzer = analyzer_class(output_dir=config["output_dir"])
            analyzers.append(analyzer)

    # Read chunk data
    read_start = time.time()
    reader = FileReader(Path(file_path))
    records = list(reader.read_records(start=start, end=end, show_progress=False))
    read_time = time.time() - read_start

    # Apply global s3_category_filter if specified
    original_count = len(records)
    filtered_count = 0
    if s3_category_filter:
        filtered_records = []
        for record in records:
            s3_path = record.get("s3_filepath", "")
            if s3_path and any(keyword.lower() in s3_path.lower() for keyword in s3_category_filter):
                filtered_records.append(record)
            else:
                filtered_count += 1
        records = filtered_records

    # Process with each analyzer
    analyzer_times = {}
    results = {
        "chunk_info": {
            "start": start,
            "end": end,
            "records_read": original_count,
            "records_filtered": filtered_count,
            "records_processed": len(records),
        }
    }

    for analyzer in analyzers:
        try:
            analyzer_start = time.time()
            results[analyzer.name] = analyzer.process_chunk(records)
            analyzer_times[analyzer.name] = time.time() - analyzer_start
        except Exception as e:
            logger.error(f"Error in {analyzer.name} for chunk {start}-{end}: {e}")
            results[analyzer.name] = {"error": str(e)}
            analyzer_times[analyzer.name] = 0

    chunk_total_time = time.time() - chunk_start_time

    # Add timing info
    results["chunk_info"]["timing"] = {
        "total_time": chunk_total_time,
        "read_time": read_time,
        "analyzer_times": analyzer_times,
        "processing_time": chunk_total_time - read_time,
    }

    return results


class ChunkProcessor:
    """Manage parallel processing of file chunks."""

    def __init__(
        self,
        file_path: Path,
        chunk_size: int = 2_000_000,
        num_workers: int = 8,
        cache_manager: Optional[CacheManager] = None,
        s3_category_filter: Optional[List[str]] = None,
    ):
        """Initialize chunk processor.

        Args:
            file_path: Path to the JSONL file
            chunk_size: Number of lines per chunk
            num_workers: Number of parallel workers
            cache_manager: Optional cache manager
            s3_category_filter: Optional list of keywords to filter records by s3_filepath
        """
        self.file_path = Path(file_path)
        self.chunk_size = chunk_size
        self.num_workers = num_workers
        self.cache_manager = cache_manager or CacheManager()
        self.s3_category_filter = s3_category_filter

        self.reader = FileReader(self.file_path, self.cache_manager)
        self.total_lines = self.reader.get_line_count()
        self.num_chunks = math.ceil(self.total_lines / chunk_size)

        if s3_category_filter:
            logger.info(
                f"Processing {self.total_lines:,} lines in {self.num_chunks} chunks with s3_category_filter: {s3_category_filter}"
            )
        else:
            logger.info(f"Processing {self.total_lines:,} lines in {self.num_chunks} chunks")

    def process_file(self, analyzers: List[Any], output_dir: Path) -> Dict[str, Any]:
        """Process entire file with all analyzers in parallel.

        Args:
            analyzers: List of analyzer instances
            output_dir: Directory for outputs

        Returns:
            Aggregated results from all analyzers
        """
        overall_start_time = time.time()

        # Prepare analyzer configurations (can't pickle analyzer objects)
        analyzer_configs = [
            {"name": analyzer.name, "output_dir": str(output_dir)} for analyzer in analyzers
        ]

        # Generate chunks
        chunks = list(self.reader.read_chunks(self.chunk_size))

        logger.info(f"Processing {len(chunks)} chunks with {self.num_workers} workers")

        # Process chunks in parallel
        chunk_results = []
        processing_start_time = time.time()
        with ProcessPoolExecutor(max_workers=self.num_workers) as executor:
            # Submit all tasks
            futures = []
            for chunk_start, chunk_end in chunks:
                future = executor.submit(
                    process_single_chunk,
                    str(self.file_path),
                    chunk_start,
                    chunk_end,
                    analyzer_configs,
                    self.s3_category_filter,
                )
                futures.append(future)

            # Collect results with progress bar
            for future in tqdm(as_completed(futures), total=len(futures), desc="Processing chunks"):
                try:
                    result = future.result(timeout=300)  # 5 min timeout per chunk
                    chunk_results.append(result)
                except Exception as e:
                    logger.error(f"Chunk processing failed: {e}")
                    chunk_results.append({"error": str(e)})

        processing_time = time.time() - processing_start_time

        # Sort results by chunk order
        chunk_results.sort(key=lambda x: x.get("chunk_info", {}).get("start", 0))

        # Aggregate results for each analyzer
        logger.info("Aggregating results...")
        aggregation_start_time = time.time()
        aggregated_results = {}

        for analyzer in analyzers:
            analyzer_name = analyzer.name
            analyzer_chunks = []

            for chunk_result in chunk_results:
                if analyzer_name in chunk_result and "error" not in chunk_result[analyzer_name]:
                    analyzer_chunks.append(chunk_result[analyzer_name])

            if analyzer_chunks:
                try:
                    logger.info(f"Aggregating {analyzer_name}... ({len(analyzer_chunks)} chunks)")
                    agg_start = time.time()
                    aggregated_results[analyzer_name] = analyzer.aggregate(analyzer_chunks)
                    agg_time = time.time() - agg_start
                    logger.info(f"  {analyzer_name} aggregation took {agg_time:.2f}s")
                except Exception as e:
                    logger.error(f"Failed to aggregate {analyzer_name}: {e}")
                    aggregated_results[analyzer_name] = {"error": str(e)}
            else:
                aggregated_results[analyzer_name] = {"error": "No valid chunks to aggregate"}

        aggregation_time = time.time() - aggregation_start_time
        overall_time = time.time() - overall_start_time

        # Collect timing statistics from chunks
        chunk_timings = []
        for chunk_result in chunk_results:
            if "chunk_info" in chunk_result and "timing" in chunk_result["chunk_info"]:
                chunk_timings.append(chunk_result["chunk_info"]["timing"])

        # Calculate timing statistics
        if chunk_timings:
            avg_chunk_time = sum(t["total_time"] for t in chunk_timings) / len(chunk_timings)
            max_chunk_time = max(t["total_time"] for t in chunk_timings)
            min_chunk_time = min(t["total_time"] for t in chunk_timings)
            avg_read_time = sum(t["read_time"] for t in chunk_timings) / len(chunk_timings)
        else:
            avg_chunk_time = max_chunk_time = min_chunk_time = avg_read_time = 0

        # Add processing metadata
        aggregated_results["_metadata"] = {
            "file_path": str(self.file_path),
            "total_lines": self.total_lines,
            "num_chunks": len(chunks),
            "chunk_size": self.chunk_size,
            "num_workers": self.num_workers,
            "chunks_processed": len(chunk_results),
            "chunks_with_errors": sum(1 for r in chunk_results if "error" in r),
            "timing": {
                "overall_time": overall_time,
                "processing_time": processing_time,
                "aggregation_time": aggregation_time,
                "avg_chunk_time": avg_chunk_time,
                "max_chunk_time": max_chunk_time,
                "min_chunk_time": min_chunk_time,
                "avg_read_time": avg_read_time,
                "records_per_second": self.total_lines / overall_time if overall_time > 0 else 0,
            },
        }

        # Log timing summary
        logger.info(f"Processing complete in {overall_time:.2f}s")
        logger.info(f"  - Chunk processing: {processing_time:.2f}s")
        logger.info(f"  - Aggregation: {aggregation_time:.2f}s")
        logger.info(
            f"  - Avg chunk time: {avg_chunk_time:.2f}s (min: {min_chunk_time:.2f}s, max: {max_chunk_time:.2f}s)"
        )
        logger.info(f"  - Processing rate: {self.total_lines / overall_time:,.0f} records/sec")

        return aggregated_results

    def process_sample(
        self, analyzers: List[Any], sample_size: int = 100_000, output_dir: Optional[Path] = None
    ) -> Dict[str, Any]:
        """Process a sample of the file for testing.

        Args:
            analyzers: List of analyzer instances
            sample_size: Number of lines to process
            output_dir: Directory for outputs

        Returns:
            Analysis results for the sample
        """
        logger.info(f"Processing sample of {sample_size:,} lines")

        # Read sample records
        records = list(self.reader.read_records(end=sample_size))

        # Process with each analyzer
        results = {}
        for analyzer in tqdm(analyzers, desc="Running analyzers"):
            try:
                chunk_result = analyzer.process_chunk(records)
                # For single chunk, aggregation is simple
                results[analyzer.name] = analyzer.aggregate([chunk_result])
            except Exception as e:
                logger.error(f"Failed to process {analyzer.name}: {e}")
                results[analyzer.name] = {"error": str(e)}

        # Add metadata
        results["_metadata"] = {
            "file_path": str(self.file_path),
            "sample_size": sample_size,
            "records_processed": len(records),
        }

        return results