#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

from __future__ import annotations

import atexit
from concurrent.futures import Executor
import queue
import threading
from typing import Self
import weakref

from futurist import _utils


class Threading:
    @staticmethod
    def event_object() -> threading.Event:
        return threading.Event()

    @staticmethod
    def lock_object() -> threading.Lock:
        return threading.Lock()

    @staticmethod
    def rlock_object() -> threading.RLock:
        return threading.RLock()

    @staticmethod
    def condition_object(
        lock: threading.Lock | threading.RLock | None = None,
    ) -> threading.Condition:
        return threading.Condition(lock=lock)


_to_be_cleaned: weakref.WeakKeyDictionary[ThreadWorker, bool] = (
    weakref.WeakKeyDictionary()
)
_dying = False


class _Stopping(Exception):
    pass


class ThreadWorker(threading.Thread):
    MAX_IDLE_FOR = 1

    def __init__(
        self,
        executor: Executor,
        work_queue: queue.Queue[_utils.WorkItem],
    ) -> None:
        super().__init__()
        self.work_queue = work_queue
        self.should_stop = False
        self.idle = False
        self.daemon = True
        # Ensure that when the owning executor gets cleaned up that these
        # threads also get shutdown (if they were not already shutdown).
        self.executor_ref = weakref.ref(executor, lambda _obj: self.stop())

    @classmethod
    def create_and_register(
        cls,
        executor: Executor,
        work_queue: queue.Queue[_utils.WorkItem],
    ) -> Self:
        w = cls(executor, work_queue)
        # Ensure that on shutdown, if threads still exist that we get
        # around to cleaning them up and waiting for them to correctly stop.
        #
        # TODO(harlowja): use a weakrefset in the future, as we don't
        # really care about the values...
        _to_be_cleaned[w] = True
        return w

    def _is_dying(self) -> bool:
        if self.should_stop or _dying:
            return True
        executor = self.executor_ref()
        if executor is None:
            return True
        # Avoid confusing the GC with cycles (since each executor
        # references its known workers)...
        del executor
        return False

    def _wait_for_work(self) -> _utils.WorkItem:
        self.idle = True
        work = None
        while work is None:
            try:
                work = self.work_queue.get(True, self.MAX_IDLE_FOR)
            except queue.Empty:
                if self._is_dying():
                    raise _Stopping()
        self.idle = False
        return work

    def stop(self) -> None:
        self.should_stop = True

    def run(self) -> None:
        while not self._is_dying():
            try:
                work = self._wait_for_work()
            except _Stopping:
                return

            try:
                work.run()
            finally:
                # Mark task as done for queue.join() support
                self.work_queue.task_done()
                # Avoid any potential (self) references to the work item
                # in tracebacks or similar...
                del work


def _clean_up() -> None:
    """Ensure all threads that were created were destroyed cleanly."""
    global _dying
    _dying = True
    threads_to_wait_for = []
    while _to_be_cleaned:
        worker, _work_val = _to_be_cleaned.popitem()
        worker.stop()
        threads_to_wait_for.append(worker)
    while threads_to_wait_for:
        worker = threads_to_wait_for.pop()
        try:
            worker.join()
        finally:
            del worker


atexit.register(_clean_up)
