# Copyright © The Debusine Developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.

"""Classes and functions supporting database tests."""

import hashlib
import threading
from collections.abc import Callable
from concurrent.futures import Future, ThreadPoolExecutor
from typing import Any

from django.db import connection, connections, transaction

from debusine.db.models import File


class RunInParallelTransaction[R: Any]:
    """
    Execute a callable in a separate transaction that gets started in a thread.

    The :py:meth:`start_transaction` starts a new thread in which it starts
    the transaction and executes the callable. The transaction then stays
    open and the thread waits until we decide to call
    :py:meth:`stop_transaction`.

    Alternatively, use :py:meth:`run_and_wait` to do both the above operations
    in a single step.
    """

    def __init__(self, to_run: Callable[[], R]) -> None:
        """
        Create a RunInParallelTransaction object.

        :param to_run: the function/callable object that gets executed in
          a transaction started in a thread.
        """
        self._to_run = to_run
        self._transaction_ready = threading.Event()
        self._transaction_close = threading.Event()
        self._executor: ThreadPoolExecutor | None = None
        self._future: Future[R] | None = None

    def start_transaction(self) -> None:
        """
        Start a transaction and execute the registered callable.

        It does this in a new thread that then waits until
        :py:meth:stop_transaction() gets called.
        """

        def run_in_thread() -> R:
            try:
                with transaction.atomic():
                    result = self._to_run()
                    self._transaction_ready.set()
                    self._transaction_close.wait()
                return result
            finally:
                # Close the DB connection that was opened in this thread
                # (otherwise the connection is left opened and the test
                # runner might not be able to destroy the debusine-test DB)
                connections.close_all()

        self._executor = ThreadPoolExecutor()
        self._future = self._executor.submit(run_in_thread)
        self._transaction_ready.wait()

    def stop_transaction(self) -> R:
        """Close the open transaction and let the thread exit."""
        assert self._executor is not None
        assert self._future is not None
        self._transaction_close.set()
        result = self._future.result()
        self._executor.shutdown(wait=True)
        return result

    def run_and_wait(self) -> R:
        """Run the callable in a transaction and wait for it to finish."""
        self.start_transaction()
        return self.stop_transaction()


def _calculate_hash_from_data(data: bytes) -> bytes:
    hasher = hashlib.new(File.current_hash_algorithm)
    hasher.update(data)
    return hasher.digest()


def raise_db_exception(condition_name: str) -> None:
    """Tell PostgreSQL to raise an exception, aborting the transaction."""
    with connection.cursor() as cursor:
        cursor.execute(f"DO $$ BEGIN RAISE {condition_name}; END; $$")
