"""Base connection protocol for all connection types."""

from __future__ import annotations

from abc import ABC
from typing import Any
from pydantic import BaseModel, ConfigDict, PrivateAttr

from ...util.docutils import include_in_docs

# include this module in documentation
__include_in_docs__ = True


@include_in_docs
class BaseConnection(BaseModel, ABC):
    """Base class for connection configuration objects.

    Subclasses define the connection-specific fields and implement :meth:`get_session` to create
    the underlying client/session object (Snowpark, DuckDB, requests, etc.). The session is cached
    on the instance; call :meth:`clear_session_cache` to force a fresh session on the next access.

    Notes
    -----
    Subclasses must define a ``type`` field as a ``typing.Literal`` for Pydantic discriminated unions.
    """
    # Note: `type` field is NOT defined here, subclasses must define it
    # with specific Literal types for discriminated unions to work correctly

    model_config = ConfigDict(arbitrary_types_allowed=True)

    # Private field for caching session (excluded from serialization)
    _cached_session: Any = PrivateAttr(default=None)

    @include_in_docs
    def get_session(self) -> Any:
        """Return the underlying session/client for this connection.

        Override this in subclasses to create the connection-specific session (for example,
        a Snowpark session, DuckDB connection, or :class:`requests.Session`).

        Returns
        -------
        Any
            The connection-specific session/client object.

        Raises
        ------
        NotImplementedError
            Always raised by :class:`~relationalai.config.connections.base.BaseConnection`.
        """
        raise NotImplementedError("Subclasses must implement get_session()")

    @include_in_docs
    def clear_session_cache(self) -> None:
        """Clear any cached session for this connection.

        After calling this, the next call to :meth:`get_session` creates a fresh
        session/client.
        """
        self._cached_session = None
