# This file is part of torf.
#
# torf is free software: you can redistribute it and/or modify it under the
# terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# torf is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
# A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with torf.  If not, see <https://www.gnu.org/licenses/>.

import errno
import logging
import os
import queue
import threading
from hashlib import sha1
from time import monotonic as time_monotonic

from . import _errors as errors
from ._stream import TorrentFileStream

QUEUE_CLOSED = object()

_debug = logging.getLogger("torf").debug


def _thread_name():
    return threading.current_thread().name


def _pretty_bytes(b):
    if isinstance(b, (bytes, bytearray)) and len(b) > 8:
        # return b[:8].hex() + '...' + b[-8:].hex()
        return b[:8] + b"..." + b[-8:]
    else:
        return b


class Worker:
    """
    :class:`threading.Thread` subclass that re-raises any exceptions from the
    thread when joined
    """

    def __init__(self, name, worker, start=True, fail_ok=False):
        self._exception = None
        self._name = str(name)
        self._worker = worker
        self._thread = threading.Thread(name=self._name, target=self._run_and_catch_exceptions)
        if start:
            self.start(fail_ok=fail_ok)

    @property
    def exception(self):
        return self._exception

    @property
    def name(self):
        return self._name

    @property
    def is_running(self):
        return self._thread.is_alive()

    def _run_and_catch_exceptions(self):
        try:
            self._worker()
        except BaseException as e:
            self._exception = e

    def start(self, fail_ok=False):
        if not self._thread.is_alive():
            try:
                self._thread.start()
            except RuntimeError as e:
                if fail_ok:
                    _debug(f"{self.name}: Failed to start thread: {e!r} - but that's ok")
                else:
                    _debug(f"{self.name}: Failed to start thread: {e!r}")
                    raise
            else:
                _debug(f"{self.name}: Started")

    def join(self, *args, **kwargs):
        if self.is_running:
            self._thread.join(*args, **kwargs)
        if self._exception:
            raise self._exception


class Reader(Worker):
    """
    :class:`Worker` subclass that reads files in pieces and pushes them to a
    queue
    """

    def __init__(self, *, torrent, queue_size, path=None):
        self._torrent = torrent
        self._path = path
        self._piece_queue = queue.Queue(maxsize=queue_size)
        self._stop = False
        self._memory_error_timestamp = -1
        super().__init__(name="reader", worker=self._push_pieces)

    def _push_pieces(self):
        stream = TorrentFileStream(self._torrent)
        try:
            iter_pieces = stream.iter_pieces(self._path, oom_callback=self._handle_oom)
            for piece_index, (piece, filepath, exceptions) in enumerate(iter_pieces):
                # _debug(f'{_thread_name()}: Read #{piece_index}')
                if self._stop:
                    _debug(f"{_thread_name()}: Stopped reading")
                    break
                elif exceptions:
                    self._push_piece(
                        piece_index=piece_index,
                        filepath=filepath,
                        exceptions=exceptions,
                    )
                elif piece:
                    self._push_piece(piece_index=piece_index, filepath=filepath, piece=piece)
                else:
                    # `piece` is None because of missing file, and the exception
                    # was already sent for the first `piece_index` of that file
                    self._push_piece(piece_index=piece_index, filepath=filepath)

                # _debug(f'{_thread_name()}: {self._piece_queue.qsize()} pieces queued')

        except BaseException as e:
            _debug(f"{_thread_name()}: Exception while reading: {e!r}")
            raise

        finally:
            self._piece_queue.put(QUEUE_CLOSED)
            _debug(f"{_thread_name()}: Piece queue is now exhausted")
            stream.close()

    def _push_piece(self, *, piece_index, filepath, piece=None, exceptions=()):
        # _debug(f'{_thread_name()}: Pushing #{piece_index}: {filepath}: {_pretty_bytes(piece)}, {exceptions!r}')
        self._piece_queue.put((piece_index, filepath, piece, exceptions))

    def _handle_oom(self, exception):
        # Reduce piece_queue.maxsize by 1 every 100ms until the MemoryErrors stop
        now = time_monotonic()
        time_diff = now - self._memory_error_timestamp
        if time_diff >= 0.1:
            old_maxsize = self._piece_queue.maxsize
            new_maxsize = max(1, int(old_maxsize * 0.9))
            if new_maxsize != old_maxsize:
                _debug(f"{_thread_name()}: Reducing piece_queue.maxsize to {new_maxsize}")
                self._piece_queue.maxsize = new_maxsize
                self._memory_error_timestamp = now
            else:
                raise errors.ReadError(errno.ENOMEM, exception)

    def stop(self):
        """Stop reading and close the piece queue"""
        if not self._stop:
            _debug(f"{_thread_name()}: {type(self).__name__}: Setting stop flag")
            self._stop = True

    @property
    def piece_queue(self):
        """
        :class:`queue.Queue` instance that gets evenly sized pieces from the
        concatenated stream of files
        """
        return self._piece_queue


class HasherPool:
    """
    Wrapper around one or more :class:`Worker` instances that each read a piece
    from :attr:`Reader.piece_queue`, feed it to :func:`~.hashlib.sha1`, and push
    the resulting hash to :attr:`hash_queue`
    """

    def __init__(self, hasher_threads, piece_queue):
        self._piece_queue = piece_queue
        self._hash_queue = queue.Queue()
        self._finalize_event = threading.Event()

        # Janitor takes care of closing the hash queue, removing idle hashers, etc
        self._janitor = Worker(
            name="janitor",
            worker=self._janitor_thread,
            start=False,
        )

        # Hashers read from piece_queue and push to hash_queue
        self._hashers = [
            Worker(
                name="hasher1",
                # One hasher is vital an may not die from boredom
                worker=lambda: self._hasher_thread(is_vital=True),
                start=False,
            ),
        ]
        for i in range(2, hasher_threads + 1):
            self._hashers.append(
                Worker(
                    name=f"hasher{i}",
                    # All other hashers should die if they are bored
                    worker=lambda: self._hasher_thread(is_vital=False),
                    start=False,
                )
            )

        # Start threads manually after they were created to prevent race
        # conditions and make sure all required threads are running
        self._janitor.start(fail_ok=False)

        # Hashers are allowed to fail (e.g. because of OS limits), but we need
        # at least one to start successfully
        self._hashers[0].start(fail_ok=False)
        for hasher in self._hashers[1:]:
            hasher.start(fail_ok=True)

    def _hasher_thread(self, is_vital=True):
        piece_queue = self._piece_queue
        handle_piece = self._handle_piece
        while True:
            # _debug(f'{_thread_name()}: Waiting for next task')
            try:
                task = piece_queue.get(timeout=0.5)
            except queue.Empty:
                if not is_vital:
                    _debug(f"{_thread_name()}: I am bored, byeee!")
                    break
                else:
                    _debug(f"{_thread_name()}: I am bored, but needed.")
            else:
                if task is QUEUE_CLOSED:
                    _debug(f"{_thread_name()}: piece_queue is closed")
                    # Repeat QUEUE_CLOSED to the next sibling. This ensures
                    # there is always one more QUEUE_CLOSED queued than running
                    # threads. Otherwise, one thread might consume multiple
                    # QUEUE_CLOSED and leave other threads running forvever.
                    piece_queue.put(QUEUE_CLOSED)
                    # Signal janitor to initiate shutdown procedure
                    self._finalize_event.set()
                    break
                else:
                    handle_piece(*task)

    def _handle_piece(self, piece_index, filepath, piece, exceptions):
        if exceptions:
            # _debug(f'{_thread_name()}: Forwarding exceptions for #{piece_index}: {exceptions!r}')
            self._hash_queue.put((piece_index, filepath, None, exceptions))

        elif piece:
            piece_hash = sha1(piece).digest()
            # _debug(f'{_thread_name()}: Hashed #{piece_index}: {_pretty_bytes(piece)} [{len(piece)} bytes] -> {piece_hash}')
            self._hash_queue.put((piece_index, filepath, piece_hash, ()))

        else:
            # _debug(f'{_thread_name()}: Nothing to hash for #{piece_index}: {piece!r}')
            self._hash_queue.put((piece_index, filepath, None, ()))

    def _janitor_thread(self):
        while True:
            _debug(f"{_thread_name()}: Waiting for finalize event")
            finalization_initiated = self._finalize_event.wait(timeout=1.0)
            if finalization_initiated:
                self._wait_for_hashers()
                _debug(f"{_thread_name()}: Closing hash queue")
                self._hash_queue.put(QUEUE_CLOSED)
                break

            else:
                # Remove terminated idle hashers
                for hasher in tuple(self._hashers):
                    if not hasher.is_running:
                        _debug(f"{_thread_name()}: Pruning {hasher.name}")
                        self._hashers.remove(hasher)

        _debug(f"{_thread_name()}: Terminating")

    def _wait_for_hashers(self):
        while True:
            # _debug(f'{_thread_name()}: Hashers running: {[h.name for h in self._hashers if h.is_running]}')
            if all(not h.is_running for h in self._hashers):
                _debug(f"{_thread_name()}: All hashers terminated")
                break

    def join(self):
        """Block until all threads have terminated"""
        for hasher in self._hashers:
            _debug(f"{_thread_name()}: Joining {hasher.name}")
            hasher.join()
        _debug(f"{_thread_name()}: Joined all hashers")

        _debug(f"{_thread_name()}: Joining {self._janitor.name}")
        self._janitor.join()
        _debug(f"{_thread_name()}: Joined {self._janitor.name}")

    @property
    def hash_queue(self):
        """:class:`queue.Queue` instance that gets piece hashes"""
        return self._hash_queue


class Collector:
    """
    Consume items from :attr:`HasherPool.hash_queue` and ensure proper
    termination of all threads if anything goes wrong or the user cancels the
    operation
    """

    def __init__(self, torrent, reader, hashers, callback=None):
        self._reader = reader
        self._hashers = hashers
        self._callback = callback
        self._hashes_unsorted = []
        self._pieces_seen = []
        self._pieces_total = torrent.pieces

    def collect(self):
        """
        Read piece hashes from :attr:`HasherPool.hash_queue`

        When this method returns, :attr:`hashes` is an ordered sequence of
        collected piece hashes.

        Exceptions from :class:`Reader`, :class:`HasherPool` or the provided
        callback are raised after all threads are terminated and joined.

        :return: the same value as :attr:`hashes`
        """
        try:
            hash_queue = self._hashers.hash_queue
            while True:
                # _debug(f'{_thread_name()}: Waiting for next piece hash')
                task = hash_queue.get()
                # _debug(f'{_thread_name()}: Got task: {task}')
                if task is QUEUE_CLOSED:
                    break
                else:
                    self._collect(*task)

        except BaseException as e:
            _debug(f"{_thread_name()}: Exception while dequeueing piece hashes: {e!r}")
            self._cancel()
            raise

        finally:
            self._finalize()

        return self.hashes

    def _collect(self, piece_index, filepath, piece_hash, exceptions):
        # _debug(f'{_thread_name()}: Collecting #{piece_index}: {_pretty_bytes(piece_hash)}, {exceptions}')

        # Remember which pieces where hashed to count them and for sanity checking
        assert piece_index not in self._pieces_seen
        self._pieces_seen.append(piece_index)

        # Collect piece
        if not exceptions and piece_hash:
            self._hashes_unsorted.append((piece_index, piece_hash))

        # If there is no callback, raise first exception
        if exceptions and not self._callback:
            raise exceptions[0]

        # Report progress/exceptions and allow callback to cancel
        if self._callback:
            # _debug(f'{_thread_name()}: Collector callback: {self._callback}')
            maybe_cancel = self._callback(
                piece_index,
                len(self._pieces_seen),
                self._pieces_total,
                filepath,
                piece_hash,
                exceptions,
            )
            # _debug(f'{_thread_name()}: Collector callback return value: {maybe_cancel}')
            if maybe_cancel is not None:
                self._cancel()

    def _cancel(self):
        # NOTE: We don't need to stop HasherPool or Collector.collect() because
        #       they will stop when Reader._push_pieces() pushes QUEUE_CLOSED.
        #       They will process the pieces in the queue, but that shouldn't
        #       take long unless the Reader's queue size is too big.
        self._reader.stop()

    def _finalize(self):
        _debug(f"{_thread_name()}: Joining {self._reader}")
        self._reader.join()
        _debug(f"{_thread_name()}: Joining {self._hashers}")
        self._hashers.join()
        _debug(f"{_thread_name()}: hash_queue has {self._hashers.hash_queue.qsize()} items left")

    @property
    def hashes(self):
        """Ordered sequence of piece hashes"""
        return tuple(hash for index, hash in sorted(self._hashes_unsorted))


class _IntervaledCallback:
    """
    Callable that calls `callback`, but only if at least `interval` seconds
    elapsed since the previous call
    """

    def __init__(self, callback, interval=0):
        self._callback = callback
        self._interval = interval
        self._prev_call_time = -1

    def __call__(self, *args, force=False):
        now = time_monotonic()
        diff = now - self._prev_call_time
        # _debug(f'{_thread_name()}: Callback? {force=} or {diff=} >= {self._interval=}')
        if force or diff >= self._interval:
            self._prev_call_time = now
            # _debug(f'{_thread_name()}: Callback! {args=}')
            return self._callback(*args)


class _TranslatingCallback:
    def __init__(self, callback, interval, torrent):
        self._callback = callback
        self._torrent = torrent
        self._intervaled_callback = _IntervaledCallback(
            callback=self._call_callback,
            interval=interval,
        )

    def __call__(self, piece_index, pieces_done, pieces_total, filepath, piece_hash, exceptions):
        force = self._force_callback(
            piece_index, pieces_done, pieces_total, filepath, piece_hash, exceptions
        )
        return self._intervaled_callback(
            piece_index,
            pieces_done,
            pieces_total,
            filepath,
            piece_hash,
            exceptions,
            force=force,
        )

    def _force_callback(
        self, piece_index, pieces_done, pieces_total, filepath, piece_hash, exceptions
    ):
        # Figure out if we must ignore the interval for this call. This method
        # is called for every hashed piece and should be as efficient as
        # possible.
        raise NotImplementedError("You must implement this method!")

    def _call_callback(
        self, piece_index, pieces_done, pieces_total, filepath, piece_hash, exceptions
    ):
        # Translate arguments for the actual callback. This method is only
        # called at intervals (e.g. once per second).
        raise NotImplementedError("You must implement this method!")


class GenerateCallback(_TranslatingCallback):
    """
    Translate arguments from :class:`Collector` to what's specified by
    :meth:`~.Torrent.generate`
    """

    def _force_callback(
        self, piece_index, pieces_done, pieces_total, filepath, piece_hash, exceptions
    ):
        return exceptions or pieces_done >= pieces_total

    def _call_callback(
        self, piece_index, pieces_done, pieces_total, filepath, piece_hash, exceptions
    ):
        if exceptions:
            # Torrent creation errors are always fatal and must be raised
            raise exceptions[0]
        elif self._callback:
            # Report progress and allow callback to cancel
            return self._callback(self._torrent, filepath, pieces_done, pieces_total)


class VerifyCallback(_TranslatingCallback):
    """
    Translate arguments from :class:`Collector` to what's specified by
    :meth:`~.Torrent.verify`
    """

    def __init__(self, *args, path, **kwargs):
        super().__init__(*args, **kwargs)

        # Store piece hashes from the torrent for quick access
        self._exp_hashes = self._torrent.hashes

        # Map expected file system paths to expected file sizes
        # NOTE: The last segment in `path` is supposed to be the torrent name so
        #       we must remove the name stored in the torrent file from each
        #       `file`. This allows verification of any renamed file/directory
        #       against a torrent.
        self._exp_file_sizes = tuple(
            (
                os.sep.join((str(path), *file.parts[1:])),
                self._torrent.partial_size(file),
            )
            for file in self._torrent.files
        )

    def _force_callback(
        self, piece_index, pieces_done, pieces_total, filepath, piece_hash, exceptions
    ):
        return (
            # Always report exceptions
            exceptions
            # Always report completion
            or pieces_done >= pieces_total
            # Always report hash mismatch
            or (piece_hash is not None and piece_hash != self._exp_hashes[piece_index])
        )

    def _call_callback(
        self, piece_index, pieces_done, pieces_total, filepath, piece_hash, exceptions
    ):
        if (
            # Don't add verification error if there are other errors
            not exceptions
            # Piece hash was calculated and doesn't match
            and piece_hash is not None
            and piece_hash != self._exp_hashes[piece_index]
        ):
            exceptions = (
                errors.VerifyContentError(
                    filepath,
                    piece_index,
                    self._torrent.piece_size,
                    self._exp_file_sizes,
                ),
            )

        if self._callback:
            # Callback can raise exception or handle it otherwise
            def call_callback(fpath, exception):
                return self._callback(
                    self._torrent,
                    fpath,
                    pieces_done,
                    pieces_total,
                    piece_index,
                    piece_hash,
                    exception,
                )

            if exceptions:
                # Call callback for each exception until it indicates
                # cancellation by returning anything truthy
                for exception in exceptions:
                    fpath = self._get_path_from_exception(exception)
                    maybe_cancel = call_callback(fpath, exception)
                    if maybe_cancel is not None:
                        return maybe_cancel
            else:
                # Report progress and return cancellation indicator
                return call_callback(filepath, None)

        elif exceptions:
            # Default to raising first exception
            raise exceptions[0]

    @staticmethod
    def _get_path_from_exception(exception):
        for attr in ("filepath", "path"):
            try:
                return getattr(exception, attr)
            except AttributeError:
                pass

        raise RuntimeError(f"Failed to get path from {exception!r}")