"""
Query related functions.
"""

from datetime import datetime, timezone
from typing import List, Tuple

import sqlparse
from sqlalchemy import text
from sqlmodel import Session, create_engine

from datajunction.config import Settings
from datajunction.models.query import (
    ColumnMetadata,
    Query,
    QueryResults,
    QueryState,
    QueryWithResults,
    StatementResults,
)
from datajunction.typing import ColumnType, Description, SQLADialect, Stream, TypeEnum


def get_columns_from_description(
    description: Description,
    dialect: SQLADialect,
) -> List[ColumnMetadata]:
    """
    Extract column metadata from the cursor description.

    For now this uses the information from the cursor description, which only allow us to
    distinguish between 4 types (see ``TypeEnum``). In the future we should use a type
    inferrer to determine the types based on the query.
    """
    type_map = {
        TypeEnum.STRING: ColumnType.STR,
        TypeEnum.BINARY: ColumnType.BYTES,
        TypeEnum.NUMBER: ColumnType.FLOAT,
        TypeEnum.DATETIME: ColumnType.DATETIME,
    }

    columns = []
    for column in description or []:
        name, native_type = column[:2]
        for dbapi_type in TypeEnum:
            if native_type == getattr(dialect.dbapi, dbapi_type.value, None):
                type_ = type_map[dbapi_type]
                break
        else:
            # fallback to string
            type_ = ColumnType.STR

        columns.append(ColumnMetadata(name=name, type=type_))

    return columns


def run_query(query: Query) -> List[Tuple[str, List[ColumnMetadata], Stream]]:
    """
    Run a query and return its results.

    For each statement we return a tuple with the statement SQL, a description of the
    columns (name and type) and a stream of rows (tuples).
    """
    engine = create_engine(query.database.URI, **query.database.extra_params)
    connection = engine.connect()

    output: List[Tuple[str, List[ColumnMetadata], Stream]] = []
    statements = sqlparse.parse(query.executed_query)
    for statement in statements:
        # Druid doesn't like statements that end in a semicolon...
        sql = str(statement).strip().rstrip(";")

        results = connection.execute(text(sql))
        stream = (tuple(row) for row in results)
        columns = get_columns_from_description(
            results.cursor.description,
            engine.dialect,
        )
        output.append((sql, columns, stream))

    return output


def process_query(
    session: Session,
    settings: Settings,
    query: Query,
) -> QueryWithResults:
    """
    Process a query.
    """
    query.scheduled = datetime.now(timezone.utc)
    query.state = QueryState.SCHEDULED
    query.executed_query = query.submitted_query

    errors = []
    query.started = datetime.now(timezone.utc)
    try:
        root = []
        for sql, columns, stream in run_query(query):
            rows = list(stream)
            root.append(
                StatementResults(
                    sql=sql,
                    columns=columns,
                    rows=rows,
                    row_count=len(rows),
                ),
            )
        results = QueryResults(__root__=root)

        query.state = QueryState.FINISHED
        query.progress = 1.0
    except Exception as ex:  # pylint: disable=broad-except
        results = QueryResults(__root__=[])
        query.state = QueryState.FAILED
        errors = [str(ex)]

    query.finished = datetime.now(timezone.utc)

    session.add(query)
    session.commit()
    session.refresh(query)

    settings.results_backend.add(str(query.id), results.json())

    return QueryWithResults(results=results, errors=errors, **query.dict())

sqlparse.parse datetime.datetime.now datajunction.models.query.QueryResults sqlalchemy.text datajunction.models.query.ColumnMetadata sqlmodel.create_engine
[((1834, 1898), 'sqlmodel.create_engine', 'create_engine', (['query.database.URI'], {}), '(query.database.URI, **query.database.extra_params)\n', (1847, 1898), False, 'from sqlmodel import Session, create_engine\n'), ((2015, 2051), 'sqlparse.parse', 'sqlparse.parse', (['query.executed_query'], {}), '(query.executed_query)\n', (2029, 2051), False, 'import sqlparse\n'), ((2657, 2683), 'datetime.datetime.now', 'datetime.now', (['timezone.utc'], {}), '(timezone.utc)\n', (2669, 2683), False, 'from datetime import datetime, timezone\n'), ((2809, 2835), 'datetime.datetime.now', 'datetime.now', (['timezone.utc'], {}), '(timezone.utc)\n', (2821, 2835), False, 'from datetime import datetime, timezone\n'), ((3490, 3516), 'datetime.datetime.now', 'datetime.now', (['timezone.utc'], {}), '(timezone.utc)\n', (3502, 3516), False, 'from datetime import datetime, timezone\n'), ((3197, 3224), 'datajunction.models.query.QueryResults', 'QueryResults', ([], {'__root__': 'root'}), '(__root__=root)\n', (3209, 3224), False, 'from datajunction.models.query import ColumnMetadata, Query, QueryResults, QueryState, QueryWithResults, StatementResults\n'), ((1479, 1516), 'datajunction.models.query.ColumnMetadata', 'ColumnMetadata', ([], {'name': 'name', 'type': 'type_'}), '(name=name, type=type_)\n', (1493, 1516), False, 'from datajunction.models.query import ColumnMetadata, Query, QueryResults, QueryState, QueryWithResults, StatementResults\n'), ((2239, 2248), 'sqlalchemy.text', 'text', (['sql'], {}), '(sql)\n', (2243, 2248), False, 'from sqlalchemy import text\n'), ((3375, 3400), 'datajunction.models.query.QueryResults', 'QueryResults', ([], {'__root__': '[]'}), '(__root__=[])\n', (3387, 3400), False, 'from datajunction.models.query import ColumnMetadata, Query, QueryResults, QueryState, QueryWithResults, StatementResults\n')]