import pyarrow as pa
from ..connector import raw_sql
from textwrap import dedent
from typing import (
    Any, 
    Literal, 
    Callable
)
from .datatypes import map_typs
from operator import itemgetter
import sqlglot as sg
import sqlglot.expressions as sge
from ..utils import (
    is_query,
    rename_col
)


def schema_query_or_table(
    driver: str,
    name: str,
    database: str,
    schema: str
) -> list[tuple]:
    
    query = is_query(name)

    if query:
       name = sge.convert(str(name)).sql('tsql')
       stmt = dedent(
           f'''
           EXEC sp_describe_first_result_set @tsql = N{name}
           '''
       )

    else:
        where = [
            sg.column("table_name").eq(sge.convert(name)),
            sg.column("table_schema").eq(sge.convert(schema))
        ]

        stmt = (
            sg.select(
                "column_name",
                "data_type",
                "is_nullable",
                "numeric_precision",
                "numeric_scale",
                "datetime_precision"
            )
            .from_(
                sg.table(
                    'columns',
                    db='information_schema',
                    catalog=database
                )
            )
            .where(*where)
            .order_by("ordinal_position")
            .sql('tsql')
        )

    with raw_sql(driver) as cur:
        rst = cur.execute(stmt).fetchall()
    
    if query:

        def alter_tuple(row):
            row = list(row)
            tipo = row[1]

            if '(' in tipo:
                row[1] = tipo[:tipo.find('(')]

            row[2] = 'YES' if row[2] else 'NO'
            return tuple(row)

        filtro = itemgetter(2, 5, 3, 7, 8, 8)
        rst = [
               *map(
                   alter_tuple, 
                   map(
                       filtro, 
                       rst
                )
            )
        ]
    
    return rst
    

def get_schema(
    driver: str,
    name: str,
    database: str,
    schema: str = 'dbo'
) -> Callable:

    
    rst = schema_query_or_table(
        driver,
        name,
        database,
        schema
    )
    
    def tipo(
        retorno: Literal['sh', 'st', 'lt']
    ) -> Any:
        
        if retorno == 'lt':
            return rst
        
        schema = []

        for (
            cols,
            types,
            is_null,
            precision,
            scale,
            dt_precision
        ) in rst:

            dtype = map_typs[types]     
            if types in ('decimal', 'numeric'):
                dtype = dtype(precision, scale)
            
            if types == 'float':
                dtype = dtype(precision)
            
            if types in (
                'datetime',
                'datetime2', 
                'smalldatetime',
                'datetimeoffset',
                'time'
            ):
                dtype = dtype(dt_precision)

            field = pa.field(
                rename_col(cols), 
                dtype, 
                nullable= is_null == 'YES'
            )
            schema.append(field)
            
        if retorno == 'st':
            return pa.struct(schema)        
        return pa.schema(schema)
    
    return tipo