"""Module for interacting with Firebolt databases"""
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence

from firebolt.async_db.connection import Connection
from prefect import task
from prefect.blocks.core import Block
from prefect.utilities.asyncutils import sync_compatible
from pydantic import VERSION as PYDANTIC_VERSION

if PYDANTIC_VERSION.startswith("2."):
    from pydantic.v1 import Field, root_validator
else:
    from pydantic import Field, root_validator

from prefect_firebolt.credentials import FireboltCredentials

if TYPE_CHECKING:
    from firebolt.async_db._types import ColType, ParameterType


class FireboltDatabase(Block):
    """
    Connects to a Firebolt database.

    Provide either `engine_name` or `engine_url`. Providing both will result in an
    error. If neither `engine_name` nor `engine_url` is provided, the default engine
    for the configured database will be used.

    Attributes:
        credentials: Credentials to use to connect to the Firebase database.
        database_name: The name of the database to connect to.
        engine_name: Name of the engine to connect to.
        engine_url: The engine endpoint to use.
        additional_parameters: Additional configuration to pass to the Firebolt
            connection.
    """

    _block_type_name = "Firebolt Database"
    _logo_url = "https://images.ctfassets.net/gm98wzqotmnx/3loU17IXqVIWl4aWQfqc78/3c7eefe5e8cf4eec870856f10d7fdcce/5e8a264ceaf4870c90478037_Favicon_128.svg.png?h=250"  # noqa
    _documentation_url = "https://prefecthq.github.io/prefect-firebolt/database/#prefect_firebolt.database.FireboltDatabase"  # noqa
    _description = "Connects to a Firebolt database."

    credentials: FireboltCredentials
    database: str = Field(
        default=...,
        title="Database Name",
        description="The name of the database to connect to.",
    )
    engine_name: Optional[str] = Field(
        default=None,
        description="Name of the engine to connect to. May not be used with "
        "engine_url. If neither engine_name nor engine_url is provided, the "
        "default engine for the configured database will be used.",
    )
    engine_url: Optional[str] = Field(
        default=None,
        title="Engine URL",
        description="The engine endpoint to use. May not be used with engine_name. "
        "If neither engine_name nor engine_url is provided, the "
        "default engine for the configured database will be used.",
    )
    additional_parameters: Optional[Dict] = Field(
        default_factory=dict,
        description="Additional configuration to pass to the Firebolt connection.",
    )

    @root_validator(pre=True)
    def _not_both_engine_name_and_engine_url(cls, values):
        """Ensures that engine_name and engine_url are not both provided"""
        if (
            values.get("engine_name") is not None
            and values.get("engine_url") is not None
        ):
            raise ValueError(
                "You have provided a value for both engine_name and engine_url. "
                "Please provide either engine_name or engine_url, but not both."
            )
        return values

    def _get_connect_params(self) -> Dict[str, Any]:
        """Assemble the parameters to pass to the Firebolt connection"""
        return {
            "database": self.database,
            "engine_name": self.engine_name,
            "engine_url": self.engine_url,
            "additional_parameters": self.additional_parameters,
        }

    @sync_compatible
    async def get_connection(self) -> Connection:
        """
        Creates and returns an authenticated Firebolt connection for the
        configured database.
        """
        return await self.credentials.get_client(**self._get_connect_params())


@task
async def query_firebolt(
    database: FireboltDatabase,
    query: str,
    parameters: "Optional[Sequence[ParameterType]]" = None,
) -> "List[List[ColType]]":
    """
    Executes a query against a Firebolt database.

    Args:
        database: Firebolt database configuration to use for query execution.
        query: SQL query to execute.
        parameters: A sequence of substitution parameters. Used to replace `?`
            placeholders inside a query with actual values.

    Returns:
        All rows retrieved by the query.

    Example:
        Execute a query against a Firebolt database:
        ```python
        from prefect import flow

        from prefect_firebolt import FireboltCredentials, FireboltDatabase, query_firebolt


        @flow
        def run_firebolt_query():
            firebolt_database_block = FireboltDatabase(
                database="travel",
                credentials=FireboltCredentials(
                    username="arthur.dent@hitchhikers.com", password="dont42panic"
                ),
            )

            results = query_firebolt(
                database=firebolt_database_block,
                query="SELECT * FROM ex_intergalactic_trips LIMIT 100",
            )

            return results


        run_firebolt_query()
        ```
    """  # noqa
    async with await database.get_connection() as connection:
        cursor = connection.cursor()
        with cursor:
            await cursor.execute(query=query, parameters=parameters)
            return await cursor.fetchall()
