"""
Batch processor with multiprocessing and multithreading support.

This module provides parallel processing capabilities for converting
large batches of equations efficiently.
"""

import os
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
from typing import List, Dict, Any, Optional, Callable
from pathlib import Path
import logging

logger = logging.getLogger(__name__)


class ParallelConverter:
    """
    Parallel equation converter using multiprocessing.

    This class handles CPU-bound equation conversion across multiple processes
    to bypass Python's GIL and maximize CPU utilization.

    Example:
        >>> converter = ParallelConverter(workers=4)
        >>> results = converter.convert_batch(mathml_equations, subject_mode='math')
        >>> successful = [r for r in results if r['success']]
    """

    def __init__(
        self,
        workers: Optional[int] = None,
        subject_mode: str = "auto",
        wrap_identifiers: bool = False,
        prefer_xlongequal: bool = True,
        universal_delimiter_scaling: bool = True,
        show_progress: bool = True,
    ):
        """
        Initialize parallel converter.

        Args:
            workers: Number of worker processes (default: CPU count)
            subject_mode: Subject-specific formatting mode
            wrap_identifiers: Whether to wrap identifiers in \\mathrm{}
            prefer_xlongequal: Whether to use \\xlongequal for arrows
            universal_delimiter_scaling: Whether to add \\left\\right to all delimiters
            show_progress: Whether to show progress bar
        """
        self.workers = workers or min(mp.cpu_count(), 8)  # Cap at 8 to avoid overhead
        self.subject_mode = subject_mode
        self.wrap_identifiers = wrap_identifiers
        self.prefer_xlongequal = prefer_xlongequal
        self.universal_delimiter_scaling = universal_delimiter_scaling
        self.show_progress = show_progress

    def convert_batch(
        self,
        mathml_equations: List[str],
        quiet: bool = False,
    ) -> List[Dict[str, Any]]:
        """
        Convert a batch of MathML equations to LaTeX in parallel.

        Args:
            mathml_equations: List of MathML equation strings
            quiet: Suppress progress output

        Returns:
            List of result dictionaries with keys:
            - 'index': Original equation index
            - 'latex': Converted LaTeX (if successful)
            - 'success': Boolean success flag
            - 'error': Error message (if failed)

        Example:
            >>> equations = ['<math><mi>x</mi></math>', '<math><mi>y</mi></math>']
            >>> results = converter.convert_batch(equations)
            >>> all(r['success'] for r in results)
            True
        """
        if not mathml_equations:
            return []

        # For small batches, sequential processing is faster due to overhead
        if len(mathml_equations) < 5:
            return self._convert_sequential(mathml_equations, quiet)

        # Prepare conversion tasks
        tasks = [
            {
                'index': i,
                'mathml': eq,
                'subject_mode': self.subject_mode,
                'wrap_identifiers': self.wrap_identifiers,
                'prefer_xlongequal': self.prefer_xlongequal,
                'universal_delimiter_scaling': self.universal_delimiter_scaling,
            }
            for i, eq in enumerate(mathml_equations)
        ]

        results = []

        try:
            with ProcessPoolExecutor(max_workers=self.workers) as executor:
                # Submit all tasks
                future_to_task = {
                    executor.submit(_convert_single_equation, task): task
                    for task in tasks
                }

                # Collect results as they complete
                completed = 0
                total = len(tasks)

                for future in as_completed(future_to_task):
                    task = future_to_task[future]
                    completed += 1

                    try:
                        result = future.result()
                        results.append(result)

                        if not quiet and self.show_progress:
                            print(f"\rProgress: {completed}/{total} equations processed", end='', flush=True)

                    except Exception as e:
                        logger.error(f"Failed to convert equation {task['index']}: {e}")
                        results.append({
                            'index': task['index'],
                            'success': False,
                            'error': str(e),
                        })

                if not quiet and self.show_progress:
                    print()  # New line after progress

        except Exception as e:
            logger.error(f"Parallel processing failed: {e}. Falling back to sequential.")
            return self._convert_sequential(mathml_equations, quiet)

        # Sort results by original index to maintain order
        results.sort(key=lambda r: r['index'])
        return results

    def _convert_sequential(
        self,
        mathml_equations: List[str],
        quiet: bool = False,
    ) -> List[Dict[str, Any]]:
        """
        Fallback sequential conversion (single-threaded).

        Args:
            mathml_equations: List of MathML equations
            quiet: Suppress output

        Returns:
            List of conversion results
        """
        # Import here to avoid circular dependencies
        from ..compat import LegacyMathMLToTeXConverter

        converter = LegacyMathMLToTeXConverter(
            wrap_identifiers_in_mathrm=self.wrap_identifiers,
            prefer_xlongequal=self.prefer_xlongequal,
            subject_mode=self.subject_mode,
            universal_delimiter_scaling=self.universal_delimiter_scaling,
        )

        results = []
        for i, mathml in enumerate(mathml_equations):
            try:
                latex = converter.convert(mathml)
                results.append({
                    'index': i,
                    'latex': latex,
                    'success': True,
                })
            except Exception as e:
                results.append({
                    'index': i,
                    'success': False,
                    'error': str(e),
                })

            if not quiet and self.show_progress:
                print(f"\rProgress: {i+1}/{len(mathml_equations)} equations processed", end='', flush=True)

        if not quiet and self.show_progress:
            print()  # New line

        return results


class BatchProcessor:
    """
    Batch file processor with parallel file and equation processing.

    Processes multiple DOCX files in parallel, with parallel equation
    conversion within each file.

    Example:
        >>> processor = BatchProcessor(file_workers=2, equation_workers=4)
        >>> results = processor.process_files(['doc1.docx', 'doc2.docx'])
    """

    def __init__(
        self,
        file_workers: Optional[int] = None,
        equation_workers: Optional[int] = None,
        subject_mode: str = "auto",
        wrap_identifiers: bool = False,
        prefer_xlongequal: bool = True,
    ):
        """
        Initialize batch processor.

        Args:
            file_workers: Workers for parallel file I/O (default: 2)
            equation_workers: Workers for parallel equation conversion (default: CPU count)
            subject_mode: Subject formatting mode
            wrap_identifiers: Wrap identifiers in \\mathrm{}
            prefer_xlongequal: Use \\xlongequal for arrows
        """
        self.file_workers = file_workers or 2
        self.equation_workers = equation_workers
        self.subject_mode = subject_mode
        self.wrap_identifiers = wrap_identifiers
        self.prefer_xlongequal = prefer_xlongequal

    def process_files(
        self,
        docx_paths: List[str],
        output_dir: Optional[str] = None,
        quiet: bool = False,
    ) -> Dict[str, Any]:
        """
        Process multiple DOCX files in parallel.

        Args:
            docx_paths: List of DOCX file paths
            output_dir: Optional output directory for results
            quiet: Suppress progress output

        Returns:
            Dictionary with processing statistics and results

        Example:
            >>> files = ['test1.docx', 'test2.docx']
            >>> results = processor.process_files(files, output_dir='output/')
            >>> print(f"Processed {results['total_equations']} equations")
        """
        if not docx_paths:
            return {'files_processed': 0, 'total_equations': 0, 'results': []}

        # For small file counts, sequential is fine
        if len(docx_paths) <= 2:
            return self._process_files_sequential(docx_paths, output_dir, quiet)

        results = []
        total_equations = 0

        try:
            with ThreadPoolExecutor(max_workers=self.file_workers) as executor:
                # Submit file processing tasks
                future_to_path = {
                    executor.submit(
                        self._process_single_file,
                        path,
                        output_dir,
                        quiet,
                    ): path
                    for path in docx_paths
                }

                # Collect results
                for future in as_completed(future_to_path):
                    path = future_to_path[future]
                    try:
                        result = future.result()
                        results.append(result)
                        total_equations += result['equation_count']

                        if not quiet:
                            print(f"✓ Processed {path}: {result['equation_count']} equations")

                    except Exception as e:
                        logger.error(f"Failed to process {path}: {e}")
                        results.append({
                            'file': path,
                            'success': False,
                            'error': str(e),
                        })

        except Exception as e:
            logger.error(f"Batch file processing failed: {e}")
            return self._process_files_sequential(docx_paths, output_dir, quiet)

        return {
            'files_processed': len([r for r in results if r.get('success', False)]),
            'total_files': len(docx_paths),
            'total_equations': total_equations,
            'results': results,
        }

    def _process_single_file(
        self,
        docx_path: str,
        output_dir: Optional[str],
        quiet: bool,
    ) -> Dict[str, Any]:
        """
        Process a single DOCX file.

        Args:
            docx_path: Path to DOCX file
            output_dir: Output directory
            quiet: Suppress output

        Returns:
            Processing result dictionary
        """
        from ..parsers import InputParserFactory

        try:
            # Parse DOCX to get MathML equations
            parser = InputParserFactory.create_parser(docx_path, parser_type='docx')
            mathml_equations = parser.parse(docx_path)

            # Convert equations in parallel
            converter = ParallelConverter(
                workers=self.equation_workers,
                subject_mode=self.subject_mode,
                wrap_identifiers=self.wrap_identifiers,
                prefer_xlongequal=self.prefer_xlongequal,
                show_progress=not quiet,
            )

            conversion_results = converter.convert_batch(mathml_equations, quiet=quiet)

            # Optionally write output
            if output_dir:
                self._write_output(docx_path, conversion_results, output_dir)

            return {
                'file': docx_path,
                'success': True,
                'equation_count': len(mathml_equations),
                'successful_conversions': sum(1 for r in conversion_results if r['success']),
                'results': conversion_results,
            }

        except Exception as e:
            logger.error(f"Error processing {docx_path}: {e}")
            return {
                'file': docx_path,
                'success': False,
                'error': str(e),
            }

    def _process_files_sequential(
        self,
        docx_paths: List[str],
        output_dir: Optional[str],
        quiet: bool,
    ) -> Dict[str, Any]:
        """Sequential fallback for file processing."""
        results = []
        total_equations = 0

        for path in docx_paths:
            result = self._process_single_file(path, output_dir, quiet)
            results.append(result)
            if result.get('success'):
                total_equations += result['equation_count']

        return {
            'files_processed': len([r for r in results if r.get('success', False)]),
            'total_files': len(docx_paths),
            'total_equations': total_equations,
            'results': results,
        }

    def _write_output(
        self,
        docx_path: str,
        results: List[Dict[str, Any]],
        output_dir: str,
    ) -> None:
        """Write conversion results to output file."""
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)

        # Generate output filename
        input_name = Path(docx_path).stem
        output_file = output_path / f"{input_name}_converted.tex"

        # Write LaTeX equations
        with open(output_file, 'w', encoding='utf-8') as f:
            for i, result in enumerate(results, 1):
                if result['success']:
                    f.write(f"% Equation {i}\n")
                    f.write(f"${result['latex']}$\n\n")
                else:
                    f.write(f"% Equation {i} - Error: {result.get('error', 'Unknown')}\n\n")


def _convert_single_equation(task: Dict[str, Any]) -> Dict[str, Any]:
    """
    Worker function for parallel equation conversion.

    This function is called in separate processes by ProcessPoolExecutor.
    Must be defined at module level for pickling.

    Args:
        task: Dictionary with conversion parameters

    Returns:
        Result dictionary with conversion outcome
    """
    # Import here (in worker process) to avoid pickling issues
    from ..compat import LegacyMathMLToTeXConverter

    try:
        # Create converter instance in worker process
        converter = LegacyMathMLToTeXConverter(
            wrap_identifiers_in_mathrm=task['wrap_identifiers'],
            prefer_xlongequal=task['prefer_xlongequal'],
            subject_mode=task['subject_mode'],
            universal_delimiter_scaling=task['universal_delimiter_scaling'],
        )

        # Convert equation
        latex = converter.convert(task['mathml'])

        return {
            'index': task['index'],
            'latex': latex,
            'success': True,
        }

    except Exception as e:
        return {
            'index': task['index'],
            'success': False,
            'error': str(e),
        }