"""

Implements a thread pool for parallel copying of files.

"""

from __future__ import unicode_literals

import typing

import threading
from six.moves.queue import Queue

from .copy import copy_file_internal, copy_modified_time
from .errors import BulkCopyFailed
from .tools import copy_file_data

if typing.TYPE_CHECKING:
    from typing import IO, List, Optional, Text, Tuple, Type

    from types import TracebackType

    from .base import FS


class _Worker(threading.Thread):
    """Worker thread that pulls tasks from a queue."""

    def __init__(self, copier):
        # type (Copier) -> None
        self.copier = copier
        super(_Worker, self).__init__()
        self.daemon = True

    def run(self):
        # type () -> None
        queue = self.copier.queue
        while True:
            task = queue.get(block=True)
            try:
                if task is None:
                    break  # Sentinel to exit thread
                task()
            except Exception as error:
                self.copier.add_error(error)
            finally:
                queue.task_done()


class _Task(object):
    """Base class for a task."""

    def __call__(self):
        # type: () -> None
        """Task implementation."""


class _CopyTask(_Task):
    """A callable that copies from one file another."""

    def __init__(self, src_file, dst_file):
        # type: (IO, IO) -> None
        self.src_file = src_file
        self.dst_file = dst_file

    def __call__(self):
        # type: () -> None
        try:
            copy_file_data(self.src_file, self.dst_file, chunk_size=1024 * 1024)
        finally:
            try:
                self.src_file.close()
            finally:
                self.dst_file.close()


class Copier(object):
    """Copy files in worker threads."""

    def __init__(self, num_workers=4, preserve_time=False):
        # type: (int, bool) -> None
        if num_workers < 0:
            raise ValueError("num_workers must be >= 0")
        self.num_workers = num_workers
        self.preserve_time = preserve_time
        self.all_tasks = []  # type: List[Tuple[FS, Text, FS, Text]]
        self.queue = None  # type: Optional[Queue[_Task]]
        self.workers = []  # type: List[_Worker]
        self.errors = []  # type: List[Exception]
        self.running = False

    def start(self):
        """Start the workers."""
        if self.num_workers:
            self.queue = Queue(maxsize=self.num_workers)
            self.workers = [_Worker(self) for _ in range(self.num_workers)]
            for worker in self.workers:
                worker.start()
        self.running = True

    def stop(self):
        """Stop the workers (will block until they are finished)."""
        if self.running and self.num_workers:
            # Notify the workers that all tasks have arrived
            # and wait for them to finish.
            for _worker in self.workers:
                self.queue.put(None)
            for worker in self.workers:
                worker.join()

            # If the "last modified" time is to be preserved, do it now.
            if self.preserve_time:
                for args in self.all_tasks:
                    copy_modified_time(*args)

            # Free up references held by workers
            del self.workers[:]
            self.queue.join()
        self.running = False

    def add_error(self, error):
        """Add an exception raised by a task."""
        self.errors.append(error)

    def __enter__(self):
        self.start()
        return self

    def __exit__(
        self,
        exc_type,  # type: Optional[Type[BaseException]]
        exc_value,  # type: Optional[BaseException]
        traceback,  # type: Optional[TracebackType]
    ):
        self.stop()
        if traceback is None and self.errors:
            raise BulkCopyFailed(self.errors)

    def copy(self, src_fs, src_path, dst_fs, dst_path, preserve_time=False):
        # type: (FS, Text, FS, Text, bool) -> None
        """Copy a file from one fs to another."""
        if self.queue is None:
            # This should be the most performant for a single-thread
            copy_file_internal(
                src_fs, src_path, dst_fs, dst_path, preserve_time=self.preserve_time
            )
        else:
            self.all_tasks.append((src_fs, src_path, dst_fs, dst_path))
            src_file = src_fs.openbin(src_path, "r")
            try:
                dst_file = dst_fs.openbin(dst_path, "w")
            except Exception:
                src_file.close()
                raise
            task = _CopyTask(src_file, dst_file)
            self.queue.put(task)