import multiprocessing
from pathlib import Path
from queue import Empty, Queue
from time import sleep, time
from typing import TYPE_CHECKING, Any, Callable, Optional, Union

from overturemaestro._rich_progress import VERBOSITY_MODE, TrackProgressSpinner

if TYPE_CHECKING:
    from multiprocessing.managers import ValueProxy
    from threading import Lock

    import pyarrow as pa
    import pyarrow.fs as fs
    from rich.progress import Progress

# Using `spawn` method to enable integration with Polars and probably other Rust-based libraries
# https://docs.pola.rs/user-guide/misc/multiprocessing/
ctx: multiprocessing.context.SpawnContext = multiprocessing.get_context("spawn")


class MultiprocessingRuntimeError(RuntimeError): ...


def _job(
    queue: Queue[tuple[str, int]],
    tracker: "ValueProxy[int]",
    tracker_lock: "Lock",
    save_path: Path,
    function: Callable[[str, int, "pa.Table"], "pa.Table"],
    columns: Optional[list[str]],
    filesystem: "fs.FileSystem",
) -> None:  # pragma: no cover
    import pyarrow.dataset as ds
    import pyarrow.parquet as pq

    current_pid = multiprocessing.current_process().pid

    filepath = save_path / f"{current_pid}.parquet"
    writer = None
    while not queue.empty():
        try:
            file_name, row_group_index = None, None
            file_name, row_group_index = queue.get_nowait()

            fragment_manual = ds.ParquetFileFormat().make_fragment(
                file=file_name,
                filesystem=filesystem,
                row_groups=[row_group_index],
            )
            row_group_table = fragment_manual.to_table(columns=columns)
            if row_group_table.num_rows == 0:
                with tracker_lock:
                    tracker.value += 1
                continue

            result_table = function(file_name, row_group_index, row_group_table)

            if result_table.num_rows == 0:
                with tracker_lock:
                    tracker.value += 1
                continue

            if not writer:
                writer = pq.ParquetWriter(filepath, result_table.schema)

            writer.write_table(result_table)

            with tracker_lock:
                tracker.value += 1
        except Empty:
            pass
        except Exception as ex:
            if file_name is not None and row_group_index is not None:
                queue.put((file_name, row_group_index))

            msg = (
                f"Error in worker (PID: {current_pid},"
                f" Parquet: {file_name}, Row group: {row_group_index})"
            )
            raise MultiprocessingRuntimeError(msg) from ex

    if writer:
        writer.close()


class WorkerProcess(ctx.Process):  # type: ignore[name-defined,misc]
    def __init__(self, *args: Any, **kwargs: Any):
        multiprocessing.Process.__init__(self, *args, **kwargs)
        self._pconn, self._cconn = multiprocessing.Pipe()
        self._exception: Optional[tuple[Exception, str]] = None

    def run(self) -> None:  # pragma: no cover
        try:
            multiprocessing.Process.run(self)
            self._cconn.send(None)
        except Exception as e:
            import traceback

            tb: str = traceback.format_exc()
            self._cconn.send((e, tb))

    @property
    def exception(self) -> Optional[tuple[Exception, str]]:
        if self._pconn.poll():
            self._exception = self._pconn.recv()
        return self._exception


def _read_row_group_number(path: str, filesystem: "fs.FileSystem") -> int:
    import pyarrow.parquet as pq

    return int(pq.ParquetFile(path, filesystem=filesystem).num_row_groups)


def map_parquet_dataset(
    dataset_path: Union[str, list[str]],
    destination_path: Path,
    function: Callable[[str, int, "pa.Table"], "pa.Table"],
    progress_description: str,
    columns: Optional[list[str]] = None,
    filesystem: Optional["fs.FileSystem"] = None,
    report_progress_as_text: bool = True,
    verbosity_mode: VERBOSITY_MODE = "transient",
) -> None:
    """
    Apply a function over parquet dataset in a multiprocessing environment.

    Will save results in multiple files in a destination path.

    Args:
        dataset_path (Union[str, list[str]]): Path(s) of the parquet dataset.
        destination_path (Path): Path of the destination.
        function (Callable[[str, pa.Table], pa.Table]): Function to apply over a parquet file name
            and a row group table. Will save resulting table in a new parquet file.
        progress_description (str): Progress bar description.
        columns (Optional[list[str]], optional): List of columns to read. Defaults to `None`.
        filesystem (Optional[fs.FileSystem], optional): Filesystem for the dataset.
            Defaults to `None`.
        report_progress_as_text (bool, optional): Whether to report task progress every minute.
            Defaults to `True`.
        verbosity_mode (Literal["silent", "transient", "verbose"], optional): Set progress
            verbosity mode. Can be one of: silent, transient and verbose. Silent disables
            output completely. Transient tracks progress, but removes output after finished.
            Verbose leaves all progress outputs in the stdout. Defaults to "transient".
    """
    with TrackProgressSpinner(
        "Preparing multiprocessing environment", verbosity_mode=verbosity_mode
    ):
        from concurrent.futures import ProcessPoolExecutor
        from functools import partial

        import pyarrow.parquet as pq

        from overturemaestro._rich_progress import TrackProgressBar

        manager = ctx.Manager()

        queue: Queue[tuple[str, int]] = manager.Queue()
        tracker: ValueProxy[int] = manager.Value("i", 0)
        tracker_lock: Lock = manager.Lock()

        dataset = pq.ParquetDataset(dataset_path, filesystem=filesystem)

        no_cpus = multiprocessing.cpu_count()
        min_no_workers = 32 if no_cpus >= 8 else 16
        no_workers = min(
            max(min_no_workers, no_cpus + 4), 64
        )  # minimum 16 / 32 workers, but not more than 64

    with TrackProgressBar(verbosity_mode=verbosity_mode) as progress:
        total_files = len(dataset.files)
        with ProcessPoolExecutor(max_workers=min(no_workers, total_files)) as ex:
            fn = partial(_read_row_group_number, filesystem=dataset.filesystem)
            row_group_numbers = list(
                progress.track(
                    ex.map(fn, dataset.files, chunksize=1),
                    description="Reading all parquet files row groups",
                    total=total_files,
                )
            )

        for pq_file, row_group_number in zip(dataset.files, row_group_numbers):
            for row_group in range(row_group_number):
                queue.put((pq_file, row_group))

    total = queue.qsize()

    destination_path.mkdir(parents=True, exist_ok=True)

    try:
        processes = [
            WorkerProcess(
                target=_job,
                args=(
                    queue,
                    tracker,
                    tracker_lock,
                    destination_path,
                    function,
                    columns,
                    dataset.filesystem,
                ),
            )
            for _ in range(min(no_cpus, total))
        ]
        with TrackProgressBar(verbosity_mode=verbosity_mode) as progress_bar:
            progress_bar.add_task(description=progress_description, total=total)
            _run_processes(
                processes=processes,
                queue=queue,
                tracker=tracker,
                tracker_lock=tracker_lock,
                total=total,
                progress_bar=progress_bar,
                report_progress_as_text=report_progress_as_text,
            )
    finally:  # pragma: no cover
        _report_exceptions(processes=processes)


def _run_processes(
    processes: list[WorkerProcess],
    queue: Queue[tuple[str, int]],
    tracker: "ValueProxy[int]",
    tracker_lock: "Lock",
    total: int,
    progress_bar: "Progress",
    report_progress_as_text: bool,
) -> None:
    # Run processes
    for p in processes:
        if queue.empty():
            break
        p.start()

    sleep_time = 0.1
    next_update_time = time()
    while any(process.is_alive() for process in processes):
        if any(p.exception for p in processes):  # pragma: no cover
            break

        with tracker_lock:
            completed = tracker.value

        progress_bar.update(task_id=progress_bar.task_ids[0], completed=completed, refresh=True)

        if report_progress_as_text:
            current_time = time()
            if time() >= next_update_time:
                progress_bar.print_progress()
                next_update_time = current_time + 60

        sleep(sleep_time)
        sleep_time = min(1.0, sleep_time + 0.1)

    progress_bar.update(task_id=progress_bar.task_ids[0], completed=total, refresh=True)


def _report_exceptions(processes: list[WorkerProcess]) -> None:
    # In case of exception
    exceptions = []
    for p in processes:
        if p.is_alive():
            p.terminate()

        if p.exception:
            exceptions.append(p.exception)

    if exceptions:
        # use ExceptionGroup in Python3.11
        _raise_multiple(exceptions)


def _raise_multiple(exceptions: list[tuple[Exception, str]]) -> None:
    if not exceptions:
        return
    try:
        error, traceback = exceptions.pop()
        msg = f"{error}\n\nOriginal {traceback}"
        raise type(error)(msg)
    finally:
        _raise_multiple(exceptions)
