import asyncio
import concurrent.futures  # For real parallelism via threads
from pathlib import Path
from typing import TYPE_CHECKING, Any

from box import Box
from opentelemetry import trace
from pandas import DataFrame
from rich.progress import (  # Import Rich Progress
    BarColumn,
    Progress,
    SpinnerColumn,
    TextColumn,
    TimeElapsedColumn,
)

from flock.config import TELEMETRY
from flock.core.context.context import FlockContext
from flock.core.context.context_vars import FLOCK_BATCH_SILENT_MODE
from flock.core.flock_agent import FlockAgent
from flock.core.logging.logging import get_logger

try:
    import pandas as pd

    PANDAS_AVAILABLE = True
except ImportError:
    pd = None
    PANDAS_AVAILABLE = False

if TYPE_CHECKING:
    from flock.core.flock import Flock

logger = get_logger("flock")
TELEMETRY.setup_tracing()  # Setup OpenTelemetry
tracer = trace.get_tracer(__name__)


class BatchProcessor:
    def __init__(self, flock_instance: "Flock"):
        self.flock = flock_instance

    async def run_batch_async(
        self,
        start_agent: FlockAgent | str,
        batch_inputs: list[dict[str, Any]] | DataFrame | str,
        input_mapping: dict[str, str] | None = None,
        static_inputs: dict[str, Any] | None = None,
        parallel: bool = True,
        max_workers: int = 5,
        use_temporal: bool | None = None,
        box_results: bool = True,
        return_errors: bool = False,
        silent_mode: bool = False,
        write_to_csv: str | None = None,
        hide_columns: list[str] | None = None,
        delimiter: str = ",",
    ) -> list[Box | dict | None | Exception]:
        """Runs the specified agent/workflow for each item in a batch asynchronously.

        Args:
            start_agent: Agent instance or name to start each run.
            batch_inputs: Input data in one of these forms:
                - List of dictionaries, each representing inputs for one run
                - Pandas DataFrame where each row is inputs for one run
                - String path to a CSV file to load as DataFrame
            input_mapping: Maps DataFrame/CSV column names to agent input keys (required for DataFrame/CSV).
            static_inputs: Dictionary of inputs constant across all batch runs.
            parallel: Whether to run local jobs in parallel (ignored if use_temporal=True).
            max_workers: Max concurrent local workers (used if parallel=True and use_temporal=False).
            use_temporal: Override Flock's 'enable_temporal' setting for this batch.
            box_results: Wrap successful dictionary results in Box objects.
            return_errors: If True, return Exception objects for failed runs instead of raising.
            silent_mode: If True, suppress output and show progress bar instead.
            write_to_csv: Path to save results as CSV file.
            hide_columns: List of column names to hide from output.

        Returns:
            List containing results (Box/dict), None (if error and not return_errors),
            or Exception objects (if error and return_errors). Order matches input.

        Raises:
            ValueError: For invalid input combinations.
            ImportError: If DataFrame/CSV used without pandas.
            Exception: First exception from a run if return_errors is False.
        """
        effective_use_temporal = (
            use_temporal
            if use_temporal is not None
            else self.flock.enable_temporal
        )
        exec_mode = (
            "Temporal"
            if effective_use_temporal
            else ("Parallel Local" if parallel else "Sequential Local")
        )
        logger.info(
            f"Starting batch run for agent '{start_agent}'. Execution: {exec_mode}, Silent: {silent_mode}"
        )

        # --- Input Preparation ---
        prepared_batch_inputs: list[dict[str, Any]] = []

        if input_mapping == {}:
            input_mapping = None
        if static_inputs == {}:
            static_inputs = None

        if isinstance(batch_inputs, str):
            # Handle CSV file input
            try:
                df = pd.read_csv(batch_inputs)
                logger.debug(
                    f"Loaded CSV file with {len(df)} rows: {batch_inputs}"
                )
                batch_inputs = df  # Convert to DataFrame for unified handling
            except Exception as e:
                raise ValueError(
                    f"Failed to load CSV file '{batch_inputs}': {e}"
                )

        if isinstance(batch_inputs, DataFrame):
            # Handle DataFrame input
            logger.debug(
                f"Converting DataFrame ({len(batch_inputs)} rows) to batch inputs."
            )
            for _, row in batch_inputs.iterrows():
                if input_mapping:
                    item_input = {
                        agent_key: row[df_col]
                        for df_col, agent_key in input_mapping.items()
                        if df_col in row
                    }
                else:
                    item_input = row.to_dict()
                prepared_batch_inputs.append(item_input)
        else:
            # Handle list of dictionaries
            if not isinstance(batch_inputs, list):
                raise ValueError(
                    "batch_inputs must be a list of dictionaries, DataFrame, or CSV file path"
                )

            if input_mapping:
                # Apply mapping to dictionary inputs
                logger.debug("Applying input mapping to dictionary inputs")
                for item in batch_inputs:
                    mapped_input = {}
                    for df_col, agent_key in input_mapping.items():
                        if df_col in item:
                            mapped_input[agent_key] = item[df_col]
                        else:
                            logger.warning(
                                f"Input mapping key '{df_col}' not found in input dictionary"
                            )
                    prepared_batch_inputs.append(mapped_input)
            else:
                # Use dictionaries as-is if no mapping provided
                prepared_batch_inputs = batch_inputs

            logger.debug(
                f"Using provided list of {len(prepared_batch_inputs)} batch inputs."
            )

        if not prepared_batch_inputs:
            return []

        # --- Setup Progress Bar if Silent ---
        progress_context = None
        progress_task_id = None
        if silent_mode:
            progress = Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                BarColumn(),
                TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
                TextColumn("({task.completed}/{task.total})"),
                TimeElapsedColumn(),
                # transient=True # Optionally remove progress bar when done
            )
            progress_context = progress  # Use as context manager
            progress_task_id = progress.add_task(
                f"Processing Batch ({exec_mode})",
                total=len(prepared_batch_inputs),
            )
            progress.start()

        results = [None] * len(prepared_batch_inputs)  # Pre-allocate results list

        # --- Worker Definitions ---
        # We implement two flavours:
        #   * async_worker: used for Temporal or sequential runs (keeps the original behaviour)
        #   * thread_worker: executes the run in a dedicated thread via ThreadPoolExecutor for true parallelism.

        async def async_worker(index: int, item_inputs: dict[str, Any]):
            """Original coroutine worker used for non-threaded execution paths."""
            full_input = {**(static_inputs or {}), **item_inputs}
            context = FlockContext()
            context.set_variable(FLOCK_BATCH_SILENT_MODE, silent_mode)

            run_desc = f"Batch item {index + 1}"
            logger.debug(f"{run_desc} started (async).")
            try:
                result = await self.flock.run_async(
                    start_agent,
                    full_input,
                    box_result=box_results,
                    context=context,
                )
                results[index] = result
                logger.debug(f"{run_desc} finished successfully.")
            except Exception as e:
                logger.error(f"{run_desc} failed: {e}", exc_info=not return_errors)
                if return_errors:
                    results[index] = e
                else:
                    raise  # Propagate to calling gather
            finally:
                if progress_context:
                    progress.update(progress_task_id, advance=1)

        # ThreadPool worker for real parallelism (suitable for blocking I/O)
        def _thread_worker(index: int, item_inputs: dict[str, Any]):
            """Synchronous helper executed inside a worker thread."""
            full_input = {**(static_inputs or {}), **item_inputs}
            run_desc = f"Batch item {index + 1}"
            logger.debug(f"{run_desc} started (thread).")
            try:
                # Use the synchronous wrapper to avoid nested event-loop issues inside threads
                result = self.flock.run(
                    start_agent=start_agent,
                    input=full_input,
                    box_result=box_results,
                )
                logger.debug(f"{run_desc} finished successfully.")
                return index, result, None
            except Exception as e:
                logger.error(f"{run_desc} failed: {e}")
                return index, None, e

        async def thread_worker(executor, index: int, item_inputs: dict[str, Any]):
            """Coroutine wrapper that submits _thread_worker to the specified executor."""
            loop = asyncio.get_running_loop()
            idx, res, err = await loop.run_in_executor(
                executor, _thread_worker, index, item_inputs
            )
            # Handle result / error on the asyncio side
            if err:
                if return_errors:
                    results[idx] = err
                else:
                    raise err
            else:
                results[idx] = res
            if progress_context:
                progress.update(progress_task_id, advance=1)

        tasks = []
        try:
            if effective_use_temporal:
                # Temporal Batching (Simplified: sequential execution for this example)
                # A real implementation might use start_workflow or signals
                logger.info(
                    "Running batch using Temporal (executing sequentially for now)..."
                )
                for i, item_data in enumerate(prepared_batch_inputs):
                    await async_worker(i, item_data)  # Run sequentially for demo
                # TODO: Implement true parallel Temporal workflow execution if needed

            elif parallel:
                # --- Real parallelism using ThreadPoolExecutor ---
                logger.info(
                    f"Running batch in parallel (threads) with max_workers={max_workers}..."
                )
                loop = asyncio.get_running_loop()
                with concurrent.futures.ThreadPoolExecutor(
                    max_workers=max_workers, thread_name_prefix="flock-batch"
                ) as executor:
                    for i, item_data in enumerate(prepared_batch_inputs):
                        tasks.append(
                            asyncio.create_task(
                                thread_worker(executor, i, item_data)
                            )
                        )

                    # Wait for all tasks allowing exceptions to propagate as needed
                    await asyncio.gather(*tasks)

            else:  # Sequential Local
                logger.info("Running batch sequentially...")
                for i, item_data in enumerate(prepared_batch_inputs):
                    await async_worker(i, item_data)  # Already handles errors internally based on return_errors

            logger.info("Batch execution finished.")

        except Exception as batch_error:
            # This catch handles errors re-raised from workers when return_errors=False
            logger.error(f"Batch execution stopped due to error: {batch_error}")
            # No need to cancel tasks here as gather would have stopped
            if not return_errors:
                raise  # Re-raise the first error encountered if not returning errors
        finally:
            if progress_context:
                progress.stop()

        if write_to_csv:
            try:
                df = pd.DataFrame(results)
                if hide_columns:
                    df = df.drop(columns=hide_columns)
                # create write_to_csv directory if it doesn't exist
                Path(write_to_csv).parent.mkdir(parents=True, exist_ok=True)
                df.to_csv(write_to_csv, index=False, sep=delimiter)
                logger.info(f"Results written to CSV file: {write_to_csv}")
            except Exception as e:
                logger.error(f"Failed to write results to CSV: {e}")

        return results

    def run_batch(  # Synchronous wrapper
        self,
        start_agent: FlockAgent | str,
        batch_inputs: list[dict[str, Any]] | DataFrame | str,
        input_mapping: dict[str, str] | None = None,
        static_inputs: dict[str, Any] | None = None,
        parallel: bool = True,
        max_workers: int = 5,
        use_temporal: bool | None = None,
        box_results: bool = True,
        return_errors: bool = False,
        silent_mode: bool = False,
        write_to_csv: str | None = None,
        hide_columns: list[str] | None = None,
        delimiter: str = ",",
    ) -> list[Box | dict | None | Exception]:
        """Synchronous wrapper for run_batch_async."""
        # (Standard asyncio run wrapper - same as in previous suggestion)
        try:
            loop = asyncio.get_running_loop()
            if loop.is_closed():
                raise RuntimeError("Event loop is closed")
        except RuntimeError:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)

        coro = self.run_batch_async(
            start_agent=start_agent,
            batch_inputs=batch_inputs,
            input_mapping=input_mapping,
            static_inputs=static_inputs,
            parallel=parallel,
            max_workers=max_workers,
            use_temporal=use_temporal,
            box_results=box_results,
            return_errors=return_errors,
            silent_mode=silent_mode,
            write_to_csv=write_to_csv,
            hide_columns=hide_columns,
            delimiter=delimiter,
        )

        if asyncio.get_event_loop() is loop and not loop.is_running():
            results = loop.run_until_complete(coro)
            # loop.close() # Avoid closing potentially shared loop
            return results
        else:
            # Run within an existing loop
            future = asyncio.ensure_future(coro)
            return loop.run_until_complete(future)
