"""A :class:`~kaya.core.KayaMixin` that drives the TortoiseORM lifecycle.

kaya calls ``KayaMixin.setup`` / ``shutdown`` synchronously from inside a
running event loop. Tortoise 1.1.7 binds database connections to a
:class:`~tortoise.context.TortoiseContext` looked up via a contextvar, and
kaya dispatches each HTTP request (and WebSocket connection) as a separate
``loop.create_task``, so a context set by an early request does not
automatically reach later requests.

This mixin therefore:

1. Lazily builds a :class:`TortoiseContext` for the active event loop
   (rebuilding it if the running loop changes, which happens in tests that
   use a fresh ``asyncio.run`` per test).
2. Per HTTP request, binds that context to the current task via the
   ``_current_context`` contextvar so the handler — running in the same
   task as the ``before_request`` hook — sees an active context.
3. Binds the same context at the start of every WebSocket connection
   (``before_websocket`` hook), because the match-result write happens at
   the end of a WebSocket match. The long-lived connection task keeps the
   context for its whole lifetime.

It deliberately avoids the global-fallback singleton
(``_enable_global_fallback``), which Tortoise only allows to be set once
per process and would therefore break across event loops.

Schema management is split by backend: in-memory sqlite databases (the
test suite) get ``generate_schemas`` on every fresh context; Postgres
schemas are owned by aerich migrations and must be applied externally
(``aerich upgrade``, run by the ``db-migrate`` compose service) before
the app serves requests.
"""
from __future__ import annotations

from asyncio import AbstractEventLoop, get_running_loop
from logging import getLogger
from typing import AbstractSet, Optional, Sequence

from kaya.core import HttpContext, KayaApp, KayaMixin, WebSocket
from tortoise.context import TortoiseContext, _current_context

log = getLogger(__name__)


class TortoiseMixin(KayaMixin):
    """Initialize and tear down a per-loop :class:`TortoiseContext`."""

    def __init__(self,
                 database_url: str,
                 models_modules: Sequence[str],
                 skip_paths: AbstractSet[str] = frozenset({"/api/health"})) -> None:
        self._database_url = database_url
        self._models_modules = list(models_modules)
        self._skip_paths = skip_paths
        self._ctx: "Optional[TortoiseContext]" = None
        self._init_loop: "Optional[AbstractEventLoop]" = None

    def apply(self, app: KayaApp) -> None:
        app.add_before_request_hook(self._ensure_context)
        app.add_before_websocket_hook(self._ensure_ws_context)

    def setup(self, loop: AbstractEventLoop) -> None:
        pass

    def shutdown(self, loop: AbstractEventLoop) -> None:
        if self._init_loop is loop and self._ctx is not None:
            loop.create_task(self._ctx.close_connections())
        self._ctx = None
        self._init_loop = None

    async def _build_context(self) -> TortoiseContext:
        ctx = TortoiseContext()
        with ctx:
            await ctx.init(
                db_url=self._database_url,
                modules={"models": self._models_modules},
            )
            # Schema creation is only done for sqlite (in-memory test
            # databases). Postgres schemas are managed by aerich migrations
            # (applied by the db-migrate compose service / `aerich upgrade`).
            if self._database_url.startswith("sqlite"):
                await ctx.generate_schemas()
        return ctx

    async def _bind(self) -> None:
        loop = get_running_loop()
        if self._init_loop is not loop:
            if self._ctx is not None:
                # A previous test loop went away; drop its context.
                self._ctx = None
            self._ctx = await self._build_context()
            self._init_loop = loop
        assert self._ctx is not None
        _current_context.set(self._ctx)

    async def _ensure_context(self, ctx: HttpContext):
        if ctx.path in self._skip_paths:
            return None
        # Only API endpoints touch the database: auth callbacks, websocket
        # handshakes (handled by _ensure_ws_context) and the static SPA
        # catch-all must not pay for a Tortoise context.
        if not ctx.path.startswith("/api/"):
            return None
        await self._bind()
        return None

    async def _ensure_ws_context(self, ws: WebSocket):
        await self._bind()
        return None