from __future__ import annotations

from typing import Any, TYPE_CHECKING, TypeVar

from .client import Client
from .client_sync import ClientSync
from .core import ClientCore
from relationalai.util.asyncio import raise_if_running_event_loop

if TYPE_CHECKING:
    from relationalai.config.config import Config


_ClientT = TypeVar("_ClientT", Client, ClientSync)


def _get_config(config: Config | None) -> Config:
    if config is not None:
        return config
    # Lazy import: keep `import relationalai.client` lightweight.
    from .wiring.wiring import default_config

    return default_config()


def _connect_lazy(
    client_type: type[_ClientT], *, config: Config, connection_name: str | None = None, token: str | None = None
) -> _ClientT:
    # Lazy by default: do not create transports here.
    core = ClientCore(config=config, _connection_name=connection_name, _token=token)
    return client_type(core=core)


def _from_session_client(
    client_type: type[_ClientT], *, session: Any, config: Config, token: str | None = None, app_name: str | None = None
) -> _ClientT:
    # Session is already provided; create SQL transport immediately.
    from relationalai.runtime.retry_classification import _is_retryable_snowflake_error

    from .middleware.compose import compose_middlewares
    from .transports.sql import SqlClient

    exec_cfg = config.execution
    middlewares, metrics = compose_middlewares(
        exec_cfg,
        is_retryable=lambda e, _req: _is_retryable_snowflake_error(e),
        logger_name="relationalai.client.execution",
    )
    sql = SqlClient(session=session, middlewares=middlewares, execution_metrics=metrics)
    core = ClientCore(config=config, _sql=sql, app_name=app_name, _connection_name=None, _token=token)
    return client_type(core=core)


async def connect(
    *, config: "Config | None" = None, connection_name: str | None = None, token: str | None = None
) -> Client:
    """Async-first entrypoint returning a `Client`.

    Use this from async code (you can `await` it). The returned `Client`:
    - exposes async lifecycle (`await client.aclose()` / `async with ...`)
    - exposes async service clients (e.g. `client.reasoners` -> `ReasonersClient`)

    Transport creation is **lazy**: no network calls are made by `connect(...)` itself.
    Transports are created on first use (SQL/HTTP/services).

    If you are in synchronous code, prefer `connect_sync()` which returns `ClientSync`.
    """
    cfg = _get_config(config)
    return _connect_lazy(Client, config=cfg, connection_name=connection_name, token=token)


def connect_sync(
    *, config: "Config | None" = None, connection_name: str | None = None, token: str | None = None
) -> ClientSync:
    """Synchronous entrypoint returning a `ClientSync`.

    Use this from synchronous code. The returned `ClientSync`:
    - exposes sync lifecycle (`client.close()` / `with ...`)
    - exposes synchronous service clients (e.g. `client.reasoners` -> `ReasonersClientSync`)

    Transport creation is **lazy**: no network calls are made by `connect_sync(...)` itself.
    Transports are created on first use (SQL/HTTP/services).

    This function raises if called while an event loop is already running in the
    current thread (use `await connect(...)` from async code instead).
    """
    raise_if_running_event_loop("connect_sync", use_in_async="`connect(...)` from async code")
    cfg = _get_config(config)
    return _connect_lazy(ClientSync, config=cfg, connection_name=connection_name, token=token)


async def from_session(
    *, session: Any, config: "Config | None" = None, token: str | None = None, app_name: str | None = None
) -> Client:
    """Construct an async `Client` from an existing Snowpark session."""
    cfg = _get_config(config)
    return _from_session_client(Client, session=session, config=cfg, token=token, app_name=app_name)


def from_session_sync(
    *, session: Any, config: "Config | None" = None, token: str | None = None, app_name: str | None = None
) -> ClientSync:
    """Construct a sync `ClientSync` from an existing Snowpark session.

    Raises if called from a thread that already has a running event loop.
    """
    raise_if_running_event_loop("from_session_sync", use_in_async="`from_session(...)` from async code")
    cfg = _get_config(config)
    return _from_session_client(ClientSync, session=session, config=cfg, token=token, app_name=app_name)


__all__ = [
    "Client",
    "ClientSync",
    "connect",
    "connect_sync",
    "from_session",
    "from_session_sync",
]

