"""ManagedConnection: never closed under the thread using it.

Closing a sqlite3 connection from one thread while another is inside a
statement on it is a use-after-free in SQLite ("Windows fatal exception:
access violation" in CI). A retired connection is closed at once when idle,
otherwise interrupted and closed by its owner as the call ends; it never
commits again.
"""

import contextlib
import os
import sqlite3
import sys
import tempfile
import threading
import unittest
import weakref
from unittest import mock

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from database.managed_connection import (  # noqa: E402
    ManagedConnection,
    ManagedCursor,
    await_closed,
    retire_connections,
)
from database.models import PromptModel  # noqa: E402
from tests.open_handles import assert_closed  # noqa: E402

WAIT = 30


class _Case(unittest.TestCase):
    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory()
        self.path = os.path.join(self._tmp.name, "managed.db")
        with contextlib.closing(sqlite3.connect(self.path)) as conn:
            conn.execute("PRAGMA journal_mode=WAL")
            conn.execute("CREATE TABLE t(x INTEGER)")
            conn.execute("INSERT INTO t VALUES (0)")
            conn.commit()
        self.opened = []

    def tearDown(self):
        for conn in self.opened:
            conn.close()
        try:
            assert_closed(self, self._tmp.name)
        finally:
            self._tmp.cleanup()

    def connect(self, **kwargs):
        conn = sqlite3.connect(
            self.path,
            check_same_thread=False,
            timeout=WAIT,
            factory=ManagedConnection,
            **kwargs,
        )
        self.opened.append(conn)
        return conn

    def value(self):
        with contextlib.closing(sqlite3.connect(self.path)) as conn:
            return conn.execute("SELECT x FROM t").fetchone()[0]


class TestRetire(_Case):
    def test_an_idle_connection_retired_from_another_thread_closes_at_once(self):
        conn = self.connect()
        conn.execute("SELECT 1").fetchone()
        closed = []
        thread = threading.Thread(target=lambda: closed.append(conn.retire()))
        thread.start()
        thread.join(WAIT)
        self.assertEqual(closed, [True])
        self.assertTrue(conn.wait_closed(0))
        with self.assertRaises(sqlite3.ProgrammingError):
            conn.execute("SELECT 1")
        self.assertTrue(conn.retire())  # retiring again is harmless

    def test_a_connection_in_use_is_closed_by_its_owner_when_the_call_ends(self):
        conn = self.connect()
        holder = sqlite3.connect(self.path, isolation_level=None)
        self.opened.append(holder)
        holder.execute("BEGIN IMMEDIATE")
        started = threading.Event()
        outcome = {}
        conn.set_trace_callback(
            lambda sql: started.set() if sql.startswith("UPDATE") else None
        )

        def write():
            try:
                with conn:
                    conn.execute("UPDATE t SET x = 1")
                outcome["ok"] = True
            except sqlite3.Error as error:
                outcome["error"] = error

        thread = threading.Thread(target=write)
        thread.start()
        self.assertTrue(started.wait(WAIT))
        # The owner is waiting on our lock inside its UPDATE: not closed here.
        self.assertFalse(conn.retire())
        self.assertFalse(conn.retire())
        self.assertFalse(conn.wait_closed(0))
        self.assertEqual(retire_connections([conn]), [conn])
        holder.execute("ROLLBACK")
        thread.join(WAIT)
        self.assertFalse(thread.is_alive())
        self.assertTrue(conn.wait_closed(WAIT))
        self.assertIsInstance(outcome.get("error"), sqlite3.Error)
        self.assertEqual(self.value(), 0)

    def test_a_write_retired_before_its_commit_is_rolled_back(self):
        conn = self.connect()
        conn.execute("UPDATE t SET x = 2")  # an implicit transaction is open
        thread = threading.Thread(target=conn.retire)
        thread.start()
        thread.join(WAIT)
        with self.assertRaises(sqlite3.ProgrammingError):
            conn.commit()
        with self.assertRaises(sqlite3.ProgrammingError):
            with conn:
                pass
        with self.assertRaises(sqlite3.ProgrammingError):
            conn.rollback()
        self.assertEqual(self.value(), 0)

    def test_close_is_retire(self):
        conn = self.connect()
        conn.close()
        self.assertTrue(conn.wait_closed(0))
        with self.assertRaises(sqlite3.ProgrammingError):
            conn.commit()

    def test_plain_connections_are_closed_directly(self):
        plain = sqlite3.connect(self.path)
        self.assertEqual(retire_connections([plain]), [])
        with self.assertRaises(sqlite3.ProgrammingError):
            plain.execute("SELECT 1")

    def test_awaiting_a_connection_still_in_use_times_out_with_a_warning(self):
        conn = self.connect()
        conn._enter()  # its owner is inside a call
        self.assertEqual(retire_connections([conn]), [conn])
        logger = mock.Mock()
        self.assertFalse(await_closed([conn], 0, logger))
        logger.warning.assert_called_once()
        conn._leave()
        self.assertTrue(await_closed([conn], WAIT, logger))


class TestCalls(_Case):
    def test_every_cursor_call_works_and_is_refused_once_retired(self):
        conn = self.connect()
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        self.assertIsInstance(cursor, ManagedCursor)
        conn.executemany("INSERT INTO t VALUES (?)", [(1,), (2,)])
        conn.executescript("INSERT INTO t VALUES (3);")
        self.assertEqual(
            [row["x"] for row in conn.execute("SELECT x FROM t ORDER BY x")],
            [0, 1, 2, 3],
        )
        cursor.execute("SELECT x FROM t ORDER BY x")
        self.assertEqual(cursor.fetchone()["x"], 0)
        self.assertEqual([r["x"] for r in cursor.fetchmany(2)], [1, 2])
        self.assertEqual([r["x"] for r in cursor.fetchall()], [3])
        cursor.executemany("UPDATE t SET x = x WHERE x = ?", [(1,)])
        cursor.executescript("SELECT 1;")
        cursor.close()
        conn.commit()
        conn.rollback()
        with conn:
            conn.execute("DELETE FROM t WHERE x = 3")
        with self.assertRaises(ZeroDivisionError):
            with conn:
                conn.execute("DELETE FROM t")
                raise ZeroDivisionError
        self.assertEqual(self.value(), 0)
        copy = sqlite3.connect(":memory:")
        self.opened.append(copy)
        conn.backup(copy)
        self.assertEqual(copy.execute("SELECT COUNT(*) FROM t").fetchone()[0], 3)

        reader = conn.execute("SELECT x FROM t")
        conn.retire()
        for call in (
            lambda: reader.fetchone(),
            lambda: reader.fetchmany(1),
            lambda: reader.fetchall(),
            lambda: next(reader),
            lambda: reader.execute("SELECT 1"),
            lambda: reader.executemany("SELECT ?", [(1,)]),
            lambda: reader.executescript("SELECT 1;"),
            lambda: conn.backup(copy),
        ):
            with self.assertRaises(sqlite3.ProgrammingError):
                call()
        # As with sqlite3's own cursors (the result varies by Python version).
        with contextlib.suppress(sqlite3.ProgrammingError):
            reader.close()


class _CursorOpenedDuringClose(weakref.WeakSet):
    """The cursor set of a connection being closed; a cursor opens mid-iteration.

    ``hook`` runs after the first cursor is yielded, where a thread switch
    could let the owner call ``conn.cursor()`` (Python 3.10-3.12 copy a
    WeakSet one item at a time). It runs in the closing thread itself, so
    the result does not depend on scheduling.
    """

    def __init__(self, items, hook):
        super().__init__(items)
        self.hook = hook

    def __iter__(self):
        hooked = False
        for item in super().__iter__():
            yield item
            if not hooked:
                hooked = True
                self.hook()


class TestReviewRound(_Case):
    """PR #321 review: I1 (cursor set race), M1 and M2 (unguarded paths)."""

    def test_a_cursor_opened_while_the_connection_closes_cannot_break_retire(self):
        conn = self.connect()
        kept = [conn.cursor() for _ in range(3)]
        refused = []

        def open_cursor():
            try:
                kept.append(conn.cursor())
            except sqlite3.ProgrammingError as error:
                refused.append(error)

        conn._cursors = _CursorOpenedDuringClose(conn._cursors, open_cursor)
        self.assertTrue(conn.retire())  # no "Set changed size during iteration"
        self.assertTrue(conn.wait_closed(0))
        self.assertEqual(len(refused), 1)
        self.assertEqual(len(kept), 3)

    def test_a_retired_connection_refuses_new_cursors_and_scripts(self):
        conn = self.connect()
        conn._enter()  # its owner is inside a call: retired, not yet closed
        try:
            self.assertFalse(conn.retire())
            for call in (
                conn.cursor,
                lambda: conn.executescript("CREATE TABLE late(x);"),
                lambda: conn.in_transaction,
                lambda: conn.create_function("f", 0, lambda: 1),
            ):
                with self.subTest(call=call):
                    with self.assertRaises(sqlite3.ProgrammingError):
                        call()
        finally:
            conn._leave()
        self.assertTrue(conn.wait_closed(0))
        with contextlib.closing(sqlite3.connect(self.path)) as check:
            tables = check.execute(
                "SELECT name FROM sqlite_master WHERE name='late'"
            ).fetchall()
        self.assertEqual(tables, [])

    def test_one_connection_failing_to_retire_does_not_skip_the_rest(self):
        class Broken(ManagedConnection):
            def retire(self):
                raise RuntimeError("retire failed")

        broken = sqlite3.connect(self.path, check_same_thread=False, factory=Broken)
        try:
            after = self.connect()
            logger = mock.Mock()
            pending = retire_connections([broken, after], logger)
            self.assertEqual(pending, [])
            self.assertTrue(after.wait_closed(0))
            logger.exception.assert_called_once()
        finally:
            ManagedConnection.retire(broken)

    def test_a_restore_marks_its_backup_target_in_use(self):
        model = PromptModel(os.path.join(self._tmp.name, "live.db"))
        try:
            self._restore_marks_its_target(model)
        finally:
            model.close_all()

    def _restore_marks_its_target(self, model):
        source_path = os.path.join(self._tmp.name, "source.db")
        PromptModel(source_path).close_all()
        target = model.get_connection()
        depths = []
        connect = sqlite3.connect

        class Source(sqlite3.Connection):
            def backup(self, into, *args, **kwargs):
                depths.append(into._depth)
                return super().backup(into, *args, **kwargs)

        def connect_source(*args, **kwargs):
            if "factory" in kwargs:
                return connect(*args, **kwargs)
            return connect(*args, **kwargs, factory=Source)

        with mock.patch("database.models.sqlite3.connect", connect_source):
            model._replace_contents_from(source_path)
        self.assertEqual(depths, [1])
        self.assertEqual(target._depth, 0)


class TestModelConnections(_Case):
    def test_the_model_hands_out_managed_connections_and_closes_them(self):
        model = PromptModel(os.path.join(self._tmp.name, "model.db"))
        conn = model.get_connection()
        self.assertIsInstance(conn, ManagedConnection)
        model.close()
        self.assertTrue(conn.wait_closed(0))
        conn = model.get_connection()
        model.close_all()
        self.assertTrue(conn.wait_closed(0))


if __name__ == "__main__":
    unittest.main()