import httpx
from snowflake.connector import SnowflakeConnection

from relationalai.util.docutils import include_in_docs
from relationalai.util.error import exc


@include_in_docs
def http_client(conn: SnowflakeConnection, timeout: int = 300) -> httpx.Client:
    """Create an authenticated HTTP client for Snowflake Cortex API requests.

    Parameters
    ----------
    conn : SnowflakeConnection
        An active Snowflake connection used to extract the account locator
        and session token.
    timeout : int
        Request timeout in seconds. Default: 300.

    Returns
    -------
    httpx.Client
        A configured HTTP client with authentication headers and the
        correct Snowflake base URL.

    Examples
    --------

    >>> from relationalai.config import Config
    >>> conn = create_config().get_session().connection
    >>> client = http_client(conn)
    """
    # Account locator is required to use in the hostname for SSL verification.
    # We can't assume the configuration will contain the locator, so we fetch
    # it using the given snowflake connection, which is assumed to be valid
    cursor = conn.cursor()
    if cursor is None:
        exc("SnowflakeConnection", "Could not get cursor from connection.")
    result = cursor.execute("select CURRENT_ACCOUNT()")
    if result is None:
        exc("SnowflakeConnection", "CURRENT_ACCOUNT() query returned None.")
    row = result.fetchone()
    if not row or len(row) != 1:
        exc("SnowflakeConnection", "Could not select CURRENT_ACCOUNT().")
    account_locator = row[0]
    host = _hostname(account_locator)

    if conn._rest is None:
        exc("SnowflakeConnection", "Connection has no REST client (_rest is None).")
    token_data = conn._rest._token_request('ISSUE')
    if token_data is None:
        exc("SnowflakeConnection", "Token request returned None.")
    token_extract = token_data['data']['sessionToken']
    token = f'\"{token_extract}\"'
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json',
        'Authorization': f"Snowflake Token={token}"
    }
    return httpx.Client(
        headers=headers,
        timeout=timeout,
        base_url=host
    )


def _hostname(account_locator: str) -> str:
    return f"https://{account_locator}.snowflakecomputing.com"
